diff --git a/dataframe-operations.cabal b/dataframe-operations.cabal --- a/dataframe-operations.cabal +++ b/dataframe-operations.cabal @@ -1,6 +1,6 @@ -cabal-version: 2.4 +cabal-version: 3.4 name: dataframe-operations -version: 1.1.1.1 +version: 2.5.0.0 synopsis: Column operations, expression DSL, and statistics for the dataframe ecosystem. description: Untyped column operations (select, filter, sort, join, groupBy, @@ -30,15 +30,15 @@ library import: warnings exposed-modules: + DataFrame.Internal.Statistics DataFrame.Functions DataFrame.Monad - DataFrame.Internal.Statistics - DataFrame.Operations.AggregateScatter + DataFrame.Operations.Aggregation.Run DataFrame.Operations.Aggregation DataFrame.Operations.Core DataFrame.Operations.Inference DataFrame.Operations.Join - DataFrame.Operations.JoinPar + DataFrame.Operations.Join.Parallel DataFrame.Operations.Merge DataFrame.Operations.Permutation DataFrame.Operations.SetOps @@ -48,19 +48,23 @@ DataFrame.Operations.Typing DataFrame.Typed.Access DataFrame.Typed.Aggregate + DataFrame.Typed.Apply DataFrame.Typed.Expr + DataFrame.Typed.Expr.Extra DataFrame.Typed.Join DataFrame.Typed.Operations + DataFrame.Typed.Sampling + DataFrame.Typed.Statistics build-depends: base >= 4 && < 5, - bytestring >= 0.11 && < 0.13, - containers >= 0.6.7 && < 0.9, - dataframe-core >= 1.1 && < 1.2, - dataframe-parsing ^>= 1.0.2, + bytestring >= 0.11 && < 0.14, + containers >= 0.6.7 && < 0.10, + dataframe-core >= 2.5 && < 2.6, + dataframe-parsing >= 2.2 && < 2.3, random >= 1.2 && < 2, regex-tdfa >= 1.3.0 && < 2, text >= 2.1 && < 3, time >= 1.12 && < 2, - vector ^>= 0.13, - vector-algorithms ^>= 0.9 - hs-source-dirs: src + vector >= 0.13 && < 0.15, + vector-algorithms >= 0.9 && < 0.11 + hs-source-dirs: src, src-internal default-language: Haskell2010 diff --git a/src-internal/DataFrame/Internal/Statistics.hs b/src-internal/DataFrame/Internal/Statistics.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/Internal/Statistics.hs @@ -0,0 +1,285 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module DataFrame.Internal.Statistics where + +import qualified Data.Vector as V +import qualified Data.Vector.Algorithms.Intro as VA +import qualified Data.Vector.Mutable as VM +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM + +import Control.Exception (throw) +import Control.Monad.ST (runST) +import DataFrame.Errors (DataFrameException (..)) + +mean' :: (Real a, VU.Unbox a) => VU.Vector a -> Double +mean' samp + | VU.null samp = throw $ EmptyDataSetException "mean" + | otherwise = rtf (VU.sum samp) / fromIntegral (VU.length samp) +{-# INLINE [0] mean' #-} + +meanDouble' :: VU.Vector Double -> Double +meanDouble' samp + | VU.null samp = throw $ EmptyDataSetException "mean" + | otherwise = VU.sum samp / fromIntegral (VU.length samp) +{-# INLINE meanDouble' #-} + +meanInt' :: VU.Vector Int -> Double +meanInt' samp + | VU.null samp = throw $ EmptyDataSetException "mean" + | otherwise = fromIntegral (VU.sum samp) / fromIntegral (VU.length samp) +{-# INLINE meanInt' #-} + +{-# RULES +"mean'/Double" [1] forall (xs :: VU.Vector Double). + mean' xs = + meanDouble' xs +"mean'/Int" [1] forall (xs :: VU.Vector Int). + mean' xs = + meanInt' xs + #-} + +median' :: (Real a, VU.Unbox a) => VU.Vector a -> Double +median' samp + | VU.null samp = throw $ EmptyDataSetException "median" + | otherwise = runST $ do + mutableSamp <- VU.thaw samp + VA.sort mutableSamp + let len = VU.length samp + middleIndex = len `div` 2 + middleElement <- VUM.read mutableSamp middleIndex + if odd len + then pure (rtf middleElement) + else do + prev <- VUM.read mutableSamp (middleIndex - 1) + pure (rtf (middleElement + prev) / 2) +{-# INLINE median' #-} + +-- accumulator: count, mean, m2 +data VarAcc + = VarAcc {-# UNPACK #-} !Int {-# UNPACK #-} !Double {-# UNPACK #-} !Double + deriving (Show) + +varianceStep :: VarAcc -> Double -> VarAcc +varianceStep (VarAcc !n !meanVal !m2) !x = + let !n' = n + 1 + !delta = x - meanVal + !meanVal' = meanVal + delta / fromIntegral n' + !m2' = m2 + delta * (x - meanVal') + in VarAcc n' meanVal' m2' +{-# INLINE varianceStep #-} + +computeVariance :: VarAcc -> Double +computeVariance (VarAcc !n _ !m2) + | n < 2 = 0 -- or error "variance of <2 samples" + | otherwise = m2 / fromIntegral (n - 1) +{-# INLINE computeVariance #-} + +variance' :: (Real a, VU.Unbox a) => VU.Vector a -> Double +variance' = computeVariance . VU.foldl' varianceStep (VarAcc 0 0 0) . VU.map rtf +{-# INLINE variance' #-} + +varianceDouble' :: VU.Vector Double -> Double +varianceDouble' = computeVariance . VU.foldl' varianceStep (VarAcc 0 0 0) +{-# INLINE varianceDouble' #-} + +-- accumulator: count, mean, m2, m3 +data SkewAcc = SkewAcc !Int !Double !Double !Double deriving (Show) + +skewnessStep :: (VU.Unbox a, Num a, Real a) => SkewAcc -> a -> SkewAcc +skewnessStep (SkewAcc !n !meanVal !m2 !m3) !x' = + let !n' = n + 1 + x = rtf x' + !k = fromIntegral n' + !delta = x - meanVal + !meanVal' = meanVal + delta / k + !m2' = m2 + (delta ^ (2 :: Int) * (k - 1)) / k + !m3' = + m3 + + (delta ^ (3 :: Int) * (k - 1) * (k - 2)) / k ^ (2 :: Int) + - (3 * delta * m2) / k + in SkewAcc n' meanVal' m2' m3' +{-# INLINE skewnessStep #-} + +computeSkewness :: SkewAcc -> Double +computeSkewness (SkewAcc n _ m2 m3) + | n < 3 = 0 -- or error "skewness of <3 samples" + | otherwise = (sqrt (fromIntegral n - 1) * m3) / sqrt (m2 ^ (3 :: Int)) +{-# INLINE computeSkewness #-} + +skewness' :: (VU.Unbox a, Real a, Num a) => VU.Vector a -> Double +skewness' = computeSkewness . VU.foldl' skewnessStep (SkewAcc 0 0 0 0) +{-# INLINE skewness' #-} + +data CorrelationStats + = CorrelationStats + {-# UNPACK #-} !Double + {-# UNPACK #-} !Double + {-# UNPACK #-} !Double + {-# UNPACK #-} !Double + {-# UNPACK #-} !Double + +correlation' :: VU.Vector Double -> VU.Vector Double -> Maybe Double +correlation' xs ys + | n < 2 = Nothing + | VU.length xs /= VU.length ys = Nothing + | otherwise = + let nf = fromIntegral n + initial = CorrelationStats 0 0 0 0 0 + (CorrelationStats sumX sumY sumXX sumYY sumXY) = VU.ifoldl' step initial xs + + !num = nf * sumXY - sumX * sumY + !den = sqrt ((nf * sumXX - sumX * sumX) * (nf * sumYY - sumY * sumY)) + in Just (num / den) + where + n = VU.length xs + step (CorrelationStats sx sy sxx syy sxy) i x = + let !y = VU.unsafeIndex ys i + in CorrelationStats (sx + x) (sy + y) (sxx + x * x) (syy + y * y) (sxy + x * y) +{-# INLINE correlation' #-} + +quantiles' :: + (VU.Unbox a, Num a, Real a) => + VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double +quantiles' qs q samp + | VU.null samp = throw $ EmptyDataSetException "quantiles" + | q < 2 = throw $ WrongQuantileNumberException q + | VU.any (\i -> i < 0 || i > q) qs = throw $ WrongQuantileIndexException qs q + | otherwise = runST $ do + let !n = VU.length samp + mutableSamp <- VU.thaw samp + VA.sort mutableSamp + VU.mapM + ( \i -> do + let !p = fromIntegral i / fromIntegral q + !position = p * fromIntegral (n - 1) :: Double + !index = floor position :: Int + !f = position - fromIntegral index + x <- fmap rtf (VUM.read mutableSamp index) + if f == 0 + then return x + else do + y <- fmap rtf (VUM.read mutableSamp (index + 1)) + return $ (1 - f) * x + f * y + ) + qs +{-# INLINE quantiles' #-} + +percentile' :: (VU.Unbox a, Num a, Real a) => Int -> VU.Vector a -> Double +percentile' n = VU.head . quantiles' (VU.fromList [n]) 100 + +quantilesOrd' :: + (Ord a, Eq a) => + VU.Vector Int -> Int -> V.Vector a -> V.Vector a +quantilesOrd' qs q samp + | V.null samp = throw $ EmptyDataSetException "quantiles" + | q < 2 = throw $ WrongQuantileNumberException q + | VU.any (\i -> i < 0 || i > q) qs = throw $ WrongQuantileIndexException qs q + | otherwise = runST $ do + let !n = V.length samp + mutableSamp <- V.thaw samp + VA.sort mutableSamp + V.mapM + ( \i -> do + let !p = fromIntegral i / fromIntegral q :: Double + !position = p * fromIntegral (n - 1) + !index = floor position :: Int + -- This is not exact for Ord instances. + -- Figure out how to make it so. + VM.read mutableSamp index + ) + (V.convert qs) + +percentileOrd' :: (Ord a, Eq a) => Int -> V.Vector a -> a +percentileOrd' n = V.head . quantilesOrd' (VU.fromList [n]) 100 + +interQuartileRange' :: (VU.Unbox a, Num a, Real a) => VU.Vector a -> Double +interQuartileRange' samp = + let quartiles = quantiles' (VU.fromList [1, 3]) 4 samp + in quartiles VU.! 1 - quartiles VU.! 0 +{-# INLINE interQuartileRange' #-} + +meanSquaredError :: VU.Vector Double -> VU.Vector Double -> Maybe Double +meanSquaredError target prediction = + let + squareDiff = VU.ifoldl' (\sq i e -> (e - target VU.! i) ^ (2 :: Int) + sq) 0 prediction + in + Just $ squareDiff / fromIntegral (max (VU.length target) (VU.length prediction)) +{-# INLINE meanSquaredError #-} + +mutualInformationBinned :: + Int -> VU.Vector Double -> VU.Vector Double -> Maybe Double +mutualInformationBinned k xs ys + | VU.length xs /= VU.length ys = Nothing + | VU.null xs = Nothing + | k < 2 = Nothing + | rx <= 0 || ry <= 0 = Just 0 + | otherwise = + let bx = VU.map (binIndex xmin xmax k) xs + by = VU.map (binIndex ymin ymax k) ys + n = fromIntegral (VU.length xs) :: Double + mx = bincount k bx + my = bincount k by + mxy = jointBincount k bx by + in Just $ + sum + [ let !cxy = fromIntegral c + !pxy = cxy / n + !px = fromIntegral (mx VU.! i) / n + !py = fromIntegral (my VU.! j) / n + in if c == 0 then 0 else pxy * logBase 2 (pxy / (px * py)) + | i <- [0 .. k - 1] + , j <- [0 .. k - 1] + , let !c = mxy VU.! (i * k + j) + ] + where + (xmin, xmax) = (VU.minimum xs, VU.maximum xs) + (ymin, ymax) = (VU.minimum ys, VU.maximum ys) + rx = xmax - xmin + ry = ymax - ymin + +binIndex :: Double -> Double -> Int -> Double -> Int +binIndex lo hi k x + | hi == lo = 0 + | otherwise = + let !t = (x - lo) / (hi - lo) + !ix = floor (fromIntegral k * t) :: Int + in max 0 (min (k - 1) ix) +{-# INLINE binIndex #-} + +bincount :: Int -> VU.Vector Int -> VU.Vector Int +bincount k bs = VU.create $ do + mv <- VU.thaw (VU.replicate k 0) + VU.forM_ bs $ \b -> do + let i + | b < 0 = 0 + | b >= k = k - 1 + | otherwise = b + x <- VUM.read mv i + VUM.write mv i (x + 1) + pure mv +{-# INLINE bincount #-} + +jointBincount :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int +jointBincount k bx by = VU.create $ do + mv <- VU.thaw (VU.replicate (k * k) 0) + VU.forM_ (VU.zip bx by) $ \(i, j) -> do + let ii = clamp i 0 (k - 1) + jj = clamp j 0 (k - 1) + ix = ii * k + jj + x <- VUM.read mv ix + VUM.write mv ix (x + 1) + pure mv + where + clamp z a b = max a (min b z) +{-# INLINE jointBincount #-} + +rtf :: (Real a) => a -> Double +rtf = realToFrac +{-# NOINLINE [1] rtf #-} + +{-# RULES +"rtf/Double" [2] forall (x :: Double). rtf x = x + #-} diff --git a/src/DataFrame/Functions.hs b/src/DataFrame/Functions.hs --- a/src/DataFrame/Functions.hs +++ b/src/DataFrame/Functions.hs @@ -11,7 +11,15 @@ {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} -module DataFrame.Functions (module DataFrame.Functions, module DataFrame.Operators) where +module DataFrame.Functions ( + module DataFrame.Functions, + module DataFrame.Expression.Operators, + add, + sub, + mult, + divide, + prettyPrint, +) where import DataFrame.Internal.Column import DataFrame.Internal.Expression @@ -27,18 +35,20 @@ import qualified Data.Maybe as Maybe import qualified Data.Text as T import Data.Time +import Data.Type.Equality (testEquality, type (:~:) (Refl)) import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU -import DataFrame.Internal.Nullable ( +import DataFrame.Expression.Operators +import DataFrame.Internal.Expression.Operators.Nullable ( BaseType, NullLift1Op (applyNull1), NullLift1Result, NullLift2Op (applyNull2), NullLift2Result, ) -import DataFrame.Operators import Text.Regex.TDFA +import Type.Reflection (typeRep) import Prelude hiding (maximum, minimum) import Prelude as P @@ -321,6 +331,49 @@ {-# SPECIALIZE mean :: Expr Int64 -> Expr Double #-} {-# INLINEABLE mean #-} +-- | The k largest values per group. +topK :: (Columnable a, Ord a) => Int -> Expr a -> Expr [a] +topK = kExtremes "topK" (>) +{-# INLINEABLE topK #-} + +-- | The k smallest values per group. +bottomK :: (Columnable a, Ord a) => Int -> Expr a -> Expr [a] +bottomK = kExtremes "bottomK" (<) +{-# INLINEABLE bottomK #-} + +{- | Top k values after applying an ordering function. Floating-point NaNs +are dropped: they have no place in a total order, and admitting them would +make the merge non-associative (results would depend on chunk boundaries in +the lazy executor). +-} +kExtremes :: + forall a. + (Columnable a, Ord a) => + T.Text -> (a -> a -> Bool) -> Int -> Expr a -> Expr [a] +kExtremes opName wins k = + Agg + ( MergeAgg + (opName <> "_" <> T.pack (show k)) + [] + step + (L.foldl' step) + id + ) + where + step acc x + | isNaNLike x = acc + | otherwise = insert x acc + insert x = P.take k . go + where + go (z : zs) | wins z x = z : go zs + go zs = x : zs + + isNaNLike :: a -> Bool + isNaNLike x + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = isNaN x + | Just Refl <- testEquality (typeRep @a) (typeRep @Float) = isNaN x + | otherwise = False + meanMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double meanMaybe = Agg (CollectAgg "meanMaybe" (mean' . optionalToDoubleVector)) {-# SPECIALIZE meanMaybe :: Expr (Maybe Double) -> Expr Double #-} @@ -476,6 +529,16 @@ {-# SPECIALIZE isNothing :: Expr (Maybe Double) -> Expr Bool #-} {-# SPECIALIZE isNothing :: Expr (Maybe Int) -> Expr Bool #-} {-# INLINEABLE isNothing #-} + +-- | SQL spelling of 'isNothing'. +isNull :: (Columnable a) => Expr (Maybe a) -> Expr Bool +isNull = isNothing +{-# INLINEABLE isNull #-} + +-- | SQL spelling of 'isJust'. +isNotNull :: (Columnable a) => Expr (Maybe a) -> Expr Bool +isNotNull = isJust +{-# INLINEABLE isNotNull #-} fromJust :: (Columnable a) => Expr (Maybe a) -> Expr a fromJust = liftDecorated Maybe.fromJust "fromJust" Nothing diff --git a/src/DataFrame/Internal/Statistics.hs b/src/DataFrame/Internal/Statistics.hs deleted file mode 100644 --- a/src/DataFrame/Internal/Statistics.hs +++ /dev/null @@ -1,285 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} - -module DataFrame.Internal.Statistics where - -import qualified Data.Vector as V -import qualified Data.Vector.Algorithms.Intro as VA -import qualified Data.Vector.Mutable as VM -import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Unboxed.Mutable as VUM - -import Control.Exception (throw) -import Control.Monad.ST (runST) -import DataFrame.Errors (DataFrameException (..)) - -mean' :: (Real a, VU.Unbox a) => VU.Vector a -> Double -mean' samp - | VU.null samp = throw $ EmptyDataSetException "mean" - | otherwise = rtf (VU.sum samp) / fromIntegral (VU.length samp) -{-# INLINE [0] mean' #-} - -meanDouble' :: VU.Vector Double -> Double -meanDouble' samp - | VU.null samp = throw $ EmptyDataSetException "mean" - | otherwise = VU.sum samp / fromIntegral (VU.length samp) -{-# INLINE meanDouble' #-} - -meanInt' :: VU.Vector Int -> Double -meanInt' samp - | VU.null samp = throw $ EmptyDataSetException "mean" - | otherwise = fromIntegral (VU.sum samp) / fromIntegral (VU.length samp) -{-# INLINE meanInt' #-} - -{-# RULES -"mean'/Double" [1] forall (xs :: VU.Vector Double). - mean' xs = - meanDouble' xs -"mean'/Int" [1] forall (xs :: VU.Vector Int). - mean' xs = - meanInt' xs - #-} - -median' :: (Real a, VU.Unbox a) => VU.Vector a -> Double -median' samp - | VU.null samp = throw $ EmptyDataSetException "median" - | otherwise = runST $ do - mutableSamp <- VU.thaw samp - VA.sort mutableSamp - let len = VU.length samp - middleIndex = len `div` 2 - middleElement <- VUM.read mutableSamp middleIndex - if odd len - then pure (rtf middleElement) - else do - prev <- VUM.read mutableSamp (middleIndex - 1) - pure (rtf (middleElement + prev) / 2) -{-# INLINE median' #-} - --- accumulator: count, mean, m2 -data VarAcc - = VarAcc {-# UNPACK #-} !Int {-# UNPACK #-} !Double {-# UNPACK #-} !Double - deriving (Show) - -varianceStep :: VarAcc -> Double -> VarAcc -varianceStep (VarAcc !n !meanVal !m2) !x = - let !n' = n + 1 - !delta = x - meanVal - !meanVal' = meanVal + delta / fromIntegral n' - !m2' = m2 + delta * (x - meanVal') - in VarAcc n' meanVal' m2' -{-# INLINE varianceStep #-} - -computeVariance :: VarAcc -> Double -computeVariance (VarAcc !n _ !m2) - | n < 2 = 0 -- or error "variance of <2 samples" - | otherwise = m2 / fromIntegral (n - 1) -{-# INLINE computeVariance #-} - -variance' :: (Real a, VU.Unbox a) => VU.Vector a -> Double -variance' = computeVariance . VU.foldl' varianceStep (VarAcc 0 0 0) . VU.map rtf -{-# INLINE variance' #-} - -varianceDouble' :: VU.Vector Double -> Double -varianceDouble' = computeVariance . VU.foldl' varianceStep (VarAcc 0 0 0) -{-# INLINE varianceDouble' #-} - --- accumulator: count, mean, m2, m3 -data SkewAcc = SkewAcc !Int !Double !Double !Double deriving (Show) - -skewnessStep :: (VU.Unbox a, Num a, Real a) => SkewAcc -> a -> SkewAcc -skewnessStep (SkewAcc !n !meanVal !m2 !m3) !x' = - let !n' = n + 1 - x = rtf x' - !k = fromIntegral n' - !delta = x - meanVal - !meanVal' = meanVal + delta / k - !m2' = m2 + (delta ^ (2 :: Int) * (k - 1)) / k - !m3' = - m3 - + (delta ^ (3 :: Int) * (k - 1) * (k - 2)) / k ^ (2 :: Int) - - (3 * delta * m2) / k - in SkewAcc n' meanVal' m2' m3' -{-# INLINE skewnessStep #-} - -computeSkewness :: SkewAcc -> Double -computeSkewness (SkewAcc n _ m2 m3) - | n < 3 = 0 -- or error "skewness of <3 samples" - | otherwise = (sqrt (fromIntegral n - 1) * m3) / sqrt (m2 ^ (3 :: Int)) -{-# INLINE computeSkewness #-} - -skewness' :: (VU.Unbox a, Real a, Num a) => VU.Vector a -> Double -skewness' = computeSkewness . VU.foldl' skewnessStep (SkewAcc 0 0 0 0) -{-# INLINE skewness' #-} - -data CorrelationStats - = CorrelationStats - {-# UNPACK #-} !Double - {-# UNPACK #-} !Double - {-# UNPACK #-} !Double - {-# UNPACK #-} !Double - {-# UNPACK #-} !Double - -correlation' :: VU.Vector Double -> VU.Vector Double -> Maybe Double -correlation' xs ys - | n < 2 = Nothing - | VU.length xs /= VU.length ys = Nothing - | otherwise = - let nf = fromIntegral n - initial = CorrelationStats 0 0 0 0 0 - (CorrelationStats sumX sumY sumXX sumYY sumXY) = VU.ifoldl' step initial xs - - !num = nf * sumXY - sumX * sumY - !den = sqrt ((nf * sumXX - sumX * sumX) * (nf * sumYY - sumY * sumY)) - in Just (num / den) - where - n = VU.length xs - step (CorrelationStats sx sy sxx syy sxy) i x = - let !y = VU.unsafeIndex ys i - in CorrelationStats (sx + x) (sy + y) (sxx + x * x) (syy + y * y) (sxy + x * y) -{-# INLINE correlation' #-} - -quantiles' :: - (VU.Unbox a, Num a, Real a) => - VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double -quantiles' qs q samp - | VU.null samp = throw $ EmptyDataSetException "quantiles" - | q < 2 = throw $ WrongQuantileNumberException q - | VU.any (\i -> i < 0 || i > q) qs = throw $ WrongQuantileIndexException qs q - | otherwise = runST $ do - let !n = VU.length samp - mutableSamp <- VU.thaw samp - VA.sort mutableSamp - VU.mapM - ( \i -> do - let !p = fromIntegral i / fromIntegral q - !position = p * fromIntegral (n - 1) :: Double - !index = floor position :: Int - !f = position - fromIntegral index - x <- fmap rtf (VUM.read mutableSamp index) - if f == 0 - then return x - else do - y <- fmap rtf (VUM.read mutableSamp (index + 1)) - return $ (1 - f) * x + f * y - ) - qs -{-# INLINE quantiles' #-} - -percentile' :: (VU.Unbox a, Num a, Real a) => Int -> VU.Vector a -> Double -percentile' n = VU.head . quantiles' (VU.fromList [n]) 100 - -quantilesOrd' :: - (Ord a, Eq a) => - VU.Vector Int -> Int -> V.Vector a -> V.Vector a -quantilesOrd' qs q samp - | V.null samp = throw $ EmptyDataSetException "quantiles" - | q < 2 = throw $ WrongQuantileNumberException q - | VU.any (\i -> i < 0 || i > q) qs = throw $ WrongQuantileIndexException qs q - | otherwise = runST $ do - let !n = V.length samp - mutableSamp <- V.thaw samp - VA.sort mutableSamp - V.mapM - ( \i -> do - let !p = fromIntegral i / fromIntegral q :: Double - !position = p * fromIntegral (n - 1) - !index = floor position :: Int - -- This is not exact for Ord instances. - -- Figure out how to make it so. - VM.read mutableSamp index - ) - (V.convert qs) - -percentileOrd' :: (Ord a, Eq a) => Int -> V.Vector a -> a -percentileOrd' n = V.head . quantilesOrd' (VU.fromList [n]) 100 - -interQuartileRange' :: (VU.Unbox a, Num a, Real a) => VU.Vector a -> Double -interQuartileRange' samp = - let quartiles = quantiles' (VU.fromList [1, 3]) 4 samp - in quartiles VU.! 1 - quartiles VU.! 0 -{-# INLINE interQuartileRange' #-} - -meanSquaredError :: VU.Vector Double -> VU.Vector Double -> Maybe Double -meanSquaredError target prediction = - let - squareDiff = VU.ifoldl' (\sq i e -> (e - target VU.! i) ^ (2 :: Int) + sq) 0 prediction - in - Just $ squareDiff / fromIntegral (max (VU.length target) (VU.length prediction)) -{-# INLINE meanSquaredError #-} - -mutualInformationBinned :: - Int -> VU.Vector Double -> VU.Vector Double -> Maybe Double -mutualInformationBinned k xs ys - | VU.length xs /= VU.length ys = Nothing - | VU.null xs = Nothing - | k < 2 = Nothing - | rx <= 0 || ry <= 0 = Just 0 - | otherwise = - let bx = VU.map (binIndex xmin xmax k) xs - by = VU.map (binIndex ymin ymax k) ys - n = fromIntegral (VU.length xs) :: Double - mx = bincount k bx - my = bincount k by - mxy = jointBincount k bx by - in Just $ - sum - [ let !cxy = fromIntegral c - !pxy = cxy / n - !px = fromIntegral (mx VU.! i) / n - !py = fromIntegral (my VU.! j) / n - in if c == 0 then 0 else pxy * logBase 2 (pxy / (px * py)) - | i <- [0 .. k - 1] - , j <- [0 .. k - 1] - , let !c = mxy VU.! (i * k + j) - ] - where - (xmin, xmax) = (VU.minimum xs, VU.maximum xs) - (ymin, ymax) = (VU.minimum ys, VU.maximum ys) - rx = xmax - xmin - ry = ymax - ymin - -binIndex :: Double -> Double -> Int -> Double -> Int -binIndex lo hi k x - | hi == lo = 0 - | otherwise = - let !t = (x - lo) / (hi - lo) - !ix = floor (fromIntegral k * t) :: Int - in max 0 (min (k - 1) ix) -{-# INLINE binIndex #-} - -bincount :: Int -> VU.Vector Int -> VU.Vector Int -bincount k bs = VU.create $ do - mv <- VU.thaw (VU.replicate k 0) - VU.forM_ bs $ \b -> do - let i - | b < 0 = 0 - | b >= k = k - 1 - | otherwise = b - x <- VUM.read mv i - VUM.write mv i (x + 1) - pure mv -{-# INLINE bincount #-} - -jointBincount :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int -jointBincount k bx by = VU.create $ do - mv <- VU.thaw (VU.replicate (k * k) 0) - VU.forM_ (VU.zip bx by) $ \(i, j) -> do - let ii = clamp i 0 (k - 1) - jj = clamp j 0 (k - 1) - ix = ii * k + jj - x <- VUM.read mv ix - VUM.write mv ix (x + 1) - pure mv - where - clamp z a b = max a (min b z) -{-# INLINE jointBincount #-} - -rtf :: (Real a) => a -> Double -rtf = realToFrac -{-# NOINLINE [1] rtf #-} - -{-# RULES -"rtf/Double" [2] forall (x :: Double). rtf x = x - #-} diff --git a/src/DataFrame/Monad.hs b/src/DataFrame/Monad.hs --- a/src/DataFrame/Monad.hs +++ b/src/DataFrame/Monad.hs @@ -1,18 +1,54 @@ {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE InstanceSigs #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TupleSections #-} -module DataFrame.Monad where +module DataFrame.Monad ( + -- * The frame-state monad + FrameM, + runFrameM, + evalFrameM, + execFrameM, + modifyM, + inspectM, + -- * Frame verbs + deriveM, + insertM, + renameM, + filterWhereM, + sampleM, + takeM, + dropM, + selectM, + excludeM, + sortByM, + SortOrder (..), + columnAsListM, + filterJustM, + imputeM, + + -- * Column-derivation pipelines + Pipeline, + letAs, + letExpr, + pipelineSteps, + toFrameM, + runPipeline, +) where + +import Control.Monad (void) import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.DataFrame (DataFrame) -import DataFrame.Internal.Expression (Expr (..)) -import DataFrame.Internal.Nullable (BaseType) +import DataFrame.Internal.Expression (Expr (..), UExpr (..), prettyPrint) +import DataFrame.Internal.Expression.Operators.Nullable (BaseType) import qualified DataFrame.Operations.Core as D +import DataFrame.Operations.Permutation (SortOrder) +import qualified DataFrame.Operations.Permutation as D import qualified DataFrame.Operations.Subset as D import DataFrame.Operations.Transformations (ImputeOp) import qualified DataFrame.Operations.Transformations as D @@ -81,6 +117,17 @@ dropM :: Int -> FrameM () dropM n = modifyM (D.drop n) +-- | Keep only the named columns. 'dropM' drops rows; this drops columns. +selectM :: [T.Text] -> FrameM () +selectM names = modifyM (D.select names) + +-- | Drop the named columns. +excludeM :: [T.Text] -> FrameM () +excludeM names = modifyM (D.exclude names) + +sortByM :: [SortOrder] -> FrameM () +sortByM orders = modifyM (D.sortBy orders) + columnAsListM :: (Columnable a) => Expr a -> FrameM [a] columnAsListM c = inspectM (D.columnAsList c) @@ -105,3 +152,55 @@ execFrameM :: DataFrame -> FrameM a -> DataFrame execFrameM df m = snd (runFrameM df m) + +newtype Pipeline a = Pipeline {unPipeline :: Int -> (Int, [(T.Text, UExpr)], a)} + +instance Functor Pipeline where + fmap f (Pipeline g) = Pipeline $ \n -> let (n', w, a) = g n in (n', w, f a) + +instance Applicative Pipeline where + pure x = Pipeline (,[],x) + Pipeline gf <*> Pipeline gx = Pipeline $ \n -> + let (n1, w1, f) = gf n + (n2, w2, x) = gx n1 + in (n2, w1 ++ w2, f x) + +instance Monad Pipeline where + Pipeline g >>= f = Pipeline $ \n -> + let (n1, w1, a) = g n + (n2, w2, b) = unPipeline (f a) n1 + in (n2, w1 ++ w2, b) + +-- | Derive a column under an explicit name, returning a reference to it. +letAs :: (Columnable a) => T.Text -> Expr a -> Pipeline (Expr a) +letAs nm e = Pipeline (,[(nm, UExpr e)],Col nm) + +-- | Derive a column under a fresh generated name. +letExpr :: (Columnable a) => Expr a -> Pipeline (Expr a) +letExpr e = Pipeline $ \n -> + let nm = T.pack ('_' : 'v' : show n) in (n + 1, [(nm, UExpr e)], Col nm) + +instance Show (Pipeline (Expr a)) where + show p = + let (_, steps, res) = unPipeline p 0 + in concatMap + (\(nm, UExpr e) -> T.unpack nm ++ " = " ++ prettyPrint e ++ "\n") + steps + ++ "return " + ++ prettyPrint res + +-- | Number of column derivations the pipeline performs. +pipelineSteps :: Pipeline a -> Int +pipelineSteps p = let (_, w, _) = unPipeline p 0 in length w + +-- | Interpret a pipeline into 'FrameM', deriving each logged column in order. +toFrameM :: Pipeline (Expr a) -> FrameM (Expr a) +toFrameM p = + let (_, steps, res) = unPipeline p 0 + in mapM_ (\(nm, UExpr e) -> void (deriveM nm e)) steps >> pure res + +{- | Run a pipeline over a frame: derive its columns and return the result +expression (a reference to the final column) and the resulting frame. +-} +runPipeline :: DataFrame -> Pipeline (Expr a) -> (Expr a, DataFrame) +runPipeline df = runFrameM df . toFrameM diff --git a/src/DataFrame/Operations/AggregateScatter.hs b/src/DataFrame/Operations/AggregateScatter.hs deleted file mode 100644 --- a/src/DataFrame/Operations/AggregateScatter.hs +++ /dev/null @@ -1,224 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ExplicitNamespaces #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - -{- | Execute a recognised aggregation plan ('AggPlan') through the vectorized -scatter kernel, producing one result column (length @nGroups@, canonical group -order). The scatter reductions live in 'DataFrame.Internal.AggKernel' (sequential) -and 'DataFrame.Internal.AggKernelPar' (parallel by disjoint group range); this -module handles the compound @max - min@ combine and the holistic grouped median. -A plan only reaches here once 'planAgg' verified the value columns are clean -unboxed Int/Double, so the @error@ branches are unreachable. - -Every reduction takes the Round-5 grouping layout @(valueIndices, offsets)@ so -the parallel kernel can split the group-id range across capabilities with no -cross-worker merge. Each group's rows stay in original-row order within one -worker's range, so results are byte-identical to the sequential path at any @-N@. --} -module DataFrame.Operations.AggregateScatter (runPlan, runMomentPlan) where - -import qualified Data.Text as T -import qualified Data.Vector.Algorithms.Intro as VA -import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Unboxed.Mutable as VUM - -import Control.Concurrent (forkIO, getNumCapabilities) -import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar) -import Control.Exception (SomeException, throwIO, try) -import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) -import DataFrame.Internal.AggKernel (Reduction (..), scatterColumnToDouble) -import DataFrame.Internal.AggKernelDirect (directReduce, directThreshold) -import DataFrame.Internal.AggKernelPar (momentScatterPar, scatterReducePar) -import DataFrame.Internal.AggPlan (AggPlan (..), MomentPlan (..), Moments (..)) -import DataFrame.Internal.Column (Column (..), fromUnboxedVector) -import DataFrame.Internal.DataFrame (GroupedDataFrame (..), getColumn) -import System.IO.Unsafe (unsafePerformIO) -import Type.Reflection (typeRep) - -runPlan :: GroupedDataFrame -> VU.Vector Int -> Int -> AggPlan -> Column -runPlan gdf rtg nGroups plan = case plan of - PlanScatter red name -> scatterColumn red name - PlanMaxMinusMin a b -> maxMinusMin vis offs nGroups (col a) (col b) - PlanMedian name -> groupedMedian vis offs nGroups (col name) - where - vis = valueIndices gdf - offs = offsets gdf - {- The low-cardinality DIRECT-INDEXED fast path: for a small dense domain the - grouping layer's @rowToGroup@ already maps row -> group, so we scatter - straight off it (no @valueIndices@ gather). 'directReduce' only admits - order-independent reductions (so the merged parallel result is byte-identical - to -N1); anything it rejects keeps the order-preserving group-range kernel. -} - scatterColumn red name = - let c = col name - direct - | nGroups <= directThreshold = directReduce red rtg nGroups c - | otherwise = Nothing - in case direct of - Just out -> out - Nothing -> case scatterReducePar red vis offs nGroups c of - Just out -> out - Nothing -> error "runPlan: scatterReducePar rejected a planned column" - col name = case getColumn name (fullDataframe gdf) of - Just c -> c - Nothing -> error ("runPlan: planned column missing: " ++ T.unpack name) - -{- | Run a recognised moment (Q9 regression) plan as one fused scatter over the -two base columns, returning each output name bound to its moment field. The six -sufficient statistics (count, Sx, Sy, Sxx, Syy, Sxy) come out of a single pass, -replacing the three derive passes and six independent scatters of the -per-expression path. Byte-identical to the sequential kernel at any @-N@. --} -runMomentPlan :: - GroupedDataFrame -> Int -> MomentPlan -> Maybe [(T.Text, Column)] -runMomentPlan gdf nGroups mp = do - ms <- momentScatterPar vis offs nGroups (col (mpColX mp)) (col (mpColY mp)) - pure - [ (mpNName mp, mN ms) - , (mpSxName mp, mSx ms) - , (mpSyName mp, mSy ms) - , (mpSxxName mp, mSxx ms) - , (mpSyyName mp, mSyy ms) - , (mpSxyName mp, mSxy ms) - ] - where - vis = valueIndices gdf - offs = offsets gdf - col name = case getColumn name (fullDataframe gdf) of - Just c -> c - Nothing -> error ("runMomentPlan: planned column missing: " ++ T.unpack name) - -{- | @max a - min b@ on the small @nGroups@ arrays. Preserves the Int element -type of the source columns (matching the interpreter), falling back to a Double -combine otherwise. --} -maxMinusMin :: - VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -> Column -maxMinusMin vis offs nGroups ca cb = - case (ca, cb) of - ( UnboxedColumn Nothing (_ :: VU.Vector x) - , UnboxedColumn Nothing (_ :: VU.Vector y) - ) - | Just Refl <- testEquality (typeRep @x) (typeRep @Int) - , Just Refl <- testEquality (typeRep @y) (typeRep @Int) -> - let mx = scatterExtremaInt RMax vis offs nGroups ca - mn = scatterExtremaInt RMin vis offs nGroups cb - in fromUnboxedVector (VU.zipWith (-) mx mn) - _ -> - let mx = scatterExtremaDbl RMax vis offs nGroups ca - mn = scatterExtremaDbl RMin vis offs nGroups cb - in fromUnboxedVector (VU.zipWith (-) mx mn) - -scatterExtremaInt :: - Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> VU.Vector Int -scatterExtremaInt red vis offs nGroups c = case scatterReducePar red vis offs nGroups c of - Just (UnboxedColumn _ (v :: VU.Vector a)) - | Just Refl <- testEquality (typeRep @a) (typeRep @Int) -> v - _ -> error "scatterExtremaInt" - -scatterExtremaDbl :: - Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> VU.Vector Double -scatterExtremaDbl red vis offs nGroups c = - case scatterReducePar red vis offs nGroups c of - Just (UnboxedColumn _ (v :: VU.Vector a)) - | Just Refl <- testEquality (typeRep @a) (typeRep @Double) -> v - | Just Refl <- testEquality (typeRep @a) (typeRep @Int) -> VU.map fromIntegral v - _ -> error "scatterExtremaDbl" - -------------------------------------------------------------------------------- --- Parallel holistic median -------------------------------------------------------------------------------- - -{- | Holistic per-group median over a single unboxed Int/Double column. The -@valueIndices@/@offsets@ layout already places each group's rows in a contiguous -run, so we copy each group's values into a scratch buffer at its own offset and -sort that slice in place — each group's slice is independent, so the per-group -sorts split across capabilities by group range with no merge. Empty groups never -occur, so the result is total. --} -groupedMedian :: VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -groupedMedian vis offs nGroups c = case scatterColumnToDouble c of - Nothing -> error "groupedMedian: non-numeric planned column" - Just vals -> fromUnboxedVector (medianByGroup vis offs nGroups vals) - -medianByGroup :: - VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double -medianByGroup vis offs nGroups vals = unsafePerformIO $ do - let !n = VU.length vis - buf <- VUM.new (max 1 n) - out <- VUM.new (max 1 nGroups) - caps <- getNumCapabilities - let !bounds = groupRangeBounds offs nGroups caps - -- Each worker fills+sorts the buffer slices of its own group range, then - -- writes that range's medians. Disjoint ranges => safe to parallelise. - forEachRange bounds caps $ \gs ge -> - let grp !g - | g >= ge = pure () - | otherwise = do - let !s = VU.unsafeIndex offs g - !e = VU.unsafeIndex offs (g + 1) - !len = e - s - fill !pos - | pos >= e = pure () - | otherwise = do - VUM.unsafeWrite buf pos (VU.unsafeIndex vals (VU.unsafeIndex vis pos)) - fill (pos + 1) - fill s - let slice = VUM.unsafeSlice s len buf - VA.sort slice - let mid = s + len `div` 2 - med <- - if odd len - then VUM.unsafeRead buf mid - else do - hi <- VUM.unsafeRead buf mid - lo <- VUM.unsafeRead buf (mid - 1) - pure ((hi + lo) / 2) - VUM.unsafeWrite out g med - grp (g + 1) - in grp gs - VU.unsafeFreeze (VUM.unsafeSlice 0 nGroups out) -{-# NOINLINE medianByGroup #-} - -------------------------------------------------------------------------------- --- Group-range partitioning (shared with the median path) -------------------------------------------------------------------------------- - -{- | Split @[0, nGroups)@ into @caps@ contiguous group ranges balanced by row -count. Identical policy to 'DataFrame.Internal.AggKernelPar.groupRangeBounds'. --} -groupRangeBounds :: VU.Vector Int -> Int -> Int -> VU.Vector Int -groupRangeBounds offs nGroups caps = VU.create $ do - b <- VUM.new (caps + 1) - let !nRows = VU.unsafeIndex offs nGroups - !per = max 1 ((nRows + caps - 1) `div` caps) - adv !target !gg - | gg >= nGroups = nGroups - | VU.unsafeIndex offs gg >= target = gg - | otherwise = adv target (gg + 1) - go !w !prev - | w >= caps = VUM.unsafeWrite b caps nGroups - | otherwise = do - let !target = min nRows (w * per) - !g = adv target prev - VUM.unsafeWrite b w g - go (w + 1) g - VUM.unsafeWrite b 0 0 - go 1 0 - pure b - -forEachRange :: VU.Vector Int -> Int -> (Int -> Int -> IO ()) -> IO () -forEachRange bounds caps act - | caps <= 1 = act (VU.unsafeIndex bounds 0) (VU.unsafeIndex bounds caps) - | otherwise = do - vars <- mapM spawn [0 .. caps - 1] - results <- mapM takeMVar vars - mapM_ (either (throwIO :: SomeException -> IO ()) pure) results - where - spawn w = do - var <- newEmptyMVar - let !s = VU.unsafeIndex bounds w - !e = VU.unsafeIndex bounds (w + 1) - _ <- forkIO (try (act s e) >>= putMVar var) - pure var diff --git a/src/DataFrame/Operations/Aggregation.hs b/src/DataFrame/Operations/Aggregation.hs --- a/src/DataFrame/Operations/Aggregation.hs +++ b/src/DataFrame/Operations/Aggregation.hs @@ -14,29 +14,50 @@ changingPoints, ) where +import qualified Data.List as L +import qualified Data.Map.Strict as MS import qualified Data.Text as T import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU import Control.Exception (throw) import DataFrame.Errors -import DataFrame.Internal.AggPlan (MomentPlan, planAgg, planMoments) +import DataFrame.Internal.Aggregation.Kernel.Fused ( + mkFusedAgg, + mkGatherAgg, + runFusedAggs, + runGatherAggs, + ) +import DataFrame.Internal.Aggregation.Kernel.Scatter (streamGroupCap) +import DataFrame.Internal.Aggregation.Plan ( + AggPlan (..), + MomentPlan, + planAgg, + planMoments, + ) +import DataFrame.Internal.Aggregation.Reduction (Reduction (..)) import DataFrame.Internal.Column ( Column (..), TypedColumn (..), atIndicesStable, + atIndicesStableMulti, ) import DataFrame.Internal.DataFrame ( DataFrame (..), GroupedDataFrame (..), columnNames, + getColumn, insertColumn, ) import DataFrame.Internal.Expression import DataFrame.Internal.Grouping (buildRowToGroup, changingPoints, groupBy) import DataFrame.Internal.Interpreter -import DataFrame.Internal.RowHash (computeRowHashesIO) -import DataFrame.Operations.AggregateScatter (runMomentPlan, runPlan) +import DataFrame.Internal.Row.RowHash (computeRowHashesIO) +import DataFrame.Operations.Aggregation.Run ( + runMedianVarFused, + runMomentPlan, + runPlan, + ) import DataFrame.Operations.Core import DataFrame.Operations.Subset import System.IO.Unsafe (unsafePerformIO) @@ -55,25 +76,143 @@ {- | Aggregate a grouped dataframe using the expressions given. All ungrouped columns will be dropped. + +NOTE: this function deliberately never pattern-matches or strictly binds the +'Grouped' per-row fields (this module is compiled with @-XStrict@, whose strict +patterns and bindings would force them): on direct-grouped frames BOTH +'valueIndices' (the placement permutation) and 'rowToGroup' are deferred +thunks, and each aggregation path needs at most one of them — always passed as +un-forced argument expressions. Key columns materialize through 'groupRepRows' +(one representative row per group) instead of gathering +@valueIndices[offsets[g]]@. -} aggregate :: [NamedExpr] -> GroupedDataFrame -> DataFrame -aggregate aggs gdf@(Grouped df groupingColumns valIndices offs rowToGroupV) = +aggregate aggs gdf = let + df = fullDataframe gdf + offs = offsets gdf + + {- Key columns materialize through ONE fused parallel gather over the + representative rows ('atIndicesStableMulti'): the 1e8-group Q10 result + was six sequential latency-bound random-gather passes as per-column + 'selectIndices'. Each result column is still identical to (and as + deferred as) the per-column gather; the row-count field uses the eager + @offsets@ so nothing here forces 'groupRepRows' early. -} df' = - selectIndices - (VU.map (valIndices VU.!) (VU.init offs)) - (select groupingColumns df) + let sub = select (groupedColumns gdf) df + in sub + { columns = + V.fromList + ( atIndicesStableMulti + (groupRepRows gdf) + (V.toList (columns sub)) + ) + , dataframeDimensions = (nGroups, snd (dataframeDimensions sub)) + } !nGroups = VU.length offs - 1 + !nRows' = fst (dataframeDimensions df) + {- Fused multi-reduction fast path (Q3/Q4/Q5-shaped aggregates): every + recognised simple scatter reduction (sum/mean/count/min/max over a clean + unboxed Int/Double column) in this aggregate runs in ONE pass instead of + one full pass per expression. At or below 'streamGroupCap' groups that + pass streams over (rowToGroup, columns) with per-worker accumulators — + never touching the (possibly lazy) valueIndices; above it (necessarily a + hash-path grouping, whose valueIndices is already eager) the accumulator + arrays would thrash, so the pass gathers by disjoint group range instead + (register accumulators, bit-identical to the unfused gather kernels, + one traversal instead of one per expression). Only taken when at + least two reductions fuse; non-fusable expressions (median, var/std, + max-min, arbitrary DSL) keep their per-expression path below. -} + {- This binding is strict (-XStrict), so it must stay empty-and-cheap + whenever the moment path below already covers the aggregate — otherwise + the pass would run redundantly before the moment result is consulted. -} + fusedScatterCols :: MS.Map T.Text Column + fusedScatterCols = case fusedMoments of + Just _ -> MS.empty + Nothing + | nGroups <= streamGroupCap -> + let cands = + [ (name, fa) + | (name, ue) <- aggs + , Just (PlanScatter red cname) <- [planAgg gdf ue] + , Just c <- [getColumn cname df] + , Just fa <- [mkFusedAgg nGroups (rowToGroup gdf) red c] + ] + in if length cands >= 2 + then + MS.fromList + (zip (map fst cands) (runFusedAggs nRows' nGroups (map snd cands))) + else MS.empty + | otherwise -> + let cands = + [ (name, ga) + | (name, ue) <- aggs + , Just (PlanScatter red cname) <- [planAgg gdf ue] + , Just c <- [getColumn cname df] + , Just ga <- + [mkGatherAgg nGroups (valueIndices gdf) offs red c] + ] + in {- Unlike the stream branch, a SINGLE candidate also + takes this path: the per-expression fallback is a + scatter over rowToGroup, which on a hash-path grouping + is now a deferred thunk — the gather kernel (documented + bit-identical to the unfused kernels) works off the + already-eager valueIndices instead and skips that whole + random-write pass. -} + if not (null cands) + then + MS.fromList + ( zip + (map fst cands) + ( runGatherAggs + (valueIndices gdf) + offs + nGroups + (map snd cands) + ) + ) + else MS.empty + + {- Fused median + std/var over one column (the Q6 shape): both are + holistic gathers over the same values, so one shared gather serves the + Welford fold and the median selection ('runMedianVarFused', + bit-identical to the separate kernels). Only built when a median and a + std/var on the same column appear together; empty-and-cheap otherwise + (same strictness caveat as 'fusedScatterCols'). -} + medianVarCols :: MS.Map T.Text Column + medianVarCols = case fusedMoments of + Just _ -> MS.empty + Nothing -> + let plans = [(name, plan) | (name, ue) <- aggs, Just plan <- [planAgg gdf ue]] + medCols = L.nub [c | (_, PlanMedian c) <- plans] + in MS.fromList + [ kv + | cname <- medCols + , let stds = [nm | (nm, PlanScatter RStd c) <- plans, c == cname] + , let vars = [nm | (nm, PlanScatter RVar c) <- plans, c == cname] + , not (null stds && null vars) + , Just c <- [getColumn cname df] + , Just (medC, varC, stdC) <- [runMedianVarFused gdf nGroups c] + , kv <- + [(nm, medC) | (nm, PlanMedian c') <- plans, c' == cname] + ++ [(nm, stdC) | nm <- stds] + ++ [(nm, varC) | nm <- vars] + ] + -- Fast path: a recognised reduction scatters in one unboxed pass. -- Anything 'planAgg' rejects keeps the existing interpreter, so the -- general typed + DSL aggregate API stays correct for arbitrary -- expressions. f ne@(name, uexpr) d = - let value = case planAgg gdf uexpr of - Just plan -> runPlan gdf rowToGroupV nGroups plan - Nothing -> interpretNamed gdf ne + let value = case MS.lookup name medianVarCols of + Just c -> c + Nothing -> case MS.lookup name fusedScatterCols of + Just c -> c + Nothing -> case planAgg gdf uexpr of + Just plan -> runPlan gdf (rowToGroup gdf) nGroups plan + Nothing -> interpretNamed gdf ne in insertColumn name value d -- Fused fast path: the Q9 regression family (count + five moment sums @@ -107,4 +246,6 @@ distinct :: DataFrame -> DataFrame distinct df = selectIndices (VU.map (indices VU.!) (VU.init os)) df where - (Grouped _ _ indices os _rtg) = groupBy (columnNames df) df + -- The trailing field stays a wildcard: under -XStrict a named pattern + -- variable would force the (possibly deferred) rowToGroup thunk. + (Grouped _ _ indices os _) = groupBy (columnNames df) df diff --git a/src/DataFrame/Operations/Aggregation/Run.hs b/src/DataFrame/Operations/Aggregation/Run.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Aggregation/Run.hs @@ -0,0 +1,389 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Execute a recognised aggregation plan ('AggPlan'), producing one result +column (length @nGroups@, canonical group order). + +This is the dispatch layer: it owns the policy of WHICH kernel runs — the dense +direct-indexed one ('DataFrame.Internal.Aggregation.Kernel.Dense') when the +group domain is small enough, otherwise the group-range scatter +('DataFrame.Internal.Aggregation.Kernel.Scatter') — and handles the compound +@max - min@ combine and the holistic grouped median itself. The kernels carry no +policy of their own. A plan only reaches here once 'planAgg' verified the value +columns are clean unboxed Int/Double, so the @error@ branches are unreachable. + +Every reduction takes the Round-5 grouping layout @(valueIndices, offsets)@ so +the parallel kernel can split the group-id range across capabilities with no +cross-worker merge. Each group's rows stay in original-row order within one +worker's range, so results are byte-identical to the sequential path at any @-N@. +(Two exceptions, both deterministic at a fixed @-N@: the direct streaming +Double sum/mean above the small-group cutoff, whose chunked partials change the +float summation order, and the direct var/std, which finalize from +(count, sum, sumsq) partials rather than the gather kernel's row-order Welford +recurrence; see 'DataFrame.Internal.AggKernelDirect'.) +-} +module DataFrame.Operations.Aggregation.Run ( + runPlan, + runMomentPlan, + runMedianVarFused, +) where + +import qualified Data.Text as T +import qualified Data.Vector.Algorithms.Intro as VA +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM + +import Control.Concurrent (getNumCapabilities) +import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) +import DataFrame.Internal.Aggregation.Kernel.Dense ( + denseMaxMinusMin, + denseReduce, + ) +import DataFrame.Internal.Aggregation.Kernel.Moments ( + Moments (..), + momentScatterPar, + momentStreamPar, + ) +import DataFrame.Internal.Aggregation.Kernel.Scatter ( + maxMinusMinScatterPar, + scatterReducePar, + streamGroupCap, + ) +import DataFrame.Internal.Aggregation.Plan ( + AggPlan (..), + MomentPlan (..), + ) +import DataFrame.Internal.Aggregation.Reduction ( + Reduction (..), + cleanDoubleVector, + ) +import DataFrame.Internal.Column (Column (..), fromUnboxedVector) +import DataFrame.Internal.Control.Concurrent (parallelBounds_) +import DataFrame.Internal.DataFrame (GroupedDataFrame (..), getColumn) +import System.IO.Unsafe (unsafePerformIO) +import Type.Reflection (typeRep) + +{- | Group-domain size at or below which the dense direct-indexed kernel is +chosen; wider domains go to the group-range scatter kernel. A @2^18@-slot +accumulator is replicated per worker there, which is what bounds this. Dispatch +policy, so it lives with the dispatcher rather than inside the kernel. +-} +denseThreshold :: Int +denseThreshold = 262144 + +runPlan :: GroupedDataFrame -> VU.Vector Int -> Int -> AggPlan -> Column +runPlan gdf rtg nGroups plan = case plan of + PlanScatter red name -> scatterColumn red name + PlanMaxMinusMin a b -> + {- min/max are order-independent, so both fused single-pass kernels + (the direct streaming one up to 'streamGroupCap', the group-range + gather one above it) are exactly the two gather extrema they replace; + anything they reject (mixed/unclean columns, small inputs) keeps the + two-pass gather path. The streaming cap extends past 'denseThreshold' + for the same reason as the fused multi-reduction pass: on a + direct-grouped frame it works off the eager @rowToGroup@ and skips the + deferred @valueIndices@ placement entirely (measured at 1e6 groups / + 1e8 rows on -N16: stream 1.1s against placement 0.7s + gather 0.7s). -} + let ca = col a + cb = col b + direct + | nGroups <= streamGroupCap = denseMaxMinusMin rtg nGroups ca cb + | otherwise = maxMinusMinScatterPar vis offs nGroups ca cb + in case direct of + Just out -> out + Nothing -> maxMinusMin vis offs nGroups ca cb + PlanMedian name -> groupedMedian vis offs nGroups (col name) + where + vis = valueIndices gdf + offs = offsets gdf + {- The low-cardinality DENSE fast path: for a small dense domain the grouping + layer's @rowToGroup@ already maps row -> group, so we scatter straight off it + (no @valueIndices@ gather). 'denseReduce' admits the order-independent + reductions (exact partial merge, byte-identical to -N1) plus the streaming + Double sum/mean/var/std variants (byte-identical sequential row order at + small group counts; deterministic chunked partials for the large-domain + sum/mean — see "DataFrame.Internal.Aggregation.Kernel.Dense"); anything it + rejects keeps the order-preserving group-range kernel. -} + scatterColumn red name = + let c = col name + dense + | nGroups <= denseThreshold = denseReduce red rtg nGroups c + {- Top-2 selection merges exactly (a multiset selection, no + float adds until finalize), so like the fused passes it streams + off @rowToGroup@ up to 'streamGroupCap': on a direct-grouped + frame that skips the deferred @valueIndices@ placement, which + costs more than the accumulator cache misses it saves + (measured at 1e6 groups / 1e8 rows on -N16: stream 1.0s + against placement 0.7s + gather 0.4s). RTop2Snd shares the + same accumulator machinery and merge-exactness. -} + | red == RTop2Sum || red == RTop2Snd + , nGroups <= streamGroupCap = + denseReduce red rtg nGroups c + | otherwise = Nothing + in case dense of + Just out -> out + Nothing -> case scatterReducePar red vis offs nGroups c of + Just out -> out + Nothing -> error "runPlan: scatterReducePar rejected a planned column" + col name = case getColumn name (fullDataframe gdf) of + Just c -> c + Nothing -> error ("runPlan: planned column missing: " ++ T.unpack name) + +{- | Run a recognised moment (Q9 regression) plan as one fused scatter over the +two base columns, returning each output name bound to its moment field. The six +sufficient statistics (count, Sx, Sy, Sxx, Syy, Sxy) come out of a single pass, +replacing the three derive passes and six independent scatters of the +per-expression path. The streaming kernel's count is exact; its five Double +sums accumulate per worker chunk in original row order and merge in fixed +worker order — deterministic at a fixed @-N@, float summation order chunk-major +rather than per-group (see 'momentStreamPar'). The gather fallback remains +byte-identical to the sequential kernel at any @-N@. +-} +runMomentPlan :: + GroupedDataFrame -> Int -> MomentPlan -> Maybe [(T.Text, Column)] +runMomentPlan gdf nGroups mp = do + let cx = col (mpColX mp) + cy = col (mpColY mp) + {- Preferred: the streaming kernel — one fused pass over rowToGroup and + the TYPED base columns (no sequential Int->Double materialization, no + valueIndices gather, so a direct-grouped frame never runs its placement + pass). Falls back to the gather kernel above 'streamGroupCap' or on + unclean columns. -} + streamed = momentStreamPar (rowToGroup gdf) nGroups cx cy + ms <- case streamed of + Just m -> Just m + Nothing -> momentScatterPar vis offs nGroups cx cy + pure + [ (mpNName mp, mN ms) + , (mpSxName mp, mSx ms) + , (mpSyName mp, mSy ms) + , (mpSxxName mp, mSxx ms) + , (mpSyyName mp, mSyy ms) + , (mpSxyName mp, mSxy ms) + ] + where + vis = valueIndices gdf + offs = offsets gdf + col name = case getColumn name (fullDataframe gdf) of + Just c -> c + Nothing -> error ("runMomentPlan: planned column missing: " ++ T.unpack name) + +{- | @max a - min b@ on the small @nGroups@ arrays. Preserves the Int element +type of the source columns (matching the interpreter), falling back to a Double +combine otherwise. +-} +maxMinusMin :: + VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -> Column +maxMinusMin vis offs nGroups ca cb = + case (ca, cb) of + ( UnboxedColumn Nothing (_ :: VU.Vector x) + , UnboxedColumn Nothing (_ :: VU.Vector y) + ) + | Just Refl <- testEquality (typeRep @x) (typeRep @Int) + , Just Refl <- testEquality (typeRep @y) (typeRep @Int) -> + let mx = scatterExtremaInt RMax vis offs nGroups ca + mn = scatterExtremaInt RMin vis offs nGroups cb + in fromUnboxedVector (VU.zipWith (-) mx mn) + _ -> + let mx = scatterExtremaDbl RMax vis offs nGroups ca + mn = scatterExtremaDbl RMin vis offs nGroups cb + in fromUnboxedVector (VU.zipWith (-) mx mn) + +scatterExtremaInt :: + Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> VU.Vector Int +scatterExtremaInt red vis offs nGroups c = case scatterReducePar red vis offs nGroups c of + Just (UnboxedColumn _ (v :: VU.Vector a)) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) -> v + _ -> error "scatterExtremaInt" + +scatterExtremaDbl :: + Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> VU.Vector Double +scatterExtremaDbl red vis offs nGroups c = + case scatterReducePar red vis offs nGroups c of + Just (UnboxedColumn _ (v :: VU.Vector a)) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) -> v + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) -> VU.map fromIntegral v + _ -> error "scatterExtremaDbl" + +------------------------------------------------------------------------------- +-- Fused holistic median + var/std over one shared gather +------------------------------------------------------------------------------- + +{- | Fused grouped median and variance family over the SAME column: one gather +into the shared scratch buffer serves both. Returns +@(median, variance, stddev)@ columns, or 'Nothing' on a non-numeric column +(the caller keeps the separate per-expression kernels). + +The Welford fold runs over each gathered slice in ascending original-row order +— exactly the recurrence, order and finalize of the var/std kernels +('DataFrame.Internal.AggKernelPar.varPar' and the sequential @varScatter@, +which agree bit-for-bit) — BEFORE the in-place median selection permutes the +slice, and the selection then proceeds exactly as 'groupedMedian'. Both +outputs are therefore bit-identical to the unfused paths; the second full +gather pass is what the fusion saves (measured ~35% off the median+sd pair at +1e4 groups / 1e8 rows on -N16). +-} +runMedianVarFused :: + GroupedDataFrame -> Int -> Column -> Maybe (Column, Column, Column) +runMedianVarFused gdf nGroups c = do + vals <- cleanDoubleVector c + let (med, var) = + medianVarByGroup (valueIndices gdf) (offsets gdf) nGroups vals + pure + ( fromUnboxedVector med + , fromUnboxedVector var + , fromUnboxedVector (VU.map sqrt var) + ) + +medianVarByGroup :: + VU.Vector Int -> + VU.Vector Int -> + Int -> + VU.Vector Double -> + (VU.Vector Double, VU.Vector Double) +medianVarByGroup vis offs nGroups vals = unsafePerformIO $ do + let !n = VU.length vis + buf <- VUM.new (max 1 n) + medOut <- VUM.new (max 1 nGroups) + varOut <- VUM.new (max 1 nGroups) + caps <- getNumCapabilities + let !bounds = groupRangeBounds offs nGroups caps + parallelBounds_ caps bounds $ \gs ge -> + let grp !g + | g >= ge = pure () + | otherwise = do + let !s = VU.unsafeIndex offs g + !e = VU.unsafeIndex offs (g + 1) + !len = e - s + fill !pos + | pos >= e = pure () + | otherwise = do + VUM.unsafeWrite buf pos (VU.unsafeIndex vals (VU.unsafeIndex vis pos)) + fill (pos + 1) + fill s + -- Welford over the gathered slice (still in row order). + let welford !pos !c !mu !mm + | pos >= e = + pure (if c < 2 then 0 else mm / fromIntegral (c - 1)) + | otherwise = do + x <- VUM.unsafeRead buf pos + let !c' = c + 1 + !delta = x - mu + !mu' = mu + delta / fromIntegral c' + !mm' = mm + delta * (x - mu') + welford (pos + 1) c' mu' mm' + var <- welford s (0 :: Int) 0 0 + VUM.unsafeWrite varOut g var + -- Median selection, as in 'medianByGroup' (permutes the slice). + let slice = VUM.unsafeSlice s len buf + !mid = len `div` 2 + VA.select slice (mid + 1) + let scan !i !hi !lo + | i > mid = pure (hi, lo) + | otherwise = do + x <- VUM.unsafeRead slice i + if x > hi + then scan (i + 1) x hi + else scan (i + 1) hi (max lo x) + (hi, lo) <- scan 0 (negate (1 / 0)) (negate (1 / 0)) + let med = if odd len then hi else (hi + lo) / 2 + VUM.unsafeWrite medOut g med + grp (g + 1) + in grp gs + med <- VU.unsafeFreeze (VUM.unsafeSlice 0 nGroups medOut) + var <- VU.unsafeFreeze (VUM.unsafeSlice 0 nGroups varOut) + pure (med, var) +{-# NOINLINE medianVarByGroup #-} + +------------------------------------------------------------------------------- +-- Parallel holistic median +------------------------------------------------------------------------------- + +{- | Holistic per-group median over a single unboxed Int/Double column. The +@valueIndices@/@offsets@ layout already places each group's rows in a contiguous +run, so we copy each group's values into a scratch buffer at its own offset and +select the median-rank order statistics in that slice in place (O(len) per +group rather than the O(len log len) full sort) — each group's slice is +independent, so the per-group selections split across capabilities by group +range with no merge. Order statistics are value-determined, so the result is +identical to the sorting variant. Empty groups never occur, so the result is +total. +-} +groupedMedian :: VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column +groupedMedian vis offs nGroups c = case cleanDoubleVector c of + Nothing -> error "groupedMedian: non-numeric planned column" + Just vals -> fromUnboxedVector (medianByGroup vis offs nGroups vals) + +medianByGroup :: + VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double +medianByGroup vis offs nGroups vals = unsafePerformIO $ do + let !n = VU.length vis + buf <- VUM.new (max 1 n) + out <- VUM.new (max 1 nGroups) + caps <- getNumCapabilities + let !bounds = groupRangeBounds offs nGroups caps + -- Each worker fills+sorts the buffer slices of its own group range, then + -- writes that range's medians. Disjoint ranges => safe to parallelise. + parallelBounds_ caps bounds $ \gs ge -> + let grp !g + | g >= ge = pure () + | otherwise = do + let !s = VU.unsafeIndex offs g + !e = VU.unsafeIndex offs (g + 1) + !len = e - s + fill !pos + | pos >= e = pure () + | otherwise = do + VUM.unsafeWrite buf pos (VU.unsafeIndex vals (VU.unsafeIndex vis pos)) + fill (pos + 1) + fill s + let slice = VUM.unsafeSlice s len buf + !mid = len `div` 2 + {- Move the least mid+1 values to the front (in no + particular order); the two largest of those are the order + statistics at sorted positions mid and mid-1. -} + VA.select slice (mid + 1) + let scan !i !hi !lo + | i > mid = pure (hi, lo) + | otherwise = do + x <- VUM.unsafeRead slice i + if x > hi + then scan (i + 1) x hi + else scan (i + 1) hi (max lo x) + (hi, lo) <- scan 0 (negate (1 / 0)) (negate (1 / 0)) + let med = if odd len then hi else (hi + lo) / 2 + VUM.unsafeWrite out g med + grp (g + 1) + in grp gs + VU.unsafeFreeze (VUM.unsafeSlice 0 nGroups out) +{-# NOINLINE medianByGroup #-} + +------------------------------------------------------------------------------- +-- Group-range partitioning (shared with the median path) +------------------------------------------------------------------------------- + +{- | Split @[0, nGroups)@ into @caps@ contiguous group ranges balanced by row +count. Identical policy to 'DataFrame.Internal.Aggregation.Kernel.Scatter.groupRangeBounds'. +-} +groupRangeBounds :: VU.Vector Int -> Int -> Int -> VU.Vector Int +groupRangeBounds offs nGroups caps = VU.create $ do + b <- VUM.new (caps + 1) + let !nRows = VU.unsafeIndex offs nGroups + !per = max 1 ((nRows + caps - 1) `div` caps) + adv !target !gg + | gg >= nGroups = nGroups + | VU.unsafeIndex offs gg >= target = gg + | otherwise = adv target (gg + 1) + go !w !prev + | w >= caps = VUM.unsafeWrite b caps nGroups + | otherwise = do + let !target = min nRows (w * per) + !g = adv target prev + VUM.unsafeWrite b w g + go (w + 1) g + VUM.unsafeWrite b 0 0 + go 1 0 + pure b diff --git a/src/DataFrame/Operations/Core.hs b/src/DataFrame/Operations/Core.hs --- a/src/DataFrame/Operations/Core.hs +++ b/src/DataFrame/Operations/Core.hs @@ -6,8 +6,47 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -module DataFrame.Operations.Core where +module DataFrame.Operations.Core ( + -- * Dimensions + dimensions, + nRows, + nColumns, + -- * Construction + fromUnnamedColumns, + fromRows, + + -- * Insertion + insert, + insertVector, + insertUnboxedVector, + insertWithDefault, + insertVectorWithDefault, + + -- * Column management + cloneColumn, + rename, + renameMany, + + -- * Inspection + describeColumns, + valueCounts, + valueProportions, + showDerivedExpressions, + + -- * Folds & matrix/vector conversions + fold, + toFloatMatrix, + toDoubleMatrix, + toIntMatrix, + columnAsVector, + columnAsList, + columnAsIntVector, + columnAsDoubleVector, + columnAsFloatVector, + columnAsUnboxedVector, +) where + import qualified Data.List as L import qualified Data.Map as M import qualified Data.Map.Strict as MS @@ -32,7 +71,9 @@ columnTypeString, fromList, fromVector, + materializeMerged, materializePacked, + mergedHead, toDoubleVector, toFloatVector, toIntVector, @@ -117,10 +158,8 @@ -- 'columnNames' is now defined in "DataFrame.Internal.DataFrame" and -- re-exported from "DataFrame" at the top level. -{- | Adds a vector to the dataframe. If the vector has less elements than the dataframe and the dataframe is not empty -the vector is converted to type `Maybe a` filled with `Nothing` to match the size of the dataframe. Similarly, -if the vector has more elements than what's currently in the dataframe, the other columns in the dataframe are -change to `Maybe ` and filled with `Nothing`. +{- | Adds a vector as a named column. Size mismatches are reconciled by making +the shorter side nullable (`Maybe a`) and padding with `Nothing`. ==== __Example__ @ @@ -160,13 +199,9 @@ insertVector name xs = insertColumn name (fromVector xs) {-# INLINE insertVector #-} -{- | Adds a foldable collection to the dataframe. If the collection has less elements than the -dataframe and the dataframe is not empty -the collection is converted to type `Maybe a` filled with `Nothing` to match the size of the dataframe. Similarly, -if the collection has more elements than what's currently in the dataframe, the other columns in the dataframe are -change to `Maybe ` and filled with `Nothing`. - -Be careful not to insert infinite collections with this function as that will crash the program. +{- | Adds a foldable collection as a named column. Size mismatches are reconciled by +making the shorter side nullable (`Maybe a`) and padding with `Nothing`. +Do not pass infinite collections: they are fully forced. ==== __Example__ @ @@ -202,7 +237,7 @@ -- | DataFrame to add column to DataFrame -> DataFrame -insert name xs = insertColumn name (fromList (Fold.foldr' (:) [] xs)) -- TODO: Do reflection on container type so we can sometimes avoid the list construction. +insert name xs = insertColumn name (fromList (Fold.foldr' (:) [] xs)) {-# INLINE insert #-} {- | Adds a vector to the dataframe and pads it with a default value if it has less elements than the number of rows. @@ -295,10 +330,8 @@ values = xs' ++ replicate (rows - length xs') defaultValue in insertColumn name (fromList values) d -{- | /O(n)/ Adds an unboxed vector to the dataframe. - -Same as insertVector but takes an unboxed vector. If you insert a vector of numbers through insertVector it will either way be converted -into an unboxed vector so this function saves that extra work/conversion. +{- | /O(n)/ Like 'insertVector' but takes an already-unboxed vector, +skipping the boxed-to-unboxed conversion 'insertVector' would do for numbers. -} insertUnboxedVector :: forall a. @@ -540,10 +573,10 @@ columnType : acc go acc i col@(PackedText _ _) = go acc i (materializePacked col) + go acc i col@(MergedColumn _ _) = go acc i (materializeMerged col) nulls :: Column -> Int nulls (BoxedColumn (Just bm) xs) = - -- count null bits in bitmap let n = VG.length xs in n - VU.foldl' (\acc b -> acc + popCount b) 0 bm nulls (BoxedColumn Nothing (xs :: V.Vector a)) = case testEquality (typeRep @a) (typeRep @T.Text) of @@ -635,9 +668,9 @@ fromRows :: [T.Text] -> [[Any]] -> DataFrame fromRows names rows = L.foldl' - (\df i -> insertColumn (names !! i) (mkColumnFromRow i rows) df) + (\df (i, name) -> insertColumn name (mkColumnFromRow name i rows) df) empty - [0 .. length names - 1] + (zip [0 ..] names) {- | O (k * n) Counts the occurences of each value in a given column. @@ -731,14 +764,8 @@ fold :: (a -> DataFrame -> DataFrame) -> [a] -> DataFrame -> DataFrame fold f xs acc = L.foldl' (flip f) acc xs -{- | Returns a dataframe as a two dimensional vector of floats. - -Converts all columns in the dataframe to float vectors and transposes them -into a row-major matrix representation. - -This is useful for handing data over into ML systems. - -Returns 'Left' with an error if any column cannot be converted to floats. +{- | The dataframe as a row-major matrix of floats, for handing data to ML systems. +'Left' if any column cannot be converted to floats. -} toFloatMatrix :: DataFrame -> Either DataFrameException (V.Vector (VU.Vector Float)) @@ -753,14 +780,8 @@ (fst (dataframeDimensions df)) (\i -> VU.generate (V.length m) (\j -> (m VG.! j) VG.! i)) -{- | Returns a dataframe as a two dimensional vector of doubles. - -Converts all columns in the dataframe to double vectors and transposes them -into a row-major matrix representation. - -This is useful for handing data over into ML systems. - -Returns 'Left' with an error if any column cannot be converted to doubles. +{- | The dataframe as a row-major matrix of doubles, for handing data to ML systems. +'Left' if any column cannot be converted to doubles. -} toDoubleMatrix :: DataFrame -> Either DataFrameException (V.Vector (VU.Vector Double)) @@ -775,14 +796,8 @@ (fst (dataframeDimensions df)) (\i -> VU.generate (V.length m) (\j -> (m VG.! j) VG.! i)) -{- | Returns a dataframe as a two dimensional vector of ints. - -Converts all columns in the dataframe to int vectors and transposes them -into a row-major matrix representation. - -This is useful for handing data over into ML systems. - -Returns 'Left' with an error if any column cannot be converted to ints. +{- | The dataframe as a row-major matrix of ints, for handing data to ML systems. +'Left' if any column cannot be converted to ints. -} toIntMatrix :: DataFrame -> Either DataFrameException (V.Vector (VU.Vector Int)) toIntMatrix df = case V.foldl' @@ -833,10 +848,8 @@ Left e -> throw e Right (TColumn col) -> toVector col -{- | Retrieves a column as an unboxed vector of 'Int' values. - -Returns 'Left' with a 'DataFrameException' if the column cannot be converted to ints. -This may occur if the column contains non-numeric data or values outside the 'Int' range. +{- | A column as an unboxed vector of 'Int' values. +'Left' if the column cannot be converted to ints (non-numeric or out of range). -} columnAsIntVector :: (Columnable a, Num a) => @@ -850,10 +863,8 @@ Left e -> throw e Right (TColumn col) -> toIntVector col -{- | Retrieves a column as an unboxed vector of 'Double' values. - -Returns 'Left' with a 'DataFrameException' if the column cannot be converted to doubles. -This may occur if the column contains non-numeric data. +{- | A column as an unboxed vector of 'Double' values. +'Left' if the column cannot be converted to doubles (e.g. non-numeric data). -} columnAsDoubleVector :: (Columnable a, Num a) => @@ -870,10 +881,8 @@ Left e -> throw e Right (TColumn col) -> toDoubleVector col -{- | Retrieves a column as an unboxed vector of 'Float' values. - -Returns 'Left' with a 'DataFrameException' if the column cannot be converted to floats. -This may occur if the column contains non-numeric data. +{- | A column as an unboxed vector of 'Float' values. +'Left' if the column cannot be converted to floats (e.g. non-numeric data). -} columnAsFloatVector :: (Columnable a, Num a) => @@ -951,4 +960,11 @@ Just (UnboxedColumn Nothing (_ :: VU.Vector a)) -> UExpr (Col @a name) Just (PackedText (Just _) _) -> UExpr (Col @(Maybe T.Text) name) Just (PackedText Nothing _) -> UExpr (Col @T.Text name) + Just c@(MergedColumn _ _) -> case mergedHead c of + BoxedColumn (Just _) (_ :: V.Vector a) -> UExpr (Col @(Maybe a) name) + BoxedColumn Nothing (_ :: V.Vector a) -> UExpr (Col @a name) + _ -> + error $ + "showDerivedExpressions: merged column did not materialize boxed: " + ++ T.unpack name Nothing -> error $ "showDerivedExpressions: column not found: " ++ T.unpack name diff --git a/src/DataFrame/Operations/Inference.hs b/src/DataFrame/Operations/Inference.hs --- a/src/DataFrame/Operations/Inference.hs +++ b/src/DataFrame/Operations/Inference.hs @@ -1,13 +1,6 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE ScopedTypeVariables #-} -{- | Single-pass sampled inference lattice (Round-2 S2). One walk over -the sampled cells maintains a candidate mask {Bool, Int, Double, Date} -cleared by attempting the WS-B byte parsers; the assumption is the -highest-priority surviving candidate. Shared by every reader (audit T1), -together with the Int -> Double prefix promotion that replaces the -full-column retry chain (audit T2). --} module DataFrame.Operations.Inference ( DateFormat, ParsingAssumption (..), @@ -57,29 +50,17 @@ dateFormat (T.unpack s) -{- | The default @%Y-%m-%d@ format takes the WS-B byte-level fast path; -custom formats keep the reference 'readByteStringDate' parser. --} byteStringDateParser :: DateFormat -> BS.ByteString -> Maybe Day byteStringDateParser "%Y-%m-%d" = parseDateField byteStringDateParser fmt = readByteStringDate fmt {-# INLINE byteStringDateParser #-} -{- | 'DataFrame.Internal.Parsing.readInt' that rejects overflow instead of -wrapping. Fields of <= 18 chars cannot overflow and keep the Text-level -parse; longer (rare) fields take the exact byte-level parser, so an -overflowing cell demotes\/promotes instead of silently wrapping. --} readIntStrict :: T.Text -> Maybe Int readIntStrict t | T.length t <= 18 = readInt t | otherwise = parseIntField (TE.encodeUtf8 t) {-# INLINE readIntStrict #-} -{- | Candidate-mask priority, reproducing the documented fallback order: -an all-null sample makes no assumption; Int wins only when the Double -mask agrees (so mixed Int\/Double samples classify as Double). --} pickAssumption :: Bool -> Bool -> Bool -> Bool -> Bool -> ParsingAssumption pickAssumption seen b i d dt @@ -90,12 +71,6 @@ | dt = DateAssumption | otherwise = TextAssumption -{- | Classify a sample of decoded 'T.Text' cells ('Nothing' = null). -Bool\/Int\/Double candidates are tested with the WS-B byte parsers on -the UTF-8 bytes (strip-tolerant, overflow-rejecting); the Date -candidate keeps 'parseTimeOpt' so custom formats behave exactly as -before. The walk exits early once every candidate is cleared. --} makeParsingAssumption :: DateFormat -> V.Vector (Maybe T.Text) -> ParsingAssumption makeParsingAssumption dfmt cells = go 0 False True True True True @@ -137,13 +112,6 @@ (d && isJust (parseDoubleField bs)) (dt && isJust (dateP bs)) -{- | Fused Int pass with in-place promotion (audit T2): on the first -non-null cell that fails Int but parses as Double, the built Int prefix -is converted by a vector map and the pass continues as Double over the -retained raw cells. @Nothing@ = some cell parses as neither (the caller -demotes the column to Text). Null slots hold sentinels (0 \/ 0.0) -guarded by the bitmap, exactly like the unpromoted passes. --} promoteIntColumn :: forall src. (Int -> src -> Bool) -> diff --git a/src/DataFrame/Operations/Join.hs b/src/DataFrame/Operations/Join.hs --- a/src/DataFrame/Operations/Join.hs +++ b/src/DataFrame/Operations/Join.hs @@ -8,13 +8,39 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -module DataFrame.Operations.Join where +module DataFrame.Operations.Join ( + -- * Join types + JoinType (..), + -- * Joins + join, + innerJoin, + leftJoin, + rightJoin, + fullOuterJoin, + + -- * Low-level join kernels + + {- | Reused by the lazy executor and the parallel-join tests; not a + stable public API (candidates for a future @Join.Internal@ split). + -} + buildHashColumn, + buildCompactIndex, + hashProbeKernel, + hashInnerKernel, + hashLeftKernel, + innerKernel, + parInnerKernel, + parLeftKernel, + assembleInner, + assembleLeft, +) where + import Control.Applicative ((<|>)) import Control.Exception (throw) import Control.Monad (when) import Control.Monad.ST (ST, runST) -import Data.Bits ((.&.)) +import Data.Bits (popCount, unsafeShiftL, unsafeShiftR, (.&.), (.|.)) import qualified Data.Map.Strict as M import Data.Maybe (fromMaybe) import Data.STRef (newSTRef, readSTRef, writeSTRef) @@ -25,15 +51,28 @@ import qualified Data.Vector.Algorithms.Merge as VA import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM +import Data.Word (Word64) import DataFrame.Errors ( DataFrameException (ColumnsNotFoundException), ) -import DataFrame.Internal.Column as D +import DataFrame.Internal.Algorithms.Sort.Radix.Parallel (parSortByHash) +import DataFrame.Internal.Column as D ( + Column (BoxedColumn, UnboxedColumn), + atIndicesStable, + columnTypeString, + fromUnboxedVector, + fromVector, + gatherWithSentinel, + isPackedText, + materializePacked, + mkMergedColumns, + ) +import DataFrame.Internal.Column.Bitmap (bitmapTestBit) import DataFrame.Internal.DataFrame as D -import DataFrame.Internal.ParRadixSort (parSortByHash) import DataFrame.Operations.Aggregation as D import DataFrame.Operations.Core as D -import DataFrame.Operations.JoinPar ( +import DataFrame.Operations.Join.Parallel ( + ProbeTable (..), parInnerProbe, parLeftProbe, shouldParallelizeJoin, @@ -54,60 +93,86 @@ join :: JoinType -> [T.Text] -> - DataFrame -> -- Right hand side - DataFrame -> -- Left hand side + DataFrame -> + DataFrame -> DataFrame join INNER xs right = innerJoin xs right join LEFT xs right = leftJoin xs right join RIGHT xs right = rightJoin xs right join FULL_OUTER xs right = fullOuterJoin xs right -{- | Row-count threshold for the build side. -When the build side exceeds this, sort-merge join is used instead of the -single-threaded hash join. The parallel chunked-probe hash join -('parInnerKernel' \/ 'parLeftKernel') is preferred above this size when more -than one capability is available (see 'shouldParallelizeJoin'): partitioning -the probe over cores beats sort-merge for the large build sides measured here -(e.g. the 2M-row build in the join pipeline). Sort-merge remains the -single-threaded fallback. +{- | Build-side row count above which the single-threaded hash join gives way +to sort-merge (single-threaded) or the parallel chunked-probe hash join when +multiple capabilities are available (see 'shouldParallelizeJoin'). -} joinStrategyThreshold :: Int joinStrategyThreshold = 500_000 -{- | A compact index mapping hash values to contiguous slices of -original row indices. All indices live in a single unboxed vector -(@ciSortedIndices@, sorted by hash). The lookup table is an open-addressing -linear-probe hash table held in three parallel unboxed vectors keyed by hash: -@ciKeys@ holds the hash at each slot, @ciStarts@ the run offset into -@ciSortedIndices@ (@-1@ marks an empty slot, since real offsets are @>= 0@), -and @ciLens@ the run length. @ciMask@ is @tableSize - 1@ (table size is a -power of two), used to map a hash to its home slot. +{- | Maps hash values to contiguous slices of original row indices. Indices +live in one hash-sorted vector; an open-addressing linear-probe table (keyed +by hash) records each run's offset and length. @-1@ starts mark empty slots. -} data CompactIndex = CompactIndex { ciSortedIndices :: {-# UNPACK #-} !(VU.Vector Int) , ciKeys :: {-# UNPACK #-} !(VU.Vector Int) - , ciStarts :: {-# UNPACK #-} !(VU.Vector Int) - , ciLens :: {-# UNPACK #-} !(VU.Vector Int) + , ciRuns :: {-# UNPACK #-} !(VU.Vector Int) + {- ^ @(start, len)@ of each run packed 32/32 into one 'Int'; @-1@ = empty + slot. One vector instead of separate starts\/lens halves the table's + memory (a 1e8-row build side is a 2^28-slot table: 2.1GB instead of + 4.3GB) and saves a random cache-line read per lookup hit. + -} , ciMask :: {-# UNPACK #-} !Int } +-- | Pack a run's @(start, len)@ 32\/32 into one non-negative 'Int'. +ciPackRun :: Int -> Int -> Int +ciPackRun !start !len = (start `unsafeShiftL` 32) .|. len +{-# INLINE ciPackRun #-} + +-- | Start field of a packed run. +ciRunStart :: Int -> Int +ciRunStart !w = w `unsafeShiftR` 32 +{-# INLINE ciRunStart #-} + +-- | Length field of a packed run. +ciRunLen :: Int -> Int +ciRunLen !w = w .&. 0xFFFF_FFFF +{-# INLINE ciRunLen #-} + +{- | Home slot of a hash: the top @log2 cap@ bits of a Fibonacci multiply. +The row hash's final FxHash step is a multiply, which leaves its LOW bits +poorly diffused (text keys cluster catastrophically: contiguous-pileup chains +in the hundreds), so the table must never index by @h .&. mask@ directly. +@shift@ is @64 - log2 cap@. +-} +ciSlot :: Int -> Int -> Int +ciSlot !shift !h = + fromIntegral + ((fromIntegral h * (0x9E37_79B9_7F4A_7C15 :: Word64)) `unsafeShiftR` shift) +{-# INLINE ciSlot #-} + +-- | @64 - log2 cap@ for a table with slot mask @mask@ (@cap@ a power of two). +ciShiftFor :: Int -> Int +ciShiftFor !mask = 64 - popCount mask +{-# INLINE ciShiftFor #-} + {- | Look up a hash in the open-addressing table. Returns @(start, len)@ of the matching run, or @(-1, 0)@ on a miss. -} ciLookup :: CompactIndex -> Int -> (Int, Int) -ciLookup ci !h = go (h .&. mask) +ciLookup ci !h = go (ciSlot shift h) where !mask = ciMask ci + !shift = ciShiftFor mask !keys = ciKeys ci - !starts = ciStarts ci - !lens = ciLens ci + !runs = ciRuns ci go !slot = - let !s = starts `VU.unsafeIndex` slot - in if s < 0 + let !w = runs `VU.unsafeIndex` slot + in if w < 0 then (-1, 0) else if keys `VU.unsafeIndex` slot == h - then (s, lens `VU.unsafeIndex` slot) + then (ciRunStart w, ciRunLen w) else go ((slot + 1) .&. mask) {-# INLINE ciLookup #-} @@ -121,45 +186,43 @@ | p > n = p | otherwise = go (p * 2) -{- | Build a compact index from a vector of row hashes. -Sorts @(hash, originalIndex)@ pairs by hash, scans for contiguous runs, then -inserts each run into an open-addressing linear-probe table sized to keep the -load factor under ~0.5. +{- | Build a compact index from a vector of row hashes: sort by hash, scan for +contiguous runs, insert each run into an open-addressing table. Capacity is +sized for the worst case (every row distinct) so building never resizes. -} buildCompactIndex :: VU.Vector Int -> CompactIndex +buildCompactIndex hashes + | VU.length hashes > 0x7FFF_FFFF = + error + "buildCompactIndex: build side exceeds 2^31 rows (packed run fields are 32-bit)" buildCompactIndex hashes = let n = VU.length hashes (sortedHashes, sortedIndices) = parSortByHash n hashes - -- Worst case every row is a distinct group; 2*n+1 keeps the table - -- sparse even then. Capacity is independent of the actual group count - -- so building stays a single pass with no resize. !cap = nextPow2Above (2 * n) !mask = cap - 1 - (keys, starts, lens) = runST $ do + !shift = ciShiftFor mask + (keys, runs) = runST $ do mKeys <- VUM.unsafeNew cap - mStarts <- VUM.replicate cap (-1) - mLens <- VUM.unsafeNew cap + mRuns <- VUM.replicate cap (-1) let insert !i | i >= n = return () | otherwise = do let !h = sortedHashes `VU.unsafeIndex` i !end = findGroupEnd sortedHashes h (i + 1) n - probe h (h .&. mask) i (end - i) + probe h (ciSlot shift h) i (end - i) insert end probe !h !slot !start !len = do - s <- VUM.unsafeRead mStarts slot - if s < 0 + w <- VUM.unsafeRead mRuns slot + if w < 0 then do VUM.unsafeWrite mKeys slot h - VUM.unsafeWrite mStarts slot start - VUM.unsafeWrite mLens slot len + VUM.unsafeWrite mRuns slot (ciPackRun start len) else probe h ((slot + 1) .&. mask) start len insert 0 - (,,) + (,) <$> VU.unsafeFreeze mKeys - <*> VU.unsafeFreeze mStarts - <*> VU.unsafeFreeze mLens - in CompactIndex sortedIndices keys starts lens mask + <*> VU.unsafeFreeze mRuns + in CompactIndex sortedIndices keys runs mask -- | Find the end of a contiguous run of equal values starting at @j@. findGroupEnd :: VU.Vector Int -> Int -> Int -> Int -> Int @@ -169,9 +232,8 @@ | otherwise = j {-# INLINE findGroupEnd #-} -{- | Sort a hash vector, returning sorted hashes and corresponding original indices. -Sorts an index array using hash values as the comparison key, avoiding the -intermediate pair vector used by the naive zip-then-sort approach. +{- | Sort a hash vector, returning sorted hashes and their original indices. +Sorts an index array keyed by hash, avoiding an intermediate pair vector. -} sortWithIndices :: VU.Vector Int -> (VU.Vector Int, VU.Vector Int) sortWithIndices hashes = runST $ do @@ -254,6 +316,9 @@ | D.null right || D.null left = D.empty | otherwise = innerJoinNonEmpty cs left right +-- Build the hash index on the smaller side and probe with the larger one; +-- the probe side is partitioned across cores, so probing the larger maximises +-- the parallel win. innerJoinNonEmpty :: [T.Text] -> DataFrame -> DataFrame -> DataFrame innerJoinNonEmpty cs left right = let @@ -267,16 +332,7 @@ rightHashes = D.computeRowHashes rightKeyIdxs right buildRows = min leftRows rightRows - -- Probe with the larger side, build on the smaller. The probe side - -- drives parallelism (it is partitioned across cores), so probing the - -- larger side maximises the win. probeRows = max leftRows rightRows - -- Parallelize the probe either when the build is large enough to spill - -- cache (the original sort-merge regime) or when the build is small - -- (cache-resident, read-only) but the probe is huge: partitioning a - -- ~1e7-row probe over cores wins even against a tiny hot build (the - -- medium-factor lever). Both routes use the same bit-identical - -- 'parInnerKernel'. useParallel = shouldParallelizeJoin probeRows buildRows || shouldParallelizeSmallBuildProbe probeRows @@ -284,10 +340,8 @@ | buildRows > joinStrategyThreshold && not useParallel = sortMergeInnerKernel leftHashes rightHashes | rightRows <= leftRows = - -- Build on right (smaller or equal), probe with left innerKernel useParallel leftHashes rightHashes | otherwise = - -- Build on left (smaller), probe with right, swap result let (!rIxs, !lIxs) = innerKernel useParallel rightHashes leftHashes in (lIxs, rIxs) in @@ -304,20 +358,30 @@ {-# INLINE innerKernel #-} {- | Parallel inner-join kernel: build the 'CompactIndex' on @buildHashes@ once, -then probe @probeHashes@ in parallel. Bit-for-bit identical output to -'hashInnerKernel' (probe-row order preserved). Runs the IO probe via -'unsafePerformIO'; the computation is pure (no observable effects, fixed result -for fixed inputs). +then probe @probeHashes@ in parallel. Output is bit-for-bit identical to +'hashInnerKernel' (probe-row order preserved). -} parInnerKernel :: VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) parInnerKernel probeHashes buildHashes = let !ci = buildCompactIndex buildHashes (pf, bf) = - unsafePerformIO (parInnerProbe (ciSortedIndices ci) (ciLookup ci) probeHashes) + unsafePerformIO (parInnerProbe (ciProbeTable ci) probeHashes) in (pf, bf) {-# NOINLINE parInnerKernel #-} +{- | The raw table fields of a 'CompactIndex', in the closure-free shape the +parallel probe kernels consume. +-} +ciProbeTable :: CompactIndex -> ProbeTable +ciProbeTable ci = + ProbeTable + { ptSorted = ciSortedIndices ci + , ptKeys = ciKeys ci + , ptRuns = ciRuns ci + , ptMask = ciMask ci + } + -- | Compute hashes for the given key column names in a DataFrame. buildHashColumn :: [T.Text] -> DataFrame -> VU.Vector Int buildHashColumn keys df = @@ -325,10 +389,9 @@ keyIdxs = validatedKeyColIndices "buildHashColumn" csSet df in D.computeRowHashes keyIdxs df -{- | Probe one batch of rows against a pre-built 'CompactIndex'. -Returns @(probeExpandedIxs, buildExpandedIxs)@. -Unlike 'hashInnerKernel', does not build the index (it is pre-built once) -and has no cross-product row guard — the caller controls probe batch size. +{- | Probe one batch of rows against a pre-built 'CompactIndex', returning +@(probeExpandedIxs, buildExpandedIxs)@. Unlike 'hashInnerKernel' it neither +builds the index nor guards cross-product size — the caller sizes batches. -} hashProbeKernel :: -- | Built once from the full right\/build side. @@ -393,16 +456,13 @@ <*> VU.unsafeFreeze (VUM.slice 0 total bv) in (VU.force pFrozen, VU.force bFrozen) -{- | Hash-based inner join kernel. -Builds compact index on @buildHashes@ (second arg), probes with -@probeHashes@ (first arg). -Returns @(probeExpandedIndices, buildExpandedIndices)@. -Uses a dynamically growing output buffer to avoid pre-allocating the full -cross-product size (which can be astronomically large for low-cardinality keys). +{- | Hash-based inner join kernel: build the index on @buildHashes@, probe with +@probeHashes@, return @(probeExpandedIndices, buildExpandedIndices)@. Grows its +output buffer dynamically rather than pre-allocating the full cross product. -} -{- | Maximum number of output rows allowed from a join kernel. -Exceeding this limit indicates a cross-product explosion (e.g. low-cardinality keys). +{- | Output-row ceiling for a join kernel; exceeding it signals a cross-product +explosion (e.g. low-cardinality keys). -} maxJoinOutputRows :: Int maxJoinOutputRows = 500_000_000 @@ -472,15 +532,11 @@ (,) <$> VU.unsafeFreeze (VUM.slice 0 total pv) <*> VU.unsafeFreeze (VUM.slice 0 total bv) - in -- VU.force copies the slice into a compact array, releasing the oversized - -- backing buffer allocated by the doubling strategy. - (VU.force pFrozen, VU.force bFrozen) + in (VU.force pFrozen, VU.force bFrozen) -{- | Sort-merge inner join kernel. -Sorts both sides by hash, walks in lockstep. -Returns @(leftExpandedIndices, rightExpandedIndices)@. -Uses a dynamically growing output buffer instead of a two-pass count-then-allocate -strategy, which OOMs when low-cardinality keys produce large cross products. +{- | Sort-merge inner join kernel: sort both sides by hash, walk in lockstep, +return @(leftExpandedIndices, rightExpandedIndices)@. Grows its output buffer +dynamically so low-cardinality cross products don't OOM. -} sortMergeInnerKernel :: VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) @@ -565,9 +621,7 @@ (,) <$> VU.unsafeFreeze (VUM.slice 0 total lv) <*> VU.unsafeFreeze (VUM.slice 0 total rv) - in -- VU.force copies the slice into a compact array, releasing the oversized - -- backing buffer allocated by the doubling strategy. - (VU.force lFrozen, VU.force rFrozen) + in (VU.force lFrozen, VU.force rFrozen) -- | Assemble the result DataFrame for an inner join from expanded index vectors. assembleInner :: @@ -582,7 +636,6 @@ leftColSet = S.fromList (D.columnNames left) rightColNames = D.columnNames right - -- Pre-expand every column once expandedLeftCols = VB.map (D.atIndicesStable leftIxs) (D.columns left) expandedRightCols = VB.map (D.atIndicesStable rightIxs) (D.columns right) @@ -594,7 +647,6 @@ idx <- M.lookup name (D.columnIndices right) return (expandedRightCols `VB.unsafeIndex` idx) - -- Base DataFrame: all left columns, expanded baseDf = left { columns = expandedLeftCols @@ -607,15 +659,15 @@ in D.fold ( \name df -> if S.member name csSet - then df -- Key column already present from left side + then df else if S.member name leftColSet - then -- Overlapping non-key column: merge with These + then insertIfPresent name - (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name) + (D.mkMergedColumns <$> getExpandedLeft name <*> getExpandedRight name) df - else -- Right-only column + else insertIfPresent name (getExpandedRight name) df ) rightColNames @@ -657,6 +709,8 @@ | D.null left || D.nRows left == 0 = D.empty | otherwise = leftJoinNonEmpty callPoint cs left right +-- The right side is always the build side; the left is probed (and drives +-- parallelism). Unmatched right rows are marked with a @-1@ sentinel. leftJoinNonEmpty :: T.Text -> [T.Text] -> DataFrame -> DataFrame -> DataFrame leftJoinNonEmpty callPoint cs left right = let @@ -669,10 +723,6 @@ rightHashes = D.computeRowHashes rightKeyIdxs right leftRows = fst (D.dimensions left) - -- Right is always the build side for left join; left is the probe side - -- (and drives parallelism). Parallelize either in the large-build regime - -- or when the build is small but the probe (left) is huge: the - -- read-only shared index is probed across cores with no synchronization. useParallel = shouldParallelizeJoin leftRows rightRows || shouldParallelizeSmallBuildProbe leftRows @@ -684,7 +734,6 @@ | otherwise = hashLeftKernel leftHashes rightHashes in - -- rightIxs uses -1 as sentinel for "no match" assembleLeft csSet left right leftIxs rightIxs {- | Parallel left-join kernel: build the 'CompactIndex' on @rightHashes@ once, @@ -695,14 +744,12 @@ VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) parLeftKernel leftHashes rightHashes = let !ci = buildCompactIndex rightHashes - in unsafePerformIO (parLeftProbe (ciSortedIndices ci) (ciLookup ci) leftHashes) + in unsafePerformIO (parLeftProbe (ciProbeTable ci) leftHashes) {-# NOINLINE parLeftKernel #-} -{- | Hash-based left join kernel. -Returns @(leftExpandedIndices, rightExpandedIndices)@ where -right indices use @-1@ as sentinel for unmatched rows. -Uses a dynamically growing output buffer to avoid pre-allocating the full -cross-product size (which can be astronomically large for low-cardinality keys). +{- | Hash-based left join kernel, returning @(leftExpandedIndices, +rightExpandedIndices)@ with @-1@ marking unmatched right rows. Grows its output +buffer dynamically rather than pre-allocating the full cross product. -} hashLeftKernel :: VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) @@ -769,10 +816,9 @@ <$> VU.unsafeFreeze (VUM.slice 0 total lv) <*> VU.unsafeFreeze (VUM.slice 0 total rv) -{- | Sort-merge left join kernel. -Returns @(leftExpandedIndices, rightExpandedIndices)@ with @-1@ sentinel. -Uses a dynamically growing output buffer instead of a two-pass count-then-allocate -strategy, which OOMs when low-cardinality keys produce large cross products. +{- | Sort-merge left join kernel, returning @(leftExpandedIndices, +rightExpandedIndices)@ with a @-1@ sentinel. Grows its output buffer +dynamically so low-cardinality cross products don't OOM. -} sortMergeLeftKernel :: VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) @@ -917,7 +963,7 @@ then insertIfPresent name - (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name) + (D.mkMergedColumns <$> getExpandedLeft name <*> getExpandedRight name) df else insertIfPresent name (getExpandedRight name) df ) @@ -967,14 +1013,12 @@ leftHashes = D.computeRowHashes leftKeyIdxs left rightHashes = D.computeRowHashes rightKeyIdxs right - -- Both sides can have nulls in full outer (leftIxs, rightIxs) | max leftRows rightRows > joinStrategyThreshold = sortMergeFullOuterKernel leftHashes rightHashes | otherwise = hashFullOuterKernel leftHashes rightHashes in - -- Both index vectors use -1 as sentinel assembleFullOuter csSet left right leftIxs rightIxs {- | Hash-based full outer join kernel. @@ -989,36 +1033,32 @@ leftSI = ciSortedIndices leftCI rightSI = ciSortedIndices rightCI leftKeys = ciKeys leftCI - leftStarts = ciStarts leftCI - leftLens = ciLens leftCI + leftRuns = ciRuns leftCI rightKeys = ciKeys rightCI - rightStarts = ciStarts rightCI - rightLens = ciLens rightCI - !leftCap = VU.length leftStarts - !rightCap = VU.length rightStarts + rightRuns = ciRuns rightCI + !leftCap = VU.length leftRuns + !rightCap = VU.length rightRuns - -- Count: matched + left-only + right-only. Iterate over occupied slots - -- (ciStarts /= -1) in each table, cross-referencing the other via ciLookup. let countLeft !slot !acc | slot >= leftCap = acc | otherwise = - let !lStart = leftStarts `VU.unsafeIndex` slot - in if lStart < 0 + let !lw = leftRuns `VU.unsafeIndex` slot + in if lw < 0 then countLeft (slot + 1) acc else let !h = leftKeys `VU.unsafeIndex` slot - !ll = leftLens `VU.unsafeIndex` slot + !ll = ciRunLen lw (!rs, !rl) = ciLookup rightCI h in countLeft (slot + 1) (acc + if rs < 0 then ll else ll * rl) countRightOnly !slot !acc | slot >= rightCap = acc | otherwise = - let !rStart = rightStarts `VU.unsafeIndex` slot - in if rStart < 0 + let !rw = rightRuns `VU.unsafeIndex` slot + in if rw < 0 then countRightOnly (slot + 1) acc else let !h = rightKeys `VU.unsafeIndex` slot - !rl = rightLens `VU.unsafeIndex` slot + !rl = ciRunLen rw (!ls, _) = ciLookup leftCI h in countRightOnly (slot + 1) (acc + if ls < 0 then rl else 0) !leftPlusMatched = countLeft 0 0 @@ -1029,16 +1069,16 @@ rv <- VUM.unsafeNew totalCount posRef <- newSTRef (0 :: Int) - -- Fill matched + left-only (iterate left slots) let fillLeft !slot | slot >= leftCap = return () | otherwise = do - let !lStart = leftStarts `VU.unsafeIndex` slot - if lStart < 0 + let !lw = leftRuns `VU.unsafeIndex` slot + if lw < 0 then fillLeft (slot + 1) else do let !h = leftKeys `VU.unsafeIndex` slot - !lLen = leftLens `VU.unsafeIndex` slot + !lStart = ciRunStart lw + !lLen = ciRunLen lw (!rStart, !rLen) = ciLookup rightCI h !p <- readSTRef posRef if rStart < 0 @@ -1066,16 +1106,16 @@ fillLeft (slot + 1) fillLeft 0 - -- Fill right-only (iterate right slots not in left) let fillRightOnly !slot | slot >= rightCap = return () | otherwise = do - let !rStart = rightStarts `VU.unsafeIndex` slot - if rStart < 0 + let !rw = rightRuns `VU.unsafeIndex` slot + if rw < 0 then fillRightOnly (slot + 1) else do let !h = rightKeys `VU.unsafeIndex` slot - !rLen = rightLens `VU.unsafeIndex` slot + !rStart = ciRunStart rw + !rLen = ciRunLen rw (!ls, _) = ciLookup leftCI h if ls >= 0 then fillRightOnly (slot + 1) @@ -1105,7 +1145,6 @@ !leftN = VU.length leftHashes !rightN = VU.length rightHashes - -- Pass 1: count let countLoop !li !ri !c | li >= leftN && ri >= rightN = c | li >= leftN = c + (rightN - ri) @@ -1121,7 +1160,6 @@ !rh = rightSH `VU.unsafeIndex` ri !totalRows = countLoop 0 0 0 - -- Pass 2: fill lv <- VUM.unsafeNew totalRows rv <- VUM.unsafeNew totalRows @@ -1164,9 +1202,8 @@ fill 0 0 0 (,) <$> VU.unsafeFreeze lv <*> VU.unsafeFreeze rv -{- | Assemble the result DataFrame for a full outer join. -Both index vectors use @-1@ sentinel; all columns gathered via -'gatherWithSentinel'. Key columns are coalesced (first non-null wins). +{- | Assemble the result DataFrame for a full outer join. Both index vectors +use a @-1@ sentinel; key columns are coalesced (first non-null wins). -} assembleFullOuter :: S.Set T.Text -> @@ -1201,8 +1238,6 @@ insertIfPresent _ Nothing df = df insertIfPresent name (Just c) df = D.insertColumn name c df - -- Coalesce two nullable columns: take first non-Nothing per row, - -- producing a non-optional column. coalesceKeyColumn :: Column -> Column -> Column coalesceKeyColumn l r | D.isPackedText l || D.isPackedText r = @@ -1252,16 +1287,15 @@ in D.fold ( \name df -> if S.member name csSet - then -- Key column: coalesce left and right - case (getExpandedLeft name, getExpandedRight name) of - (Just lc, Just rc) -> D.insertColumn name (coalesceKeyColumn lc rc) df - _ -> df + then case (getExpandedLeft name, getExpandedRight name) of + (Just lc, Just rc) -> D.insertColumn name (coalesceKeyColumn lc rc) df + _ -> df else if S.member name leftColSet then insertIfPresent name - (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name) + (D.mkMergedColumns <$> getExpandedLeft name <*> getExpandedRight name) df else insertIfPresent name (getExpandedRight name) df ) diff --git a/src/DataFrame/Operations/Join/Parallel.hs b/src/DataFrame/Operations/Join/Parallel.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Join/Parallel.hs @@ -0,0 +1,222 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ScopedTypeVariables #-} + +{- | +Parallel chunked-probe join kernels. The build side is indexed once into a +shared, read-only 'CompactIndex' (open-addressing, from +"DataFrame.Operations.Join"); the probe side is split into @caps@ /contiguous/ +row ranges and probed in parallel by 'forkFinally' workers (no sparks). Each worker +makes two passes over its range — a count pass to size its slice, then a fill +pass — and writes into the single shared output buffers at a precomputed +prefix-sum offset. Because ranges are contiguous and laid out in range order, +the produced @(probeIxs, buildIxs)@ vectors are /bit-for-bit identical/ to the +sequential 'hashInnerKernel' \/ 'hashLeftKernel': probe rows appear in original +order and, within a probe row, build matches in @ciSortedIndices@ order. +-} +module DataFrame.Operations.Join.Parallel ( + ProbeTable (..), + parInnerProbe, + parLeftProbe, + shouldParallelizeJoin, + shouldParallelizeSmallBuildProbe, + parJoinThreshold, + parBuildThreshold, + parProbeThreshold, +) where + +import Control.Concurrent (getNumCapabilities) +import Data.Bits (popCount, unsafeShiftR, (.&.)) +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM +import Data.Word (Word64) +import DataFrame.Internal.Control.Concurrent (capabilities, forkJoin_) + +{- | Below this many probe rows the fork/coordination overhead is not worth it; +the caller uses its sequential 'ST' kernel instead. +-} +parJoinThreshold :: Int +parJoinThreshold = 200000 + +{- | Below this many build rows the shared 'CompactIndex' is small and hot, so +the sequential hash probe is already memory-bound-fast and the fork overhead +loses (measured: a 1e4-row Text-key build probed by 1e7 rows is /slower/ in +parallel). Parallelism only pays once the build index is large enough to spill +cache — exactly the regime where sort-merge used to be chosen. +-} +parBuildThreshold :: Int +parBuildThreshold = 500000 + +{- | Whether a join should take the parallel probe path: more than one +capability, a probe side of at least 'parJoinThreshold' rows, and a build side +of at least 'parBuildThreshold' rows (a small/hot index is faster probed +sequentially). +-} +shouldParallelizeJoin :: Int -> Int -> Bool +shouldParallelizeJoin probeRows buildRows = + probeRows >= parJoinThreshold + && buildRows >= parBuildThreshold + && capabilities > 1 +{-# NOINLINE shouldParallelizeJoin #-} + +{- | Above this many probe rows the probe-side row hashing and table lookups +dominate the join, so partitioning the probe across cores wins even when the +build side is small and cache-resident (the regime 'shouldParallelizeJoin' +deliberately leaves sequential). Sized at 1e6: below it the per-question gain is +swamped by 'forkFinally'/coordination overhead, so small/medium-inner joins stay +sequential (measured). This is the small-build large-probe lever closing the +medium-factor 1e7 join (1e7 probe x ~1e4 build). +-} +parProbeThreshold :: Int +parProbeThreshold = 1000000 + +{- | Whether a /small-build/ join (build below 'parBuildThreshold', so radix +partitioning / sort-merge is not used) should take the parallel probe path: a +very large probe side (at least 'parProbeThreshold') and more than one +capability. The shared build index is read-only across threads, so probing it in +parallel needs no synchronization. Independent of build size on purpose: the +build is already tiny; the cost is the 1e7-row probe hashing, which parallelizes +cleanly. +-} +shouldParallelizeSmallBuildProbe :: Int -> Bool +shouldParallelizeSmallBuildProbe probeRows = + probeRows >= parProbeThreshold + && capabilities > 1 +{-# NOINLINE shouldParallelizeSmallBuildProbe #-} + +{- | A read-only view of the build-side index needed by the probe: the raw +open-addressing table vectors of the @CompactIndex@ (which lives in +"DataFrame.Operations.Join"; passing the fields avoids an import cycle). +Passing concrete vectors instead of a lookup closure keeps the per-row probe +loop free of unknown calls and boxed-tuple allocation — the lookup is inlined +into the count and fill loops. +-} +data ProbeTable = ProbeTable + { ptSorted :: !(VU.Vector Int) + , ptKeys :: !(VU.Vector Int) + , ptRuns :: !(VU.Vector Int) + -- ^ @(start, len)@ packed 32\/32 per slot; @-1@ = empty (see @ciRuns@). + , ptMask :: {-# UNPACK #-} !Int + } + +{- | Parallel inner-join probe. @parInnerProbe table probeHashes@ returns +@(probeIxs, buildIxs)@ identical to a sequential probe of the same index. The +build index must already be constructed from the build side. +-} +parInnerProbe :: + ProbeTable -> + VU.Vector Int -> + IO (VU.Vector Int, VU.Vector Int) +parInnerProbe = runProbe False + +{- | Parallel left-join probe. Like 'parInnerProbe' but every probe row emits at +least one output row; unmatched rows carry a @-1@ sentinel in the build column. +-} +parLeftProbe :: + ProbeTable -> + VU.Vector Int -> + IO (VU.Vector Int, VU.Vector Int) +parLeftProbe = runProbe True + +{- | Shared two-pass parallel probe. @keepUnmatched@ selects left- vs +inner-join semantics. Splits @[0, probeN)@ into @caps@ contiguous ranges, counts +each range's output, prefix-sums to global offsets, then fills the single output +buffers in parallel. +-} +runProbe :: + Bool -> + ProbeTable -> + VU.Vector Int -> + IO (VU.Vector Int, VU.Vector Int) +runProbe keepUnmatched pt probeHashes = do + caps <- getNumCapabilities + let !probeN = VU.length probeHashes + !nChunks = max 1 (min caps probeN) + !sorted = ptSorted pt + !keys = ptKeys pt + !runs = ptRuns pt + !mask = ptMask pt + !shift = 64 - popCount mask + -- Packed (start,len) run for hash @h@, or -1 on a miss. Home slot is + -- the top log2(cap) bits of a Fibonacci multiply (the row hash's low + -- bits are poorly diffused); must match the build-side ciSlot exactly. + findRun !h = + go + ( fromIntegral + ((fromIntegral h * (0x9E3779B97F4A7C15 :: Word64)) `unsafeShiftR` shift) + ) + where + go !slot = + let !w = runs `VU.unsafeIndex` slot + in if w < 0 + then -1 + else + if keys `VU.unsafeIndex` slot == h + then w + else go ((slot + 1) .&. mask) + chunkBounds k = (lo, hi) + where + !lo = (probeN * k) `div` nChunks + !hi = (probeN * (k + 1)) `div` nChunks + -- Count pass: output rows produced by probe range [lo, hi). + countRange !lo !hi = + let go !i !acc + | i >= hi = acc + | otherwise = + let !w = findRun (VU.unsafeIndex probeHashes i) + in if w < 0 + then go (i + 1) (if keepUnmatched then acc + 1 else acc) + else go (i + 1) (acc + (w .&. 0xFFFFFFFF)) + in go lo 0 + chunkCounts <- VUM.new (nChunks + 1) + forkRanges nChunks $ \k -> + let (lo, hi) = chunkBounds k + in VUM.unsafeWrite chunkCounts k (countRange lo hi) + -- Exclusive prefix sum -> per-chunk global start offsets; total at [nChunks]. + let scan !k !acc + | k > nChunks = pure acc + | otherwise = do + c <- if k < nChunks then VUM.unsafeRead chunkCounts k else pure 0 + VUM.unsafeWrite chunkCounts k acc + scan (k + 1) (acc + c) + !total <- scan 0 0 + pv <- VUM.unsafeNew (max 1 total) + bv <- VUM.unsafeNew (max 1 total) + -- Fill pass: each chunk writes from its prefix-sum offset. + offs <- VU.unsafeFreeze chunkCounts + forkRanges nChunks $ \k -> do + let (lo, hi) = chunkBounds k + !base = VU.unsafeIndex offs k + fill !i !p + | i >= hi = pure () + | otherwise = do + let !w = findRun (VU.unsafeIndex probeHashes i) + if w < 0 + then + if keepUnmatched + then do + VUM.unsafeWrite pv p i + VUM.unsafeWrite bv p (-1) + fill (i + 1) (p + 1) + else fill (i + 1) p + else do + let !start = w `unsafeShiftR` 32 + !len = w .&. 0xFFFFFFFF + writeMatch !j !q + | j >= len = pure () + | otherwise = do + VUM.unsafeWrite pv q i + VUM.unsafeWrite bv q (VU.unsafeIndex sorted (start + j)) + writeMatch (j + 1) (q + 1) + writeMatch 0 p + fill (i + 1) (p + len) + fill lo base + pf <- VU.unsafeFreeze (VUM.slice 0 total pv) + bf <- VU.unsafeFreeze (VUM.slice 0 total bv) + pure (pf, bf) + +{- | Run @body k@ for @k@ in @[0, nChunks)@, one chunk per task, on @nChunks@ +forked threads; rethrow the first failure. Chunk @k@ is owned by exactly one +thread, so concurrent writes to disjoint output regions are race-free. +-} +forkRanges :: Int -> (Int -> IO ()) -> IO () +forkRanges nChunks body = forkJoin_ (map body [0 .. nChunks - 1]) diff --git a/src/DataFrame/Operations/JoinPar.hs b/src/DataFrame/Operations/JoinPar.hs deleted file mode 100644 --- a/src/DataFrame/Operations/JoinPar.hs +++ /dev/null @@ -1,215 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ScopedTypeVariables #-} - -{- | -Parallel chunked-probe join kernels. The build side is indexed once into a -shared, read-only 'CompactIndex' (open-addressing, from -"DataFrame.Operations.Join"); the probe side is split into @caps@ /contiguous/ -row ranges and probed in parallel by 'forkIO' workers (no sparks). Each worker -makes two passes over its range — a count pass to size its slice, then a fill -pass — and writes into the single shared output buffers at a precomputed -prefix-sum offset. Because ranges are contiguous and laid out in range order, -the produced @(probeIxs, buildIxs)@ vectors are /bit-for-bit identical/ to the -sequential 'hashInnerKernel' \/ 'hashLeftKernel': probe rows appear in original -order and, within a probe row, build matches in @ciSortedIndices@ order. - -This is the parallel==sequential correctness gate (see -@tests/Operations/ParallelJoin.hs@). A sequential fallback is used when there is -a single capability or the probe side is below 'parJoinThreshold'; the caller -('innerJoin' \/ 'leftJoin') decides via 'shouldParallelizeJoin'. --} -module DataFrame.Operations.JoinPar ( - parInnerProbe, - parLeftProbe, - shouldParallelizeJoin, - shouldParallelizeSmallBuildProbe, - parJoinThreshold, - parBuildThreshold, - parProbeThreshold, -) where - -import Control.Concurrent (forkIO, getNumCapabilities) -import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar) -import Control.Exception (SomeException, throwIO, try) -import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Unboxed.Mutable as VUM -import System.IO.Unsafe (unsafePerformIO) - -{- | Below this many probe rows the fork/coordination overhead is not worth it; -the caller uses its sequential 'ST' kernel instead. --} -parJoinThreshold :: Int -parJoinThreshold = 200000 - -{- | Below this many build rows the shared 'CompactIndex' is small and hot, so -the sequential hash probe is already memory-bound-fast and the fork overhead -loses (measured: a 1e4-row Text-key build probed by 1e7 rows is /slower/ in -parallel). Parallelism only pays once the build index is large enough to spill -cache — exactly the regime where sort-merge used to be chosen. --} -parBuildThreshold :: Int -parBuildThreshold = 500000 - -{- | Whether a join should take the parallel probe path: more than one -capability, a probe side of at least 'parJoinThreshold' rows, and a build side -of at least 'parBuildThreshold' rows (a small/hot index is faster probed -sequentially). --} -shouldParallelizeJoin :: Int -> Int -> Bool -shouldParallelizeJoin probeRows buildRows = - probeRows >= parJoinThreshold - && buildRows >= parBuildThreshold - && capabilities > 1 -{-# NOINLINE shouldParallelizeJoin #-} - -{- | Above this many probe rows the probe-side row hashing and table lookups -dominate the join, so partitioning the probe across cores wins even when the -build side is small and cache-resident (the regime 'shouldParallelizeJoin' -deliberately leaves sequential). Sized at 1e6: below it the per-question gain is -swamped by 'forkIO'/coordination overhead, so small/medium-inner joins stay -sequential (measured). This is the small-build large-probe lever closing the -medium-factor 1e7 join (1e7 probe x ~1e4 build). --} -parProbeThreshold :: Int -parProbeThreshold = 1000000 - -{- | Whether a /small-build/ join (build below 'parBuildThreshold', so radix -partitioning / sort-merge is not used) should take the parallel probe path: a -very large probe side (at least 'parProbeThreshold') and more than one -capability. The shared build index is read-only across threads, so probing it in -parallel needs no synchronization. Independent of build size on purpose: the -build is already tiny; the cost is the 1e7-row probe hashing, which parallelizes -cleanly. --} -shouldParallelizeSmallBuildProbe :: Int -> Bool -shouldParallelizeSmallBuildProbe probeRows = - probeRows >= parProbeThreshold - && capabilities > 1 -{-# NOINLINE shouldParallelizeSmallBuildProbe #-} - -capabilities :: Int -capabilities = unsafePerformIO getNumCapabilities -{-# NOINLINE capabilities #-} - -{- | A read-only view of the build-side index needed by the probe: the lookup -returns @(start, len)@ of the matching run in @sortedIndices@, or @(-1, 0)@ on a -miss. Passed in by the caller so this module need not depend on the -'CompactIndex' record directly. --} -data ProbeIndex = ProbeIndex - { piSorted :: !(VU.Vector Int) - , piLookup :: !(Int -> (Int, Int)) - } - -{- | Parallel inner-join probe. @parInnerProbe sortedIdxs lookup probeHashes@ -returns @(probeIxs, buildIxs)@ identical to a sequential probe of the same -index. The build index must already be constructed from the build side. --} -parInnerProbe :: - VU.Vector Int -> - (Int -> (Int, Int)) -> - VU.Vector Int -> - IO (VU.Vector Int, VU.Vector Int) -parInnerProbe sortedIdxs lookupFn = - runProbe False (ProbeIndex sortedIdxs lookupFn) - -{- | Parallel left-join probe. Like 'parInnerProbe' but every probe row emits at -least one output row; unmatched rows carry a @-1@ sentinel in the build column. --} -parLeftProbe :: - VU.Vector Int -> - (Int -> (Int, Int)) -> - VU.Vector Int -> - IO (VU.Vector Int, VU.Vector Int) -parLeftProbe sortedIdxs lookupFn = - runProbe True (ProbeIndex sortedIdxs lookupFn) - -{- | Shared two-pass parallel probe. @keepUnmatched@ selects left- vs -inner-join semantics. Splits @[0, probeN)@ into @caps@ contiguous ranges, counts -each range's output, prefix-sums to global offsets, then fills the single output -buffers in parallel. --} -runProbe :: - Bool -> - ProbeIndex -> - VU.Vector Int -> - IO (VU.Vector Int, VU.Vector Int) -runProbe keepUnmatched pidx probeHashes = do - caps <- getNumCapabilities - let !probeN = VU.length probeHashes - !nChunks = max 1 (min caps probeN) - !sorted = piSorted pidx - !lookupFn = piLookup pidx - chunkBounds k = (lo, hi) - where - !lo = (probeN * k) `div` nChunks - !hi = (probeN * (k + 1)) `div` nChunks - -- Count pass: output rows produced by probe range [lo, hi). - countRange !lo !hi = - let go !i !acc - | i >= hi = acc - | otherwise = - let (!start, !len) = lookupFn (VU.unsafeIndex probeHashes i) - in if start < 0 - then go (i + 1) (if keepUnmatched then acc + 1 else acc) - else go (i + 1) (acc + len) - in go lo 0 - chunkCounts <- VUM.new (nChunks + 1) - forkRanges nChunks $ \k -> - let (lo, hi) = chunkBounds k - in VUM.unsafeWrite chunkCounts k (countRange lo hi) - -- Exclusive prefix sum -> per-chunk global start offsets; total at [nChunks]. - let scan !k !acc - | k > nChunks = pure acc - | otherwise = do - c <- if k < nChunks then VUM.unsafeRead chunkCounts k else pure 0 - VUM.unsafeWrite chunkCounts k acc - scan (k + 1) (acc + c) - !total <- scan 0 0 - pv <- VUM.unsafeNew (max 1 total) - bv <- VUM.unsafeNew (max 1 total) - -- Fill pass: each chunk writes from its prefix-sum offset. - offs <- VU.unsafeFreeze chunkCounts - forkRanges nChunks $ \k -> do - let (lo, hi) = chunkBounds k - !base = VU.unsafeIndex offs k - fill !i !p - | i >= hi = pure () - | otherwise = do - let (!start, !len) = lookupFn (VU.unsafeIndex probeHashes i) - if start < 0 - then - if keepUnmatched - then do - VUM.unsafeWrite pv p i - VUM.unsafeWrite bv p (-1) - fill (i + 1) (p + 1) - else fill (i + 1) p - else do - let writeMatch !j !q - | j >= len = pure () - | otherwise = do - VUM.unsafeWrite pv q i - VUM.unsafeWrite bv q (VU.unsafeIndex sorted (start + j)) - writeMatch (j + 1) (q + 1) - writeMatch 0 p - fill (i + 1) (p + len) - fill lo base - pf <- VU.unsafeFreeze (VUM.slice 0 total pv) - bf <- VU.unsafeFreeze (VUM.slice 0 total bv) - pure (pf, bf) - -{- | Run @body k@ for @k@ in @[0, nChunks)@, one chunk per task, on @nChunks@ -forked threads; rethrow the first failure. Chunk @k@ is owned by exactly one -thread, so concurrent writes to disjoint output regions are race-free. --} -forkRanges :: Int -> (Int -> IO ()) -> IO () -forkRanges nChunks body = do - vars <- mapM spawn [0 .. nChunks - 1] - results <- mapM takeMVar vars - mapM_ (either (throwIO :: SomeException -> IO ()) pure) results - where - spawn k = do - var <- newEmptyMVar - _ <- forkIO (try (body k) >>= putMVar var) - pure var diff --git a/src/DataFrame/Operations/Merge.hs b/src/DataFrame/Operations/Merge.hs --- a/src/DataFrame/Operations/Merge.hs +++ b/src/DataFrame/Operations/Merge.hs @@ -1,7 +1,11 @@ {-# LANGUAGE InstanceSigs #-} {-# OPTIONS_GHC -Wno-orphans #-} -module DataFrame.Operations.Merge where +module DataFrame.Operations.Merge ( + -- * Horizontal (side-by-side) merge + (|||), + -- Also exports the orphan @Semigroup@/@Monoid DataFrame@ instances. +) where import qualified Data.List as L import qualified Data.Text as T @@ -38,9 +42,6 @@ case optB of Nothing -> case optA of Nothing -> - -- N.B. this case should never happen, because we're dealing with columns coming from - -- union of column names of both dataframes. Nothing + Nothing would mean column - -- wasn't in either dataframe, which shouldn't happen D.insertColumn name (D.fromList ([] :: [T.Text])) df Just a'' -> D.insertColumn name (D.expandColumn sumRows a'') df @@ -48,7 +49,7 @@ Nothing -> D.insertColumn name (D.leftExpandColumn sumRows b'') df Just a'' -> - let concatedColumns = D.concatColumnsEither a'' b'' + let concatedColumns = D.mappendColumnsEither a'' b'' in D.insertColumn name concatedColumns df result = L.foldl' (addColumns a b) D.empty (D.columnNames a `L.union` D.columnNames b) in diff --git a/src/DataFrame/Operations/Permutation.hs b/src/DataFrame/Operations/Permutation.hs --- a/src/DataFrame/Operations/Permutation.hs +++ b/src/DataFrame/Operations/Permutation.hs @@ -5,8 +5,16 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -module DataFrame.Operations.Permutation where +module DataFrame.Operations.Permutation ( + -- * Sorting + SortOrder (..), + sortBy, + -- * Shuffling + shuffle, + shuffledIndices, +) where + import qualified Data.List as L import qualified Data.Text as T import qualified Data.Vector as V @@ -20,14 +28,19 @@ import Data.Type.Equality (testEquality, (:~:) (Refl)) import Data.Vector.Internal.Check (HasCallStack) import DataFrame.Errors (DataFrameException (..)) -import DataFrame.Internal.Column (Column (..), Columnable, atIndicesStable) +import DataFrame.Internal.Column ( + Column (..), + Columnable, + atIndicesStable, + materializeMerged, + ) +import DataFrame.Internal.Data.PackedText (packedSlice, sliceCmpBytes) import DataFrame.Internal.DataFrame ( DataFrame (..), columnNames, unsafeGetColumn, ) import DataFrame.Internal.Expression (Expr (Col), getColumns) -import DataFrame.Internal.PackedText (packedSlice, sliceCmpBytes) import DataFrame.Operations.Core (dimensions) import DataFrame.Operations.Transformations (derive) import System.Random (Random (randomR), RandomGen) @@ -124,6 +137,11 @@ (aj, oj, lj) = packedSlice p j in sliceCmpBytes ai oi li aj oj lj Nothing -> \_ _ -> EQ + c@(MergedColumn _ _) -> case materializeMerged c of + BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> \i j -> compare (v `V.unsafeIndex` i) (v `V.unsafeIndex` j) + Nothing -> \_ _ -> EQ + _ -> \_ _ -> EQ sortOrderComparator (Desc (Col name :: Expr a)) df = case unsafeGetColumn name df of BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of @@ -138,6 +156,11 @@ (aj, oj, lj) = packedSlice p j in sliceCmpBytes aj oj lj ai oi li Nothing -> \_ _ -> EQ + c@(MergedColumn _ _) -> case materializeMerged c of + BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> \i j -> compare (v `V.unsafeIndex` j) (v `V.unsafeIndex` i) + Nothing -> \_ _ -> EQ + _ -> \_ _ -> EQ sortOrderComparator _ _ = error "Sorting on compound column" -- | Sort row indices using a comparator function. @@ -168,14 +191,12 @@ shuffleVec :: (RandomGen g) => g -> VU.Vector Int shuffleVec g = runST $ do vm <- VUM.generate k id - let (n, nGen) = randomR (1, k - 1) g - go vm n nGen + go vm (k - 1) g VU.unsafeFreeze vm - go _v (-1) _ = pure () - go _v 0 _ = pure () - go v maxInd gen = + go _v i _ | i <= 0 = pure () + go v i gen = let - (n, nextGen) = randomR (1, maxInd) gen + (j, nextGen) = randomR (0, i) gen in - VUM.swap v 0 n *> go (VUM.tail v) (maxInd - 1) nextGen + VUM.swap v i j *> go v (i - 1) nextGen diff --git a/src/DataFrame/Operations/SetOps.hs b/src/DataFrame/Operations/SetOps.hs --- a/src/DataFrame/Operations/SetOps.hs +++ b/src/DataFrame/Operations/SetOps.hs @@ -1,19 +1,6 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ScopedTypeVariables #-} -{- | -Module : DataFrame.Operations.SetOps -Description : Set-theoretic ("topos") row operations. - -These treat a 'DataFrame' as a /set/ of rows and implement the subobject -lattice from relational algebra: 'union', 'intersect', 'difference', and -'symmetricDifference'. Every result is deduplicated, so each operation has the -schema-preserving shape @DataFrame -> DataFrame -> DataFrame@. - -Row equality is the same hash-based notion used by 'distinct' (see -"DataFrame.Operations.Aggregation"), so these operations and 'distinct' agree -on what "the same row" means. Both inputs are expected to share a schema; the -typed layer ('DataFrame.Typed') enforces that statically. --} module DataFrame.Operations.SetOps ( union, intersect, @@ -79,8 +66,8 @@ chosen = [ VU.head members | k <- [0 .. nGroups - 1] - , let s = VU.unsafeIndex offs k - e = VU.unsafeIndex offs (k + 1) + , let !s = VU.unsafeIndex offs k + !e = VU.unsafeIndex offs (k + 1) members = VU.slice s (e - s) vis inLeft = VU.any (< leftRows) members inRight = VU.any (>= leftRows) members diff --git a/src/DataFrame/Operations/Statistics.hs b/src/DataFrame/Operations/Statistics.hs --- a/src/DataFrame/Operations/Statistics.hs +++ b/src/DataFrame/Operations/Statistics.hs @@ -9,8 +9,30 @@ {-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_GHC -Wno-orphans #-} -module DataFrame.Operations.Statistics where +module DataFrame.Operations.Statistics ( + -- * Summaries + summarize, + frequencies, + -- * Aggregate statistics + mean, + meanMaybe, + median, + medianMaybe, + percentile, + genericPercentile, + standardDeviation, + skewness, + variance, + interQuartileRange, + correlation, + sum, + + -- * Imputation + imputeWith, + -- Also exports the orphan @ImputeOp (Maybe b)@ instance. +) where + import qualified Data.List as L import qualified Data.Map as M import qualified Data.Text as T @@ -18,7 +40,8 @@ import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU -import Prelude as P +import Prelude hiding (sum) +import qualified Prelude as P import Control.Exception (throw) import Data.Function ((&)) @@ -34,11 +57,10 @@ getColumn, ) import DataFrame.Internal.Expression +import DataFrame.Internal.Expression.Operators.Nullable (BaseType) import DataFrame.Internal.Interpreter -import DataFrame.Internal.Nullable (BaseType) import DataFrame.Internal.Row (showValue, toAny) import DataFrame.Internal.Statistics -import DataFrame.Internal.Types import DataFrame.Operations.Core import DataFrame.Operations.Subset (filterJust) import DataFrame.Operations.Transformations (ImputeOp (..), imputeCore) @@ -227,7 +249,7 @@ Nothing -> throw $ ColumnsNotFoundException [name] "_getColumnAsDouble" (M.keys $ columnIndices df) - _ -> Nothing -- Return a type mismatch error here. + _ -> Nothing {-# INLINE _getColumnAsDouble #-} optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double @@ -249,6 +271,7 @@ Just Refl -> VG.sum column Nothing -> 0 Just (PackedText _ _) -> 0 + Just (MergedColumn _ _) -> 0 -- matches the old eager These column (type never Num) sum expr df = case interpret df expr of Left e -> throw e Right (TColumn xs) -> case toVector @a @V.Vector xs of @@ -304,7 +327,7 @@ if all (== h) (toList @b value) then imputeCore col h df else error "Impute expression returned more than one value" - runImputeWith _ _ df = df + runImputeWith _ expr _ = throw (NonColumnReferenceException (T.pack (show expr))) imputeWith :: forall a. diff --git a/src/DataFrame/Operations/Subset.hs b/src/DataFrame/Operations/Subset.hs --- a/src/DataFrame/Operations/Subset.hs +++ b/src/DataFrame/Operations/Subset.hs @@ -5,12 +5,53 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -module DataFrame.Operations.Subset where +module DataFrame.Operations.Subset ( + -- * Row slicing + take, + takeLast, + drop, + dropLast, + range, + cube, + selectRows, + -- * Filtering + filter, + filterBy, + filterWhere, + filterJust, + filterNothing, + filterAllJust, + filterAllNothing, + + -- * Column selection + select, + selectBy, + exclude, + SelectionCriteria, + byName, + byProperty, + byNameProperty, + byNameRange, + byIndexRange, + + -- * Sampling & splitting + SplittableGen, + sample, + randomSplit, + kFolds, + stratifiedSample, + stratifiedSplit, + + -- * Label helpers (exported for satellite packages) + rowsAtIndices, +) where + import qualified Data.List as L import qualified Data.Map as M import qualified Data.Set as S @@ -33,26 +74,54 @@ DataFrameException (..), TypeErrorContext (..), ) -import DataFrame.Internal.Column +import DataFrame.Expression.Operators ( + col, + name, + (.<.), + (.<=.), + (.>.), + (.>=.), + ) +import DataFrame.Internal.Column ( + Column (..), + Columnable, + TypedColumn (TColumn), + atIndicesStable, + columnToTextVec, + findIndices, + hasMissing, + materializeMerged, + materializePacked, + mkRandom, + sliceColumn, + takeColumn, + takeLastColumn, + ) +import DataFrame.Internal.Column.Bitmap (bitmapTestBit) import DataFrame.Internal.DataFrame ( DataFrame (..), columnNames, + dataframeDimensions, derivingExpressions, empty, getColumn, insertColumn, unsafeGetColumn, ) -import DataFrame.Internal.Expression -import DataFrame.Internal.Interpreter -import DataFrame.Internal.PackedText (packedIndexText, packedLength) +import DataFrame.Internal.Expression (Expr (Col, Lit), normalize) +import DataFrame.Internal.Interpreter (Ctx (..), eval, interpret, materialize) import DataFrame.Operations.Core () import DataFrame.Operations.Merge () import DataFrame.Operations.Transformations (apply) -import DataFrame.Operators -import System.Random -import Type.Reflection -import Prelude hiding (filter, take) +import System.Random (RandomGen, SplitGen (..)) +import Type.Reflection ( + eqTypeRep, + typeRep, + pattern App, + type (:~:) (Refl), + type (:~~:) (HRefl), + ) +import Prelude hiding (drop, filter, take) #if MIN_VERSION_random(1,3,0) type SplittableGen g = (SplitGen g, RandomGen g) @@ -107,12 +176,16 @@ range :: (Int, Int) -> DataFrame -> DataFrame range (start, end) d = d - { columns = V.map (sliceColumn (clip start 0 r) n') (columns d) + { columns = V.map (sliceColumn start' n') (columns d) , dataframeDimensions = (n', c) } where (r, c) = dataframeDimensions d - n' = clip (end - start) 0 r + start' = clip start 0 r + -- Clamp both endpoints before subtracting: end - start' on an unclamped + -- end wraps for very negative values and reopens the range. + end' = clip end start' r + n' = end' - start' clip :: Int -> Int -> Int -> Int clip n left right = min right $ max n left @@ -137,8 +210,9 @@ ColumnsNotFoundException [filterColumnName] "filter" (M.keys $ columnIndices df) Just c@(PackedText _ _) -> filter e condition (insertColumn filterColumnName (materializePacked c) df) + Just c@(MergedColumn _ _) -> + filter e condition (insertColumn filterColumnName (materializeMerged c) df) Just _col@(BoxedColumn bm (column :: V.Vector b)) -> - -- Check direct type match first, then try Maybe b match for nullable columns case testEquality (typeRep @a) (typeRep @b) of Just Refl -> filterByVector filterColumnName column condition df Nothing -> case (bm, typeRep @a) of @@ -476,26 +550,6 @@ in map (exclude [name cRand]) (go (folds - 1) withRand) --- | Convert any Column to a vector of Text labels (one per row). -columnToTextVec :: Column -> V.Vector T.Text -columnToTextVec (BoxedColumn bm (col' :: V.Vector a)) = - case bm of - Nothing -> case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> col' - Nothing -> V.map (T.pack . show) col' - Just bitmap -> - V.imap (\i x -> if bitmapTestBit bitmap i then T.pack (show x) else "null") col' -columnToTextVec (UnboxedColumn bm col') = - case bm of - Nothing -> V.map (T.pack . show) (V.convert col') - Just bitmap -> - V.generate (VU.length col') $ \i -> - if bitmapTestBit bitmap i then T.pack (show (col' VU.! i)) else "null" -columnToTextVec (PackedText bm p) = - V.generate (packedLength p) $ \i -> case bm of - Just bitmap | not (bitmapTestBit bitmap i) -> "null" - _ -> packedIndexText p i - -- | Build a map from stringified label to row indices. groupByIndices :: Column -> M.Map T.Text (VU.Vector Int) groupByIndices col' = @@ -530,7 +584,9 @@ stratifiedSample gen p strataCol df = let col' = case strataCol of Col colName -> unsafeGetColumn colName df - _ -> unwrapTypedColumn (either throw id (interpret @a df strataCol)) + _ -> either throw id $ do + v <- eval (FlatCtx df) strataCol + pure $ materialize @a (fst (dataframeDimensions df)) v groups = M.elems (groupByIndices col') go _ [] = mempty go g (ixs : rest) = @@ -554,7 +610,9 @@ stratifiedSplit gen p strataCol df = let col' = case strataCol of Col colName -> unsafeGetColumn colName df - _ -> unwrapTypedColumn (either throw id (interpret @a df strataCol)) + _ -> either throw id $ do + v <- eval (FlatCtx df) strataCol + pure $ materialize @a (fst (dataframeDimensions df)) v groups = M.elems (groupByIndices col') go _ [] = (mempty, mempty) go g (ixs : rest) = diff --git a/src/DataFrame/Operations/Transformations.hs b/src/DataFrame/Operations/Transformations.hs --- a/src/DataFrame/Operations/Transformations.hs +++ b/src/DataFrame/Operations/Transformations.hs @@ -1,5 +1,5 @@ {-# LANGUAGE ConstrainedClassMethods #-} -{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} @@ -7,11 +7,31 @@ {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE UndecidableSuperClasses #-} -module DataFrame.Operations.Transformations where +module DataFrame.Operations.Transformations ( + -- * Apply + apply, + safeApply, + applyMany, + applyInt, + applyDouble, + applyWhere, + applyAtIndex, + -- * Derive + derive, + deriveWithExpr, + deriveMany, + + -- * Impute + ImputeOp (..), + impute, + imputeCore, +) where + import qualified Data.List as L import qualified Data.Map as M import qualified Data.Text as T @@ -21,8 +41,10 @@ import Data.Maybe import DataFrame.Errors (DataFrameException (..), TypeErrorContext (..)) import DataFrame.Internal.Column ( + Column, Columnable, TypedColumn (..), + columnTypeString, hasMissing, ifoldrColumn, imapColumn, @@ -30,9 +52,11 @@ ) import DataFrame.Internal.DataFrame (DataFrame (..), getColumn, insertColumn) import DataFrame.Internal.Expression +import DataFrame.Internal.Expression.Operators.Nullable (BaseType) import DataFrame.Internal.Interpreter -import DataFrame.Internal.Nullable (BaseType) import DataFrame.Operations.Core +import GHC.TypeLits (ErrorMessage (..), TypeError) +import Type.Reflection (typeRep) -- | O(k) Apply a function to a given column in a dataframe. apply :: @@ -199,7 +223,9 @@ Left e -> throw e Right column' -> insertColumn columnName column' df --- | Core impute implementation for nullable columns. Silently no-ops on non-nullable columns. +{- | Core impute implementation for nullable columns. A column with no null +bitmap cannot hold a missing value, so imputing it is a mistake, not a no-op. +-} imputeCore :: forall b. (Columnable b) => @@ -215,9 +241,24 @@ Left (TypeMismatchException context) -> throw $ TypeMismatchException (context{callingFunctionName = Just "impute"}) Left exception -> throw exception Right res -> res - _ -> df -imputeCore _ _ df = df + Just col -> + throw + (nonNullableColumnError columnName ("Maybe " ++ show (typeRep @b)) col) +imputeCore expr _ _ = throw (NonColumnReferenceException (T.pack (show expr))) +{- | 'impute' was handed a column that carries no null bitmap, so there is +nothing it could replace. +-} +nonNullableColumnError :: T.Text -> String -> Column -> DataFrameException +nonNullableColumnError columnName wanted col = + TypeMismatchException @() @() + MkTypeErrorContext + { userType = Left wanted + , expectedType = Left (columnTypeString col) + , errorColumnName = Just (T.unpack columnName) + , callingFunctionName = Just "impute" + } + class (Columnable a) => ImputeOp a where runImpute :: Expr a -> BaseType a -> DataFrame -> DataFrame runImputeWith :: @@ -227,12 +268,24 @@ DataFrame -> DataFrame -instance {-# OVERLAPPABLE #-} (Columnable a) => ImputeOp a where - runImpute _ _ df = df - runImputeWith _ _ df = df +instance + {-# OVERLAPPABLE #-} + ( Columnable a + , TypeError + ( 'Text "impute needs a nullable column, but was given Expr " + ':<>: 'ShowType a + ':$$: 'Text "Try F.col @(Maybe " + ':<>: 'ShowType a + ':<>: 'Text ") instead." + ) + ) => + ImputeOp a + where + runImpute = errorWithoutStackTrace "impute: unreachable" + runImputeWith = errorWithoutStackTrace "impute: unreachable" {- | Replace all instances of `Nothing` in a column with the given value. -When the column is already non-nullable, this is a silent no-op. +Throws when the column carries no nulls to replace. -} impute :: forall a. diff --git a/src/DataFrame/Operations/Typing.hs b/src/DataFrame/Operations/Typing.hs --- a/src/DataFrame/Operations/Typing.hs +++ b/src/DataFrame/Operations/Typing.hs @@ -29,33 +29,29 @@ import Data.Time import Data.Type.Equality (TestEquality (..)) import DataFrame.Internal.Column ( - Bitmap, Column (..), Columnable, - bitmapTestBit, ensureOptional, finalizeParseResult, fromVector, materializePacked, ) +import DataFrame.Internal.Column.Bitmap (Bitmap, bitmapTestBit) import DataFrame.Internal.DataFrame ( DataFrame (..), insertColumn, unsafeGetColumn, ) import DataFrame.Internal.Parsing -import DataFrame.Internal.Schema import DataFrame.Operations.Core () import DataFrame.Operations.Inference +import DataFrame.Schema import Text.Read import Type.Reflection -{- | How parse failures are surfaced in the resulting column. - -* 'NoSafeRead' — strict parsing: failures throw (via 'read'). -* 'MaybeRead' — failures become 'Nothing'; columns are wrapped as @Maybe a@. -* 'EitherRead' — failures become @Left rawText@; columns are wrapped as - @Either Text a@, preserving the original input so callers can inspect it. +{- | How parse failures are surfaced: 'NoSafeRead' throws, 'MaybeRead' yields +@Nothing@ (column wrapped @Maybe a@), 'EitherRead' yields @Left rawText@ +(column wrapped @Either Text a@, preserving the original input). -} data SafeReadMode = NoSafeRead @@ -70,16 +66,10 @@ , sampleSize :: Int -- ^ Number of rows to inspect when inferring a column's type (0 = all rows). , parseSafe :: SafeReadMode - {- ^ Default 'SafeReadMode' applied to every column that does not have an - entry in 'parseSafeOverrides'. 'NoSafeRead' only treats empty strings as - missing; 'MaybeRead' additionally treats 'missingValues' and nullish - strings as @Nothing@; 'EitherRead' wraps the resulting column as - @Either Text a@ with the raw input preserved on failure. - -} + -- ^ Default 'SafeReadMode' for columns without a 'parseSafeOverrides' entry. , parseSafeOverrides :: [(T.Text, SafeReadMode)] - {- ^ Per-column overrides. When a column name is present here, its value - takes precedence over 'parseSafe'. Typical use: strict IDs - (@NoSafeRead@) alongside lenient fields (@MaybeRead@/@EitherRead@). + {- ^ Per-column overrides taking precedence over 'parseSafe' — e.g. strict + IDs (@NoSafeRead@) alongside lenient fields (@MaybeRead@/@EitherRead@). -} , parseDateFormat :: DateFormat -- ^ Date format string as accepted by "Data.Time.Format" (e.g. @\"%Y-%m-%d\"@). @@ -108,7 +98,6 @@ parseDefaults :: ParseOptions -> DataFrame -> DataFrame parseDefaults opts df = df{columns = V.imap forCol (columns df)} where - -- Index -> column name: reverse the columnIndices map once. nameAt = let inverted = M.fromList [(i, n) | (n, i) <- M.toList (columnIndices df)] in \i -> M.findWithDefault "" i inverted @@ -144,10 +133,6 @@ let isNull = case parseSafe opts of NoSafeRead -> T.null _ -> isNullishOrMissing (missingValues opts) - -- `examples` is small (≤ sampleSize, default 100), so the - -- Maybe-wrap allocation here is ignorable. The full-column - -- equivalent (`asMaybeText = V.map ... cols`) has been removed: - -- handlers now walk `cols` directly with `isNull`. examples = V.map (classify isNull) (V.take (sampleSize opts) cols) dfmt = parseDateFormat opts assumption = makeParsingAssumption dfmt examples @@ -165,10 +150,9 @@ where classify p t = if p t then Nothing else Just t -{- | For 'EitherRead' mode: take the chosen parsing assumption and produce an -@Either Text a@ column. Successful parses become @Right@; any row that fails -to parse as the chosen type (including null/missing cells) becomes @Left@ -carrying the raw input text verbatim. +{- | For 'EitherRead' mode: parse under the chosen assumption into an +@Either Text a@ column. Successful parses become @Right@; failures (including +null/missing cells) become @Left@ carrying the raw input verbatim. -} handleEitherAssumption :: DateFormat -> ParsingAssumption -> V.Vector T.Text -> Column @@ -177,9 +161,6 @@ IntAssumption -> fromVector (V.map (toEither readInt) raw) DoubleAssumption -> fromVector (V.map (toEither readDouble) raw) DateAssumption -> fromVector (V.map (toEither (parseTimeOpt dfmt)) raw) - -- TextAssumption and NoAssumption degenerate to Either Text Text; treat - -- empty strings as Left "" so the convention (Left = missing/failure) stays - -- consistent across column types. TextAssumption -> fromVector (V.map textToEither raw) NoAssumption -> fromVector (V.map textToEither raw) where @@ -234,10 +215,9 @@ (parseUnboxedColumnWithPred False isNull readBool cols) (handleTextAssumption isNull cols) -{- | Int columns: one fused pass with in-place Int -> Double promotion -('promoteIntColumn'); a cell parsing as neither demotes to Text over -the retained raw cells. 'readIntStrict' rejects overflow so a huge -integer promotes to its true 'Double' value instead of wrapping. +{- | Int columns: one fused pass with in-place Int -> Double promotion; a cell +parsing as neither demotes the column to Text. 'readIntStrict' rejects overflow +so a huge integer promotes to 'Double' rather than wrapping. -} handleIntAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column handleIntAssumption isNull cols = @@ -251,9 +231,8 @@ (parseUnboxedColumnWithPred 0 isNull readDouble cols) (handleTextAssumption isNull cols) -{- | Text columns: no parse, just null-marking. When the whole column -is non-null we return a plain 'V.Vector T.Text'; otherwise we emit a -@V.Vector (Maybe T.Text)@ the same shape the old code produced. +{- | Text columns: no parse, just null-marking. An all-non-null column stays a +plain @V.Vector T.Text@; otherwise it becomes @V.Vector (Maybe T.Text)@. -} handleTextAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column handleTextAssumption isNull cols @@ -262,19 +241,15 @@ (V.map (\t -> if isNull t then Nothing else Just t) cols) | otherwise = fromVector cols -{- | Date: single parse pass, boxed because 'Day' is not unboxable. -Bails to 'handleTextAssumption' the moment a non-null cell fails to -parse as a 'Day'. Still avoids the outer @V.Vector (Maybe T.Text)@ -allocation — we walk @cols@ directly with @isNull@. +{- | Date: single boxed parse pass ('Day' is not unboxable). Bails to +'handleTextAssumption' the moment a non-null cell fails to parse as a 'Day'. +A column with no nulls keeps type 'Day' rather than 'Maybe Day'. -} handleDateAssumption :: DateFormat -> (T.Text -> Bool) -> V.Vector T.Text -> Column handleDateAssumption dateFormat isNull cols = case parseBoxedMaybeColumn isNull (parseTimeOpt dateFormat) cols of Just (anyNull, vec) - -- `vec :: V.Vector (Maybe Day)`. If no nulls, strip the - -- outer 'Maybe' (every cell is guaranteed 'Just') so the - -- column type stays 'Day' rather than becoming 'Maybe Day'. | anyNull -> fromVector vec | otherwise -> fromVector (V.mapMaybe id vec) Nothing -> handleTextAssumption isNull cols @@ -304,11 +279,11 @@ Nothing -> return Nothing loop 0 False +-- Reached only when the sample was all-null: try each concrete type in turn, +-- falling back to Text. A column with no nulls keeps type 'Day', not 'Maybe Day'. handleNoAssumption :: DateFormat -> (T.Text -> Bool) -> V.Vector T.Text -> Column handleNoAssumption dateFormat isNull cols - -- Only reached when the 100-row sample was all-null. Try each - -- concrete type in turn; fall back to Text otherwise. | V.all isNull cols = fromVector (V.map (const (Nothing :: Maybe T.Text)) cols) | Just (mbm, vec) <- parseUnboxedColumnWithPred False isNull readBool cols = @@ -319,19 +294,12 @@ UnboxedColumn mbm vec | otherwise = case parseBoxedMaybeColumn isNull (parseTimeOpt dateFormat) cols of Just (anyNull, vec) - -- `vec :: V.Vector (Maybe Day)`. If no nulls, strip the - -- outer 'Maybe' (every cell is guaranteed 'Just') so the - -- column type stays 'Day' rather than becoming 'Maybe Day'. | anyNull -> fromVector vec | otherwise -> fromVector (V.mapMaybe id vec) Nothing -> handleTextAssumption isNull cols -{- | Predicate matching what 'parseSafe == NoSafeRead' previously used: -only empty strings are treated as missing. - -We still expose 'convertNullish' \/ 'convertOnlyEmpty' below because -other parts of the library reference them, but neither is used by -'parseFromExamples' any longer. +{- | True for nullish or explicitly-listed missing strings. ('convertNullish' +and 'convertOnlyEmpty' below are kept only for external callers.) -} isNullishOrMissing :: [T.Text] -> T.Text -> Bool isNullishOrMissing missing v = isNullish v || v `elem` missing @@ -344,7 +312,7 @@ unsafeParseTime :: DateFormat -> T.Text -> Day unsafeParseTime dateFormat s = - parseTimeOrError {- Accept leading/trailing whitespace -} + parseTimeOrError True defaultTimeLocale dateFormat @@ -361,10 +329,8 @@ hasSameConstructor Nothing Nothing = True hasSameConstructor _ _ = False -{- | Re-type columns of a 'DataFrame' according to the supplied schema map. -The caller provides a @resolveMode@ function that maps a column name to its -'SafeReadMode' — typically built from a global default plus an overrides map -via 'effectiveSafeRead'. +{- | Re-type columns of a 'DataFrame' according to a schema map. @resolveMode@ +maps a column name to its 'SafeReadMode' (typically via 'effectiveSafeRead'). -} parseWithTypes :: (T.Text -> SafeReadMode) -> @@ -379,9 +345,6 @@ df ts where - -- \| Re-parse a plain (non-Maybe, non-Either) target type according to the - -- 'SafeReadMode'. @toStr@ converts column elements to a 'String' ready for - -- 'Read'. plainType :: forall a b. (Columnable a, Read a) => @@ -392,8 +355,6 @@ EitherRead -> fromVector (V.map ((readEitherRaw @a) . toStr) col) asType :: SafeReadMode -> SchemaType -> Column -> Column - -- A raw CSV string column may arrive as PackedText; decode to boxed Text - -- so the re-parse arms below can read the cells. asType mode st c@(PackedText _ _) = asType mode st (materializePacked c) asType mode (SType (_ :: P.Proxy a)) c@(BoxedColumn _ (col :: V.Vector b)) = case typeRep @a of App t1 _t2 -> case eqTypeRep t1 (typeRep @Maybe) of diff --git a/src/DataFrame/Typed/Access.hs b/src/DataFrame/Typed/Access.hs --- a/src/DataFrame/Typed/Access.hs +++ b/src/DataFrame/Typed/Access.hs @@ -10,18 +10,35 @@ -- * Typed column access columnAsVector, columnAsList, + + -- * Numeric vector extraction + columnAsIntVector, + columnAsDoubleVector, + columnAsFloatVector, + columnAsUnboxedVector, + + -- * Matrix extraction + toDoubleMatrix, + toFloatMatrix, + toIntMatrix, ) where import Control.Exception (throw) import Data.Proxy (Proxy (..)) import qualified Data.Text as T import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU import GHC.TypeLits (KnownSymbol, symbolVal) import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.Expression (Expr (Col)) import qualified DataFrame.Operations.Core as D -import DataFrame.Typed.Schema (AssertPresent, SafeLookup) +import DataFrame.Typed.Schema ( + AllColumnsReal, + AssertPresent, + AssertRealColumn, + SafeLookup, + ) import DataFrame.Typed.Types (TypedDataFrame (..)) {- | Retrieve a column as a boxed 'Vector', with the type determined by @@ -53,3 +70,94 @@ D.columnAsList (Col @a colName) df where colName = T.pack (symbolVal (Proxy @name)) + +{- | Retrieve a column coerced to an unboxed 'Int' vector, named by type +application. The column must exist and be numeric — both are compile-time +checks via 'SafeLookup', so this is total (no 'Either', no runtime throw). +-} +columnAsIntVector :: + forall name cols a. + ( KnownSymbol name + , a ~ SafeLookup name cols + , Columnable a + , AssertRealColumn "columnAsIntVector" name a + , Real a + , VU.Unbox a + , AssertPresent name cols + ) => + TypedDataFrame cols -> VU.Vector Int +columnAsIntVector (TDF df) = either throw id (D.columnAsIntVector (Col @a colName) df) + where + colName = T.pack (symbolVal (Proxy @name)) + +-- | Retrieve a column coerced to an unboxed 'Double' vector. See 'columnAsIntVector'. +columnAsDoubleVector :: + forall name cols a. + ( KnownSymbol name + , a ~ SafeLookup name cols + , Columnable a + , AssertRealColumn "columnAsDoubleVector" name a + , Real a + , VU.Unbox a + , AssertPresent name cols + ) => + TypedDataFrame cols -> VU.Vector Double +columnAsDoubleVector (TDF df) = + either throw id (D.columnAsDoubleVector (Col @a colName) df) + where + colName = T.pack (symbolVal (Proxy @name)) + +-- | Retrieve a column coerced to an unboxed 'Float' vector. See 'columnAsIntVector'. +columnAsFloatVector :: + forall name cols a. + ( KnownSymbol name + , a ~ SafeLookup name cols + , Columnable a + , AssertRealColumn "columnAsFloatVector" name a + , Real a + , VU.Unbox a + , AssertPresent name cols + ) => + TypedDataFrame cols -> VU.Vector Float +columnAsFloatVector (TDF df) = + either throw id (D.columnAsFloatVector (Col @a colName) df) + where + colName = T.pack (symbolVal (Proxy @name)) + +{- | Retrieve a column as an unboxed vector of its own element type. The column +must exist and be unboxable — both compile-time checks, so this is total. +-} +columnAsUnboxedVector :: + forall name cols a. + ( KnownSymbol name + , a ~ SafeLookup name cols + , Columnable a + , VU.Unbox a + , AssertPresent name cols + ) => + TypedDataFrame cols -> VU.Vector a +columnAsUnboxedVector (TDF df) = + either throw id (D.columnAsUnboxedVector (Col @a colName) df) + where + colName = T.pack (symbolVal (Proxy @name)) + +{- | Convert every column to 'Double' and transpose into a row-major matrix. +Total: 'AllColumnsReal' proves at compile time that every column is numeric and +unboxed, so the conversion cannot fail. +-} +toDoubleMatrix :: + (AllColumnsReal "toDoubleMatrix" cols) => + TypedDataFrame cols -> V.Vector (VU.Vector Double) +toDoubleMatrix (TDF df) = either throw id (D.toDoubleMatrix df) + +-- | Convert every column to 'Float' and transpose into a row-major matrix. See 'toDoubleMatrix'. +toFloatMatrix :: + (AllColumnsReal "toFloatMatrix" cols) => + TypedDataFrame cols -> V.Vector (VU.Vector Float) +toFloatMatrix (TDF df) = either throw id (D.toFloatMatrix df) + +-- | Convert every column to 'Int' and transpose into a row-major matrix. See 'toDoubleMatrix'. +toIntMatrix :: + (AllColumnsReal "toIntMatrix" cols) => + TypedDataFrame cols -> V.Vector (VU.Vector Int) +toIntMatrix (TDF df) = either throw id (D.toIntMatrix df) diff --git a/src/DataFrame/Typed/Aggregate.hs b/src/DataFrame/Typed/Aggregate.hs --- a/src/DataFrame/Typed/Aggregate.hs +++ b/src/DataFrame/Typed/Aggregate.hs @@ -83,7 +83,7 @@ (KnownSymbol name, Columnable a) => TExpr cols a -> TAgg keys cols aggs -> - TAgg keys cols (Column name a ': aggs) + TAgg keys cols ('(name, a) ': aggs) as = TAggCons (T.pack (symbolVal (Proxy @name))) {- | Run a typed aggregation against a grouped DataFrame. @@ -99,9 +99,9 @@ . as \@\"orders\" (count (col \@\"order_id\")) ) -- result :: TypedDataFrame --- '[ Column \"region\" Text --- , Column \"total\" Double --- , Column \"orders\" Int +-- '[ '(\"region\", Text) +-- , '(\"total\", Double) +-- , '(\"orders\", Int) -- ] @ -} diff --git a/src/DataFrame/Typed/Apply.hs b/src/DataFrame/Typed/Apply.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Apply.hs @@ -0,0 +1,218 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} + +{- | Typed column transformations: the @apply@\/@derive@ family and +default-valued inserts, plus the horizontal merge @('|||')@. Schema changes are +tracked at the type level (e.g. 'applyColumn' rewrites a column's element type +via 'SetColumnType'). +-} +module DataFrame.Typed.Apply ( + applyColumn, + applyMany, + applyWhere, + applyAtIndex, + safeApply, + deriveWithExpr, + insertWithDefault, + insertVectorWithDefault, + insertUnboxedVector, + (|||), +) where + +import Data.Proxy (Proxy (..)) +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import GHC.TypeLits (KnownSymbol, Symbol, symbolVal) + +import DataFrame.Errors (DataFrameException) +import DataFrame.Internal.Column (Columnable) +import qualified DataFrame.Operations.Core as D +import qualified DataFrame.Operations.Merge as D +import qualified DataFrame.Operations.Transformations as D +import DataFrame.Typed.Freeze (unsafeFreeze) +import DataFrame.Typed.Schema ( + AllKnownSymbol, + Append, + AssertAbsent, + AssertAllColumnsHaveType, + AssertDisjoint, + AssertPresent, + SafeLookup, + SetColumnType, + Snoc, + symbolVals, + ) +import DataFrame.Typed.Types (TExpr (..), TypedDataFrame (..)) + +{- | Map a function over a column, rewriting its element type from @a@ to @b@. +The schema's entry for @name@ is updated via 'SetColumnType'. + +@ +df' = applyColumn \@\"age\" (show :: Int -> String) df +-- the \"age\" column is now String-typed +@ +-} +applyColumn :: + forall name a b cols. + ( KnownSymbol name + , a ~ SafeLookup name cols + , Columnable a + , Columnable b + , AssertPresent name cols + ) => + (a -> b) -> + TypedDataFrame cols -> + TypedDataFrame (SetColumnType name b cols) +applyColumn f (TDF df) = unsafeFreeze (D.apply f colName df) + where + colName = T.pack (symbolVal (Proxy @name)) + +-- | Like 'applyColumn' but returns the error instead of throwing. +safeApply :: + forall name a b cols. + ( KnownSymbol name + , a ~ SafeLookup name cols + , Columnable a + , Columnable b + , AssertPresent name cols + ) => + (a -> b) -> + TypedDataFrame cols -> + Either DataFrameException (TypedDataFrame (SetColumnType name b cols)) +safeApply f (TDF df) = fmap unsafeFreeze (D.safeApply f colName df) + where + colName = T.pack (symbolVal (Proxy @name)) + +{- | Apply a type-preserving function to several columns at once. Every named +column must already share the element type @a@ (enforced by +'AssertAllColumnsHaveType'). +-} +applyMany :: + forall (names :: [Symbol]) a cols. + (AllKnownSymbol names, Columnable a, AssertAllColumnsHaveType names a cols) => + (a -> a) -> + TypedDataFrame cols -> + TypedDataFrame cols +applyMany f (TDF df) = TDF (D.applyMany f (symbolVals @names) df) + +{- | Apply a function to a target column only on rows where a condition holds on +a filter column. Both columns are named by type application; the target keeps +its type. + +@ +applyWhere \@\"flagged\" \@\"score\" id (* 2) df +@ +-} +applyWhere :: + forall filterName targetName a b cols. + ( KnownSymbol filterName + , KnownSymbol targetName + , a ~ SafeLookup filterName cols + , b ~ SafeLookup targetName cols + , Columnable a + , Columnable b + , AssertPresent filterName cols + , AssertPresent targetName cols + ) => + (a -> Bool) -> + (b -> b) -> + TypedDataFrame cols -> + TypedDataFrame cols +applyWhere cond f (TDF df) = TDF (D.applyWhere cond filterName f targetName df) + where + filterName = T.pack (symbolVal (Proxy @filterName)) + targetName = T.pack (symbolVal (Proxy @targetName)) + +-- | Apply a type-preserving function to a single row of a column. +applyAtIndex :: + forall name a cols. + ( KnownSymbol name + , a ~ SafeLookup name cols + , Columnable a + , AssertPresent name cols + ) => + Int -> + (a -> a) -> + TypedDataFrame cols -> + TypedDataFrame cols +applyAtIndex i f (TDF df) = TDF (D.applyAtIndex i f colName df) + where + colName = T.pack (symbolVal (Proxy @name)) + +{- | Derive a new column and also return a typed reference to it. The returned +expression lives in the extended schema, so it can feed later operations. +-} +deriveWithExpr :: + forall name a cols. + ( KnownSymbol name + , Columnable a + , AssertAbsent name cols + ) => + TExpr cols a -> + TypedDataFrame cols -> + ( TExpr (Snoc cols '(name, a)) a + , TypedDataFrame (Snoc cols '(name, a)) + ) +deriveWithExpr (TExpr expr) (TDF df) = + let (e', df') = D.deriveWithExpr colName expr df + in (TExpr e', unsafeFreeze df') + where + colName = T.pack (symbolVal (Proxy @name)) + +-- | Insert a column from a 'Foldable', padding missing rows with a default. +insertWithDefault :: + forall name a cols t. + ( KnownSymbol name + , Columnable a + , Foldable t + , AssertAbsent name cols + ) => + a -> t a -> TypedDataFrame cols -> TypedDataFrame ('(name, a) ': cols) +insertWithDefault def xs (TDF df) = + unsafeFreeze (D.insertWithDefault def colName xs df) + where + colName = T.pack (symbolVal (Proxy @name)) + +-- | Insert a boxed 'V.Vector', padding missing rows with a default. +insertVectorWithDefault :: + forall name a cols. + ( KnownSymbol name + , Columnable a + , AssertAbsent name cols + ) => + a -> V.Vector a -> TypedDataFrame cols -> TypedDataFrame ('(name, a) ': cols) +insertVectorWithDefault def vec (TDF df) = + unsafeFreeze (D.insertVectorWithDefault def colName vec df) + where + colName = T.pack (symbolVal (Proxy @name)) + +-- | Insert an unboxed 'VU.Vector' as a new column. +insertUnboxedVector :: + forall name a cols. + ( KnownSymbol name + , Columnable a + , VU.Unbox a + , AssertAbsent name cols + ) => + VU.Vector a -> TypedDataFrame cols -> TypedDataFrame ('(name, a) ': cols) +insertUnboxedVector vec (TDF df) = + unsafeFreeze (D.insertUnboxedVector colName vec df) + where + colName = T.pack (symbolVal (Proxy @name)) + +{- | Horizontal merge: place two DataFrames side by side. The schemas must be +disjoint (no shared column names), enforced by 'AssertDisjoint'; the result +schema is their concatenation. +-} +(|||) :: + (AssertDisjoint left right) => + TypedDataFrame left -> + TypedDataFrame right -> + TypedDataFrame (Append left right) +(TDF a) ||| (TDF b) = unsafeFreeze (a D.||| b) diff --git a/src/DataFrame/Typed/Expr.hs b/src/DataFrame/Typed/Expr.hs --- a/src/DataFrame/Typed/Expr.hs +++ b/src/DataFrame/Typed/Expr.hs @@ -25,7 +25,7 @@ == Example @ -type Schema = '[Column \"age\" Int, Column \"salary\" Double] +type Schema = '[ '(\"age\", Int), '(\"salary\", Double)] -- This compiles: goodExpr :: TExpr Schema Double @@ -123,6 +123,9 @@ -- * Sort helpers asc, desc, + + -- * Additional expression functions + module DataFrame.Typed.Expr.Extra, ) where import Data.Either (fromRight) @@ -134,12 +137,13 @@ import qualified DataFrame.Functions as F import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.Column.Types (Promote, PromoteDiv) import DataFrame.Internal.Expression ( BinUDF (..), Expr (..), UnUDF (..), ) -import DataFrame.Internal.Nullable ( +import DataFrame.Internal.Expression.Operators.Nullable ( BaseType, DivWidenOp, NullCmpResult, @@ -155,9 +159,9 @@ widenArithOp, widenCmpOp, ) -import DataFrame.Internal.Types (Promote, PromoteDiv) import qualified Data.Vector.Unboxed as VU +import DataFrame.Typed.Expr.Extra import DataFrame.Typed.Schema ( AllKnownSymbol, AssertAllPresent, @@ -174,7 +178,7 @@ Both checks happen at compile time via type families. @ -salary :: TExpr '[Column \"salary\" Double] Double +salary :: TExpr '[(\"salary\", Double)] Double salary = col \@\"salary\" @ -} @@ -217,10 +221,6 @@ TExpr cols Bool -> TExpr cols a -> TExpr cols a -> TExpr cols a ifThenElse (TExpr c) (TExpr t) (TExpr e) = TExpr (If c t e) -------------------------------------------------------------------------------- --- Numeric instances (mirror Expr's instances) -------------------------------------------------------------------------------- - instance (Num a, Columnable a) => Num (TExpr cols a) where (TExpr a) + (TExpr b) = TExpr (a + b) (TExpr a) - (TExpr b) = TExpr (a - b) @@ -256,10 +256,6 @@ instance (IsString a, Columnable a) => IsString (TExpr cols a) where fromString = TExpr . fromString -------------------------------------------------------------------------------- --- Lifting arbitrary functions -------------------------------------------------------------------------------- - -- | Lift a unary function into a typed expression. lift :: (Columnable a, Columnable b) => (a -> b) -> TExpr cols a -> TExpr cols b @@ -377,10 +373,6 @@ (.||) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (nullCmpOp (||)) "nullor" (Just ".||") True 2) a b) -------------------------------------------------------------------------------- --- Nullable-aware arithmetic operators -------------------------------------------------------------------------------- - infixl 6 .+, .- infixl 7 .*, ./ @@ -481,10 +473,6 @@ TExpr cols a -> TExpr cols b -> TExpr cols a (.^) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (applyNull2 (^)) "pow" (Just ".^") False 8) a b) - -------------------------------------------------------------------------------- --- Nullable-aware comparison operators (three-valued logic) -------------------------------------------------------------------------------- {- | Nullable-aware equality. Widens numeric operands to their common type, so @TExpr cols Double .== TExpr cols Int@ typechecks. Returns @Maybe Bool@ diff --git a/src/DataFrame/Typed/Expr/Extra.hs b/src/DataFrame/Typed/Expr/Extra.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Expr/Extra.hs @@ -0,0 +1,221 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE ScopedTypeVariables #-} + +{- | Typed counterparts of the remaining "DataFrame.Functions" expression +combinators not already provided by "DataFrame.Typed.Expr". Each wraps the +untyped combinator 1:1, replacing @Expr@ with @'TExpr' cols@ so column +references stay schema-checked. Re-exported from "DataFrame.Typed.Expr". +-} +module DataFrame.Typed.Expr.Extra ( + div, + mod, + mode, + sumMaybe, + meanMaybe, + variance, + medianMaybe, + percentile, + stddev, + stddevMaybe, + zScore, + pow, + relu, + min, + max, + reduce, + toMaybe, + fromMaybe, + isJust, + isNothing, + fromJust, + whenPresent, + whenBothPresent, + recode, + recodeWithCondition, + recodeWithDefault, + firstOrNothing, + lastOrNothing, + splitOn, + match, + matchAll, + parseDate, + daysBetween, + bind, +) where + +import qualified Data.Text as T +import Data.Time (Day, ParseTime) +import qualified Data.Vector.Unboxed as VU +import Prelude hiding (div, max, min, mod) + +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column (Columnable) +import DataFrame.Typed.Types (TExpr (..)) + +-- | Integer division. +div :: + (Integral a, Columnable a) => TExpr cols a -> TExpr cols a -> TExpr cols a +div (TExpr a) (TExpr b) = TExpr (F.div a b) + +-- | Integer modulus. +mod :: + (Integral a, Columnable a) => TExpr cols a -> TExpr cols a -> TExpr cols a +mod (TExpr a) (TExpr b) = TExpr (F.mod a b) + +-- | Most frequent value (aggregation). +mode :: (Ord a, Columnable a, Eq a) => TExpr cols a -> TExpr cols a +mode (TExpr e) = TExpr (F.mode e) + +-- | Sum of a nullable column, ignoring 'Nothing' (aggregation). +sumMaybe :: (Columnable a, Num a) => TExpr cols (Maybe a) -> TExpr cols a +sumMaybe (TExpr e) = TExpr (F.sumMaybe e) + +-- | Mean of a nullable column, ignoring 'Nothing' (aggregation). +meanMaybe :: (Columnable a, Real a) => TExpr cols (Maybe a) -> TExpr cols Double +meanMaybe (TExpr e) = TExpr (F.meanMaybe e) + +-- | Variance (aggregation). +variance :: + (Columnable a, Real a, VU.Unbox a) => TExpr cols a -> TExpr cols Double +variance (TExpr e) = TExpr (F.variance e) + +-- | Median of a nullable column, ignoring 'Nothing' (aggregation). +medianMaybe :: + (Columnable a, Real a) => TExpr cols (Maybe a) -> TExpr cols Double +medianMaybe (TExpr e) = TExpr (F.medianMaybe e) + +-- | The @n@-th percentile (aggregation). +percentile :: Int -> TExpr cols Double -> TExpr cols Double +percentile n (TExpr e) = TExpr (F.percentile n e) + +-- | Standard deviation (aggregation). +stddev :: + (Columnable a, Real a, VU.Unbox a) => TExpr cols a -> TExpr cols Double +stddev (TExpr e) = TExpr (F.stddev e) + +-- | Standard deviation of a nullable column, ignoring 'Nothing' (aggregation). +stddevMaybe :: + (Columnable a, Real a) => TExpr cols (Maybe a) -> TExpr cols Double +stddevMaybe (TExpr e) = TExpr (F.stddevMaybe e) + +-- | Z-score (value minus group mean, over standard deviation). +zScore :: TExpr cols Double -> TExpr cols Double +zScore (TExpr e) = TExpr (F.zScore e) + +-- | Raise an expression to an integer power. +pow :: (Columnable a, Num a) => TExpr cols a -> Int -> TExpr cols a +pow (TExpr e) i = TExpr (F.pow e i) + +-- | Rectified linear unit: @max 0@. +relu :: (Columnable a, Num a, Ord a) => TExpr cols a -> TExpr cols a +relu (TExpr e) = TExpr (F.relu e) + +-- | Element-wise minimum of two expressions. +min :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols a +min (TExpr a) (TExpr b) = TExpr (F.min a b) + +-- | Element-wise maximum of two expressions. +max :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols a +max (TExpr a) (TExpr b) = TExpr (F.max a b) + +-- | Fold a column into a single value with a seed and step function (aggregation). +reduce :: + (Columnable a, Columnable b) => + TExpr cols b -> a -> (a -> b -> a) -> TExpr cols a +reduce (TExpr e) start f = TExpr (F.reduce e start f) + +-- | Wrap each value in 'Just'. +toMaybe :: (Columnable a) => TExpr cols a -> TExpr cols (Maybe a) +toMaybe (TExpr e) = TExpr (F.toMaybe e) + +-- | Replace 'Nothing' with a default. +fromMaybe :: (Columnable a) => a -> TExpr cols (Maybe a) -> TExpr cols a +fromMaybe d (TExpr e) = TExpr (F.fromMaybe d e) + +-- | True where the value is 'Just'. +isJust :: (Columnable a) => TExpr cols (Maybe a) -> TExpr cols Bool +isJust (TExpr e) = TExpr (F.isJust e) + +-- | True where the value is 'Nothing'. +isNothing :: (Columnable a) => TExpr cols (Maybe a) -> TExpr cols Bool +isNothing (TExpr e) = TExpr (F.isNothing e) + +-- | Unwrap a 'Just', erroring on 'Nothing'. +fromJust :: (Columnable a) => TExpr cols (Maybe a) -> TExpr cols a +fromJust (TExpr e) = TExpr (F.fromJust e) + +-- | Apply a function only where the value is present. +whenPresent :: + (Columnable a, Columnable b) => + (a -> b) -> TExpr cols (Maybe a) -> TExpr cols (Maybe b) +whenPresent f (TExpr e) = TExpr (F.whenPresent f e) + +-- | Apply a binary function only where both values are present. +whenBothPresent :: + (Columnable a, Columnable b, Columnable c) => + (a -> b -> c) -> + TExpr cols (Maybe a) -> + TExpr cols (Maybe b) -> + TExpr cols (Maybe c) +whenBothPresent f (TExpr a) (TExpr b) = TExpr (F.whenBothPresent f a b) + +-- | Map values through a lookup table, yielding 'Nothing' for misses. +recode :: + (Columnable a, Columnable b, Show a, Show b, Show (a, b)) => + [(a, b)] -> TExpr cols a -> TExpr cols (Maybe b) +recode mapping (TExpr e) = TExpr (F.recode mapping e) + +-- | Pick the first value whose condition holds, else a fallback. +recodeWithCondition :: + (Columnable a, Columnable b) => + TExpr cols b -> + [(TExpr cols a -> TExpr cols Bool, b)] -> + TExpr cols a -> + TExpr cols b +recodeWithCondition (TExpr fallback) conds (TExpr e) = + TExpr (F.recodeWithCondition fallback (map untype conds) e) + where + untype (p, v) = (unTExpr . p . TExpr, v) + +-- | Map values through a lookup table, with a default for misses. +recodeWithDefault :: + (Columnable a, Columnable b, Show (a, b)) => + b -> [(a, b)] -> TExpr cols a -> TExpr cols b +recodeWithDefault d mapping (TExpr e) = TExpr (F.recodeWithDefault d mapping e) + +-- | First element of a list column, or 'Nothing'. +firstOrNothing :: (Columnable a) => TExpr cols [a] -> TExpr cols (Maybe a) +firstOrNothing (TExpr e) = TExpr (F.firstOrNothing e) + +-- | Last element of a list column, or 'Nothing'. +lastOrNothing :: (Columnable a) => TExpr cols [a] -> TExpr cols (Maybe a) +lastOrNothing (TExpr e) = TExpr (F.lastOrNothing e) + +-- | Split text on a delimiter. +splitOn :: T.Text -> TExpr cols T.Text -> TExpr cols [T.Text] +splitOn delim (TExpr e) = TExpr (F.splitOn delim e) + +-- | First regex match, or 'Nothing'. +match :: T.Text -> TExpr cols T.Text -> TExpr cols (Maybe T.Text) +match regex (TExpr e) = TExpr (F.match regex e) + +-- | All regex matches. +matchAll :: T.Text -> TExpr cols T.Text -> TExpr cols [T.Text] +matchAll regex (TExpr e) = TExpr (F.matchAll regex e) + +-- | Parse text into a time value with the given format. +parseDate :: + (ParseTime t, Columnable t) => + T.Text -> TExpr cols T.Text -> TExpr cols (Maybe t) +parseDate format (TExpr e) = TExpr (F.parseDate format e) + +-- | Number of days between two dates. +daysBetween :: TExpr cols Day -> TExpr cols Day -> TExpr cols Int +daysBetween (TExpr a) (TExpr b) = TExpr (F.daysBetween a b) + +-- | Monadic bind over a column of monadic values. +bind :: + (Columnable a, Columnable (m a), Monad m, Columnable b, Columnable (m b)) => + (a -> m b) -> TExpr cols (m a) -> TExpr cols (m b) +bind f (TExpr e) = TExpr (F.bind f e) diff --git a/src/DataFrame/Typed/Operations.hs b/src/DataFrame/Typed/Operations.hs --- a/src/DataFrame/Typed/Operations.hs +++ b/src/DataFrame/Typed/Operations.hs @@ -17,6 +17,7 @@ filterAllJust, filterJust, filterNothing, + filterAllNothing, sortBy, take, takeLast, @@ -42,6 +43,10 @@ dropColumn, replaceColumn, + -- * Frequencies + valueCounts, + valueProportions, + -- * Metadata dimensions, nRows, @@ -80,12 +85,7 @@ import DataFrame.Typed.Freeze (unsafeFreeze) import DataFrame.Typed.Schema import DataFrame.Typed.Types (TExpr (..), TSortOrder (..), TypedDataFrame (..)) -import qualified DataFrame.Typed.Types as T -------------------------------------------------------------------------------- --- Schema-preserving operations -------------------------------------------------------------------------------- - {- | Filter rows where a boolean expression evaluates to True. The expression is validated against the schema at compile time. -} @@ -108,8 +108,8 @@ Strips 'Maybe' from all column types in the result schema. @ -df :: TDF '[Column \"x\" (Maybe Double), Column \"y\" Int] -filterAllJust df :: TDF '[Column \"x\" Double, Column \"y\" Int] +df :: TDF '[ '(\"x\", Maybe Double), '(\"y\", Int)] +filterAllJust df :: TDF '[ '(\"x\", Double), '(\"y\", Int)] @ -} filterAllJust :: TypedDataFrame cols -> TypedDataFrame (StripAllMaybe cols) @@ -145,6 +145,12 @@ where colName = T.pack (symbolVal (Proxy @name)) +{- | Keep only rows where every nullable column has Nothing. +Schema is preserved. +-} +filterAllNothing :: TypedDataFrame cols -> TypedDataFrame cols +filterAllNothing (TDF df) = TDF (D.filterAllNothing df) + {- | Sort by the given typed sort orders. Sort orders reference columns that are validated against the schema. -} @@ -192,17 +198,13 @@ shuffle :: (RandomGen g) => g -> TypedDataFrame cols -> TypedDataFrame cols shuffle g (TDF df) = TDF (D.shuffle g df) -------------------------------------------------------------------------------- --- Schema-modifying operations -------------------------------------------------------------------------------- - {- | Derive a new column from a typed expression. The column name must NOT already exist in the schema (enforced at compile time via 'AssertAbsent'). The expression is validated against the current schema. @ df' = derive \@\"total\" (col \@\"price\" * col \@\"qty\") df --- df' :: TDF (Column \"total\" Double ': originalCols) +-- df' :: TDF ('(\"total\", Double ': originalCols)) @ -} derive :: @@ -213,7 +215,7 @@ ) => TExpr cols a -> TypedDataFrame cols -> - TypedDataFrame (Snoc cols (T.Column name a)) + TypedDataFrame (Snoc cols '(name, a)) derive (TExpr expr) (TDF df) = unsafeFreeze (D.derive colName expr df) where colName = T.pack (symbolVal (Proxy @name)) @@ -275,7 +277,7 @@ , Foldable t , AssertAbsent name cols ) => - t a -> TypedDataFrame cols -> TypedDataFrame (T.Column name a ': cols) + t a -> TypedDataFrame cols -> TypedDataFrame ('(name, a) ': cols) insert xs (TDF df) = unsafeFreeze (D.insert colName xs df) where colName = T.pack (symbolVal (Proxy @name)) @@ -287,7 +289,7 @@ , Columnable a , AssertAbsent name cols ) => - C.Column -> TypedDataFrame cols -> TypedDataFrame (T.Column name a ': cols) + C.Column -> TypedDataFrame cols -> TypedDataFrame ('(name, a) ': cols) insertColumn col (TDF df) = unsafeFreeze (D.insertColumn colName col df) where colName = T.pack (symbolVal (Proxy @name)) @@ -299,7 +301,7 @@ , Columnable a , AssertAbsent name cols ) => - V.Vector a -> TypedDataFrame cols -> TypedDataFrame (T.Column name a ': cols) + V.Vector a -> TypedDataFrame cols -> TypedDataFrame ('(name, a) ': cols) insertVector vec (TDF df) = unsafeFreeze (D.insertVector colName vec df) where colName = T.pack (symbolVal (Proxy @name)) @@ -312,7 +314,7 @@ , AssertPresent old cols , AssertAbsent new cols ) => - TypedDataFrame cols -> TypedDataFrame (T.Column new (Lookup old cols) ': cols) + TypedDataFrame cols -> TypedDataFrame ('(new, Lookup old cols) ': cols) cloneColumn (TDF df) = unsafeFreeze (D.cloneColumn oldName newName df) where oldName = T.pack (symbolVal (Proxy @old)) @@ -348,13 +350,6 @@ append :: TypedDataFrame cols -> TypedDataFrame cols -> TypedDataFrame cols append (TDF a) (TDF b) = TDF (a <> b) -------------------------------------------------------------------------------- --- Set algebra (topos operations) --- --- Each treats a DataFrame as a /set/ of rows and is schema-preserving: --- the output type equals the input type; only which rows are present changes. -------------------------------------------------------------------------------- - -- | Rows appearing in either DataFrame, deduplicated (set union). union :: TypedDataFrame cols -> TypedDataFrame cols -> TypedDataFrame cols union (TDF a) (TDF b) = TDF (DS.union a b) @@ -374,10 +369,6 @@ TypedDataFrame cols -> TypedDataFrame cols -> TypedDataFrame cols symmetricDifference (TDF a) (TDF b) = TDF (DS.symmetricDifference a b) -------------------------------------------------------------------------------- --- Metadata (pass-through) -------------------------------------------------------------------------------- - dimensions :: TypedDataFrame cols -> (Int, Int) dimensions (TDF df) = D.dimensions df @@ -390,9 +381,15 @@ columnNames :: TypedDataFrame cols -> [T.Text] columnNames (TDF df) = D.columnNames df -------------------------------------------------------------------------------- --- Internal helpers -------------------------------------------------------------------------------- +-- | Count occurrences of each distinct value in a column. +valueCounts :: + (Ord a, Columnable a) => TExpr cols a -> TypedDataFrame cols -> [(a, Int)] +valueCounts (TExpr e) (TDF df) = D.valueCounts e df + +-- | Proportion of each distinct value in a column. +valueProportions :: + (Ord a, Columnable a) => TExpr cols a -> TypedDataFrame cols -> [(a, Double)] +valueProportions (TExpr e) (TDF df) = D.valueProportions e df -- | Helper class for extracting [(Text, Text)] from type-level pairs. class AllKnownPairs (pairs :: [(Symbol, Symbol)]) where diff --git a/src/DataFrame/Typed/Sampling.hs b/src/DataFrame/Typed/Sampling.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Sampling.hs @@ -0,0 +1,55 @@ +{-# LANGUAGE FlexibleContexts #-} + +{- | Typed sampling and splitting. All operations are schema-preserving: they +change which rows are present, never the columns, so every result reuses the +input schema @cols@. +-} +module DataFrame.Typed.Sampling ( + randomSplit, + kFolds, + selectRows, + stratifiedSample, + stratifiedSplit, +) where + +import System.Random (RandomGen) + +import DataFrame.Internal.Column (Columnable) +import DataFrame.Operations.Subset (SplittableGen) +import qualified DataFrame.Operations.Subset as D +import DataFrame.Typed.Types (TExpr (..), TypedDataFrame (..)) + +-- | Split rows into two DataFrames by a fraction. +randomSplit :: + (RandomGen g) => + g -> Double -> TypedDataFrame cols -> (TypedDataFrame cols, TypedDataFrame cols) +randomSplit g p (TDF df) = let (a, b) = D.randomSplit g p df in (TDF a, TDF b) + +-- | Partition rows into @k@ folds. +kFolds :: + (RandomGen g) => g -> Int -> TypedDataFrame cols -> [TypedDataFrame cols] +kFolds g k (TDF df) = map TDF (D.kFolds g k df) + +{- | Select rows by index. +| This may fail if the indices are out of bounds; +| use with caution or use 'filter' to select rows by a predicate instead. +-} +selectRows :: [Int] -> TypedDataFrame cols -> TypedDataFrame cols +selectRows ixs (TDF df) = TDF (D.selectRows ixs df) + +-- | Sample a fraction of rows, preserving the distribution of a strata column. +stratifiedSample :: + (SplittableGen g, Columnable a) => + g -> Double -> TExpr cols a -> TypedDataFrame cols -> TypedDataFrame cols +stratifiedSample g p (TExpr e) (TDF df) = TDF (D.stratifiedSample g p e df) + +-- | Split rows by a fraction, preserving the distribution of a strata column. +stratifiedSplit :: + (SplittableGen g, Columnable a) => + g -> + Double -> + TExpr cols a -> + TypedDataFrame cols -> + (TypedDataFrame cols, TypedDataFrame cols) +stratifiedSplit g p (TExpr e) (TDF df) = + let (a, b) = D.stratifiedSplit g p e df in (TDF a, TDF b) diff --git a/src/DataFrame/Typed/Statistics.hs b/src/DataFrame/Typed/Statistics.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Statistics.hs @@ -0,0 +1,180 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} + +{- | Typed statistical reducers over a 'TypedDataFrame'. + +These mirror the untyped reducers in "DataFrame.Operations.Statistics", taking +a schema-checked 'TExpr' instead of a raw @Expr@. The names (@sum@, @mean@, +@median@, …) deliberately collide with the aggregation-expression combinators +in "DataFrame.Typed.Expr", so this module is meant to be imported qualified: + +@ +import qualified DataFrame.Typed.Statistics as TS + +avg = TS.mean (col \@\"salary\") employees +@ +-} +module DataFrame.Typed.Statistics ( + mean, + meanMaybe, + median, + medianMaybe, + percentile, + genericPercentile, + standardDeviation, + skewness, + variance, + interQuartileRange, + sum, + correlation, + frequencies, + imputeWith, + summarize, + describeColumns, +) where + +import Data.Proxy (Proxy (..)) +import qualified Data.Text as T +import qualified Data.Vector.Unboxed as VU +import GHC.TypeLits (KnownSymbol, symbolVal) +import Prelude hiding (sum) + +import DataFrame.Internal.Column (Columnable) +import qualified DataFrame.Internal.DataFrame as D +import DataFrame.Internal.Expression.Operators.Nullable (BaseType) +import qualified DataFrame.Operations.Core as Core +import qualified DataFrame.Operations.Statistics as Stats +import DataFrame.Operations.Transformations (ImputeOp) +import DataFrame.Typed.Schema (AssertPresent, SafeLookup) +import DataFrame.Typed.Types (TExpr (..), TypedDataFrame (..)) + +-- | Mean of a column. +mean :: + (Columnable a, Real a, VU.Unbox a) => + TExpr cols a -> TypedDataFrame cols -> Double +mean (TExpr e) (TDF df) = Stats.mean e df + +-- | Mean of a nullable column, ignoring 'Nothing'. +meanMaybe :: + (Columnable a, Real a) => + TExpr cols (Maybe a) -> TypedDataFrame cols -> Double +meanMaybe (TExpr e) (TDF df) = Stats.meanMaybe e df + +-- | Median of a column. +median :: + (Columnable a, Real a, VU.Unbox a) => + TExpr cols a -> TypedDataFrame cols -> Double +median (TExpr e) (TDF df) = Stats.median e df + +-- | Median of a nullable column, ignoring 'Nothing'. +medianMaybe :: + (Columnable a, Real a) => + TExpr cols (Maybe a) -> TypedDataFrame cols -> Double +medianMaybe (TExpr e) (TDF df) = Stats.medianMaybe e df + +-- | The @n@-th percentile of a column. +percentile :: + (Columnable a, Real a, VU.Unbox a) => + Int -> TExpr cols a -> TypedDataFrame cols -> Double +percentile n (TExpr e) (TDF df) = Stats.percentile n e df + +-- | The @n@-th percentile of a column of any 'Ord' type. +genericPercentile :: + (Columnable a, Ord a) => + Int -> TExpr cols a -> TypedDataFrame cols -> a +genericPercentile n (TExpr e) (TDF df) = Stats.genericPercentile n e df + +-- | Standard deviation of a column. +standardDeviation :: + (Columnable a, Real a, VU.Unbox a) => + TExpr cols a -> TypedDataFrame cols -> Double +standardDeviation (TExpr e) (TDF df) = Stats.standardDeviation e df + +-- | Skewness of a column. +skewness :: + (Columnable a, Real a, VU.Unbox a) => + TExpr cols a -> TypedDataFrame cols -> Double +skewness (TExpr e) (TDF df) = Stats.skewness e df + +-- | Variance of a column. +variance :: + (Columnable a, Real a, VU.Unbox a) => + TExpr cols a -> TypedDataFrame cols -> Double +variance (TExpr e) (TDF df) = Stats.variance e df + +-- | Inter-quartile range of a column. +interQuartileRange :: + (Columnable a, Real a, VU.Unbox a) => + TExpr cols a -> TypedDataFrame cols -> Double +interQuartileRange (TExpr e) (TDF df) = Stats.interQuartileRange e df + +-- | Sum of a column. +sum :: (Columnable a, Num a) => TExpr cols a -> TypedDataFrame cols -> a +sum (TExpr e) (TDF df) = Stats.sum e df + +{- | Pearson's correlation coefficient between two columns, named by type +application. Both columns must exist in the schema and be numeric — these are +checked at compile time via 'SafeLookup' on each name. + +@ +TS.correlation \@\"height\" \@\"weight\" people +@ +-} +correlation :: + forall c1 c2 a b cols. + ( KnownSymbol c1 + , KnownSymbol c2 + , a ~ SafeLookup c1 cols + , b ~ SafeLookup c2 cols + , Columnable a + , Columnable b + , Real a + , Real b + , VU.Unbox a + , VU.Unbox b + , AssertPresent c1 cols + , AssertPresent c2 cols + ) => + TypedDataFrame cols -> Maybe Double +correlation (TDF df) = + Stats.correlation + (T.pack (symbolVal (Proxy @c1))) + (T.pack (symbolVal (Proxy @c2))) + df + +{- | Frequency table for a column. The result schema is data-dependent +(one column per distinct value), so an untyped 'D.DataFrame' is returned. +-} +frequencies :: + (Columnable a, Ord a) => TExpr cols a -> TypedDataFrame cols -> D.DataFrame +frequencies (TExpr e) (TDF df) = Stats.frequencies e df + +{- | Impute missing values in a column using a derived scalar (e.g. the mean). +Schema-preserving: the imputed column keeps its type-level @Maybe@ even though +its runtime values are now fully populated. +-} +imputeWith :: + (ImputeOp a, Columnable (BaseType a)) => + (TExpr cols (BaseType a) -> TExpr cols (BaseType a)) -> + TExpr cols a -> + TypedDataFrame cols -> + TypedDataFrame cols +imputeWith f (TExpr e) (TDF df) = + TDF (Stats.imputeWith (unTExpr . f . TExpr) e df) + +{- | Descriptive statistics of the numeric columns. Returns an untyped +'D.DataFrame' (the result is a fixed set of statistic rows, not the input schema). +-} +summarize :: TypedDataFrame cols -> D.DataFrame +summarize (TDF df) = Stats.summarize df + +{- | Per-column summary (non-null\/null counts, unique values, type). Returns an +untyped 'D.DataFrame'. +-} +describeColumns :: TypedDataFrame cols -> D.DataFrame +describeColumns (TDF df) = Core.describeColumns df