packages feed

dataframe-operations 1.1.1.1 → 2.5.0.0

raw patch · 28 files changed

Files

dataframe-operations.cabal view
@@ -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
+ src-internal/DataFrame/Internal/Statistics.hs view
@@ -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+    #-}
src/DataFrame/Functions.hs view
@@ -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
− src/DataFrame/Internal/Statistics.hs
@@ -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-    #-}
src/DataFrame/Monad.hs view
@@ -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
− src/DataFrame/Operations/AggregateScatter.hs
@@ -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
src/DataFrame/Operations/Aggregation.hs view
@@ -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
+ src/DataFrame/Operations/Aggregation/Run.hs view
@@ -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
src/DataFrame/Operations/Core.hs view
@@ -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 <Type>` 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 <Type>` 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
src/DataFrame/Operations/Inference.hs view
@@ -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) ->
src/DataFrame/Operations/Join.hs view
@@ -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             )
+ src/DataFrame/Operations/Join/Parallel.hs view
@@ -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])
− src/DataFrame/Operations/JoinPar.hs
@@ -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
src/DataFrame/Operations/Merge.hs view
@@ -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
src/DataFrame/Operations/Permutation.hs view
@@ -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
src/DataFrame/Operations/SetOps.hs view
@@ -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
src/DataFrame/Operations/Statistics.hs view
@@ -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.
src/DataFrame/Operations/Subset.hs view
@@ -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) =
src/DataFrame/Operations/Transformations.hs view
@@ -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.
src/DataFrame/Operations/Typing.hs view
@@ -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
src/DataFrame/Typed/Access.hs view
@@ -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)
src/DataFrame/Typed/Aggregate.hs view
@@ -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) --      ] @ -}
+ src/DataFrame/Typed/Apply.hs view
@@ -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)
src/DataFrame/Typed/Expr.hs view
@@ -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@
+ src/DataFrame/Typed/Expr/Extra.hs view
@@ -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)
src/DataFrame/Typed/Operations.hs view
@@ -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
+ src/DataFrame/Typed/Sampling.hs view
@@ -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)
+ src/DataFrame/Typed/Statistics.hs view
@@ -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