diff --git a/dataframe-operations.cabal b/dataframe-operations.cabal --- a/dataframe-operations.cabal +++ b/dataframe-operations.cabal @@ -1,6 +1,6 @@ -cabal-version: 2.4 +cabal-version: 3.4 name: dataframe-operations -version: 1.1.1.1 +version: 2.6.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,16 @@ library import: warnings exposed-modules: + DataFrame.Internal.Statistics + DataFrame.Expression.Parse 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 +49,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.6 && < 2.7, + dataframe-parsing >= 2.2 && < 2.3, random >= 1.2 && < 2, regex-tdfa >= 1.3.0 && < 2, text >= 2.1 && < 3, time >= 1.12 && < 2, - vector ^>= 0.13, - vector-algorithms ^>= 0.9 - hs-source-dirs: src + vector >= 0.13 && < 0.15, + vector-algorithms >= 0.9 && < 0.11 + hs-source-dirs: src, src-internal default-language: Haskell2010 diff --git a/src-internal/DataFrame/Internal/Statistics.hs b/src-internal/DataFrame/Internal/Statistics.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/Internal/Statistics.hs @@ -0,0 +1,287 @@ +{-# 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) * 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 + -- Not exact for Ord instances; TODO 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 + | VU.length target /= VU.length prediction = Nothing + | VU.null target = Nothing + | otherwise = + Just + ( VU.sum (VU.zipWith (\t p -> (p - t) ^ (2 :: Int)) target prediction) + / fromIntegral (VU.length target) + ) +{-# INLINE meanSquaredError #-} + +mutualInformationBinned :: + Int -> VU.Vector Double -> VU.Vector Double -> Maybe Double +mutualInformationBinned k xs ys + | VU.length xs /= VU.length ys = Nothing + | VU.null xs = Nothing + | k < 2 = Nothing + | rx <= 0 || ry <= 0 = Just 0 + | otherwise = + let bx = VU.map (binIndex xmin xmax k) xs + by = VU.map (binIndex ymin ymax k) ys + n = fromIntegral (VU.length xs) :: Double + mx = bincount k bx + my = bincount k by + mxy = jointBincount k bx by + in Just $ + sum + [ let !cxy = fromIntegral c + !pxy = cxy / n + !px = fromIntegral (mx VU.! i) / n + !py = fromIntegral (my VU.! j) / n + in if c == 0 then 0 else pxy * logBase 2 (pxy / (px * py)) + | i <- [0 .. k - 1] + , j <- [0 .. k - 1] + , let !c = mxy VU.! (i * k + j) + ] + where + (xmin, xmax) = (VU.minimum xs, VU.maximum xs) + (ymin, ymax) = (VU.minimum ys, VU.maximum ys) + rx = xmax - xmin + ry = ymax - ymin + +binIndex :: Double -> Double -> Int -> Double -> Int +binIndex lo hi k x + | hi == lo = 0 + | otherwise = + let !t = (x - lo) / (hi - lo) + !ix = floor (fromIntegral k * t) :: Int + in max 0 (min (k - 1) ix) +{-# INLINE binIndex #-} + +bincount :: Int -> VU.Vector Int -> VU.Vector Int +bincount k bs = VU.create $ do + mv <- VU.thaw (VU.replicate k 0) + VU.forM_ bs $ \b -> do + let i + | b < 0 = 0 + | b >= k = k - 1 + | otherwise = b + x <- VUM.read mv i + VUM.write mv i (x + 1) + pure mv +{-# INLINE bincount #-} + +jointBincount :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int +jointBincount k bx by = VU.create $ do + mv <- VU.thaw (VU.replicate (k * k) 0) + VU.forM_ (VU.zip bx by) $ \(i, j) -> do + let ii = clamp i 0 (k - 1) + jj = clamp j 0 (k - 1) + ix = ii * k + jj + x <- VUM.read mv ix + VUM.write mv ix (x + 1) + pure mv + where + clamp z a b = max a (min b z) +{-# INLINE jointBincount #-} + +rtf :: (Real a) => a -> Double +rtf = realToFrac +{-# NOINLINE [1] rtf #-} + +{-# RULES +"rtf/Double" [2] forall (x :: Double). rtf x = x + #-} diff --git a/src/DataFrame/Expression/Parse.hs b/src/DataFrame/Expression/Parse.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Expression/Parse.hs @@ -0,0 +1,477 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Reload expressions from their 'Show' and 'prettyPrint' printed forms. +module DataFrame.Expression.Parse ( + ColType (..), + Schema, + schemaOf, + parseShown, + parsePretty, + parseShownAs, + parsePrettyAs, +) where + +import Data.Char (isAlphaNum, isDigit, isSpace) +import Data.List (isPrefixOf, sortBy) +import qualified Data.Map.Strict as M +import Data.Ord (Down (..), comparing) +import qualified Data.Text as T + +import Data.Type.Equality (testEquality, (:~:) (Refl)) +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.DataFrame (DataFrame, columnNames) +import Type.Reflection (typeRep) + +import qualified DataFrame.Functions as F +import DataFrame.Internal.Expression (Expr (..), SomeExpr (..), fromSomeExpr) +import DataFrame.Internal.Expression.Operators ( + (.&&.), + (./=.), + (.<.), + (.<=.), + (.==.), + (.>.), + (.>=.), + (.||.), + ) +import DataFrame.Internal.Expression.Reflect ( + Dict (..), + fracDict, + integralDict, + numDict, + ordDict, + realDict, + ) +import DataFrame.Operations.Core (columnAsVector) + +-- | Column types the parser can build expressions over. +data ColType = TInt | TDouble | TText | TBool + deriving (Eq, Show) + +type Schema = M.Map T.Text ColType + +-- | Column types of a frame, for the columns whose type the parser supports. +schemaOf :: DataFrame -> Schema +schemaOf df = M.fromList [(c, t) | c <- columnNames df, Just t <- [typeOf c]] + where + typeOf c + | ok (F.col @Int c) = Just TInt + | ok (F.col @Double c) = Just TDouble + | ok (F.col @T.Text c) = Just TText + | ok (F.col @Bool c) = Just TBool + | otherwise = Nothing + ok :: forall a. (Columnable a) => Expr a -> Bool + ok e = either (const False) (const True) (columnAsVector e df) + +-- | The untyped syntax both printers share. +data Ast + = ACol T.Text (Maybe ColType) + | -- | The value, and whether it was a Double (point or exponent). + ANum Double Bool + | AStr T.Text + | ABool Bool + | AIf Ast Ast Ast + | ACall T.Text [Ast] + deriving (Show) + +-- | Two expressions of one type (the result of making operands agree). +data Same where + Same :: (Columnable t) => Expr t -> Expr t -> Same + +typeName :: SomeExpr -> String +typeName (SomeExpr (_ :: Expr t)) = show (typeRep @t) + +-- | Use evidence for a class on an expression's type, or name the operator. +need :: String -> Maybe (Dict c) -> ((c) => Either String r) -> Either String r +need what d k = maybe (Left (what ++ ": operand type not supported")) (\Dict -> k) d + +elab :: Schema -> Ast -> Either String SomeExpr +elab schema ast = case ast of + ACol c (Just t) -> pure (colOf c t) + ACol c Nothing -> + maybe (Left ("unknown column " ++ show c)) (pure . colOf c) (M.lookup c schema) + ANum v True -> pure (SomeExpr (F.lit v)) + ANum v False -> pure (SomeExpr (F.lit (round v :: Int))) + AStr s -> pure (SomeExpr (F.lit s)) + ABool b -> pure (SomeExpr (F.lit b)) + AIf c t e -> do + c' <- elab schema c + cb <- + maybe + (Left ("if: condition has type " ++ typeName c')) + pure + (fromSomeExpr @Bool c') + t' <- elab schema t + e' <- elab schema e + Same a b <- unify t' e' + pure (SomeExpr (F.ifThenElse cb a b)) + ACall name args -> do + args' <- mapM (elab schema) args + call name args' + where + colOf c t = case t of + TInt -> SomeExpr (F.col @Int c) + TDouble -> SomeExpr (F.col @Double c) + TText -> SomeExpr (F.col @T.Text c) + TBool -> SomeExpr (F.col @Bool c) + +-- | Make two operands the same type; a numeric literal adopts the other's. +unify :: SomeExpr -> SomeExpr -> Either String Same +unify a@(SomeExpr (x :: Expr p)) b@(SomeExpr (y :: Expr q)) + | Just Refl <- testEquality (typeRep @p) (typeRep @q) = pure (Same x y) + | Just (Lit i) <- fromSomeExpr @Int a + , Just yd <- fromSomeExpr @Double b = + pure (Same (F.lit (fromIntegral i)) yd) + | Just xd <- fromSomeExpr @Double a + , Just (Lit i) <- fromSomeExpr @Int b = + pure (Same xd (F.lit (fromIntegral i))) + | Just (Lit d) <- fromSomeExpr @Double a + , Just yi <- fromSomeExpr @Int b + , whole d = + pure (Same (F.lit (round d)) yi) + | Just xi <- fromSomeExpr @Int a + , Just (Lit d) <- fromSomeExpr @Double b + , whole d = + pure (Same xi (F.lit (round d))) + | otherwise = + Left ("operand types differ: " ++ typeName a ++ " vs " ++ typeName b) + where + whole d = fromIntegral (round d :: Int) == (d :: Double) + +call :: T.Text -> [SomeExpr] -> Either String SomeExpr +call name args = case (name, args) of + ("toDouble", [SomeExpr (e :: Expr t)]) -> need "toDouble" (realDict @t) (pure (SomeExpr (F.toDouble e))) + ("not", [a]) -> + maybe + (Left "not needs a Bool operand") + (pure . SomeExpr . F.not) + (fromSomeExpr @Bool a) + ("negate", [SomeExpr (e :: Expr t)]) -> need "negate" (numDict @t) (pure (SomeExpr (negate e))) + ("abs", [SomeExpr (e :: Expr t)]) -> need "abs" (numDict @t) (pure (SomeExpr (abs e))) + (_, [a, b]) -> do + Same x y <- unify a b + binary name x y + _ -> + Left + ( "unsupported call " + ++ T.unpack name + ++ " with " + ++ show (length args) + ++ " arguments" + ) + +binary :: + forall t. (Columnable t) => T.Text -> Expr t -> Expr t -> Either String SomeExpr +binary name x y = case name of + "add" -> need "add" (numDict @t) (pure (SomeExpr (x + y))) + "sub" -> need "sub" (numDict @t) (pure (SomeExpr (x - y))) + "mult" -> need "mult" (numDict @t) (pure (SomeExpr (x * y))) + "divide" -> need "divide" (fracDict @t) (pure (SomeExpr (x / y))) + "div" -> need "div" (integralDict @t) (pure (SomeExpr (F.div x y))) + "mod" -> need "mod" (integralDict @t) (pure (SomeExpr (F.mod x y))) + "min" -> need "min" (ordDict @t) (pure (SomeExpr (F.min x y))) + "max" -> need "max" (ordDict @t) (pure (SomeExpr (F.max x y))) + "eq" -> pure (SomeExpr (x .==. y)) + "neq" -> pure (SomeExpr (x ./=. y)) + "lt" -> need "lt" (ordDict @t) (pure (SomeExpr (x .<. y))) + "gt" -> need "gt" (ordDict @t) (pure (SomeExpr (x .>. y))) + "leq" -> need "leq" (ordDict @t) (pure (SomeExpr (x .<=. y))) + "geq" -> need "geq" (ordDict @t) (pure (SomeExpr (x .>=. y))) + "and" -> bool (.&&.) + "or" -> bool (.||.) + _ -> Left ("unsupported operator " ++ T.unpack name) + where + bool f = case testEquality (typeRep @t) (typeRep @Bool) of + Just Refl -> pure (SomeExpr (f x y)) + Nothing -> Left (T.unpack name ++ " needs Bool operands") + +data SExp = SAtom String | SList [SExp] + deriving (Show) + +sexp :: String -> Either String (SExp, String) +sexp s = case dropWhile isSpace s of + '(' : rest -> list rest [] + '"' : _ -> case reads (dropWhile isSpace s) :: [(String, String)] of + [(str, rest)] -> pure (SAtom (show str), rest) + _ -> Left "bad string literal" + rest -> + let (tok, more) = span (\c -> not (isSpace c) && c /= '(' && c /= ')') rest + in if null tok + then Left ("unexpected input near " ++ take 20 rest) + else pure (SAtom tok, more) + where + list str acc = case dropWhile isSpace str of + ')' : rest -> pure (SList (reverse acc), rest) + "" -> Left "unbalanced parentheses" + more -> do + (e, rest) <- sexp more + list rest (e : acc) + +fromSExp :: SExp -> Either String Ast +fromSExp e = case e of + SList [SAtom "col", SAtom ('@' : ty), SAtom name] -> ACol <$> strLit name <*> (Just <$> colType ty) + SList [SAtom "col", SAtom "@Maybe", SAtom _, SAtom _] -> Left "nullable columns are not supported" + SList [SAtom "lit", inner] -> literal inner + SList [SAtom "ifThenElse", c, t, f] -> AIf <$> fromSExp c <*> fromSExp t <*> fromSExp f + SList (SAtom fn : args) -> ACall (T.pack fn) <$> mapM fromSExp args + SAtom a -> Left ("bare atom " ++ a) + SList [] -> Left "empty list" + SList (SList _ : _) -> Left "list head must be an atom" + where + strLit s = case reads s :: [(String, String)] of + [(v, "")] -> pure (T.pack v) + _ -> Left ("bad name " ++ s) + colType ty = case ty of + "Int" -> pure TInt + "Double" -> pure TDouble + "Text" -> pure TText + "Bool" -> pure TBool + _ -> Left ("unsupported column type " ++ ty) + -- (lit (3.5)), (lit ("x")); the inner parentheses come from show. + literal inner = case inner of + SList [SAtom v] -> atomLit v + SAtom v -> atomLit v + SList [SAtom "-", SAtom v] -> atomLit ('-' : v) + _ -> Left ("bad literal " ++ show inner) + atomLit v + | v == "True" = pure (ABool True) + | v == "False" = pure (ABool False) + | take 1 v == "\"" = AStr <$> strLit v + | otherwise = case reads v :: [(Double, String)] of + [(d, "")] -> pure (ANum d (any (`elem` (".eE" :: String)) v)) + _ -> Left ("bad literal " ++ v) + +-- | Parse the output of 'show' on an 'Expr'; columns carry their own types. +parseShown :: String -> Either String SomeExpr +parseShown s = do + (e, rest) <- sexp s + if all isSpace rest then pure () else Left ("trailing input: " ++ take 30 rest) + ast <- fromSExp e + elab M.empty ast + +data Tok + = TNum Double Bool + | TStr T.Text + | TName T.Text + | TFun T.Text + | TOp String + | TLParen + | TRParen + | TComma + | TIf + | TElse + deriving (Eq, Show) + +operators :: [String] +-- Dotted forms come from Operators, bare forms from Functions; same names. +operators = + sortBy + (comparing (Down . length)) + [ ".<=." + , ".>=." + , ".==." + , "./=." + , ".&&." + , ".||." + , ".<." + , ".>." + , "<=" + , ">=" + , "==" + , "/=" + , "<" + , ">" + , "//" + , "+" + , "-" + , "*" + , "/" + , "~" + ] + +functions :: [T.Text] +functions = ["toDouble", "min", "max", "mod", "abs", "negate", "not", "isJust"] + +lexPretty :: Schema -> String -> Either String [Tok] +lexPretty schema = go + where + names = sortBy (comparing (Down . T.length)) (M.keys schema) + go [] = pure [] + go s@(c : rest) + | isSpace c = go rest + | c == '(' = (TLParen :) <$> go rest + | c == ')' = (TRParen :) <$> go rest + | c == ',' = (TComma :) <$> go rest + | c == '"' = case reads s :: [(String, String)] of + [(str, more)] -> (TStr (T.pack str) :) <$> go more + _ -> Left "bad string literal" + | Just (n, more) <- keyword s = (n :) <$> go more + | Just (n, more) <- longest names s = (TName n :) <$> go more + | Just (f, more) <- longest functions s + , take 1 (dropWhile isSpace more) == "(" = + (TFun f :) <$> go more + | isDigit c = number s + -- A '-' touching a digit is a negative literal; binary '-' is spaced. + | c == '-', (d : _) <- rest, isDigit d = number s + | Just op <- firstPrefix operators s = (TOp op :) <$> go (drop (length op) s) + | otherwise = Left ("cannot lex near: " ++ take 30 s) + keyword s + | "else" `isPrefixOf` s && boundary (drop 4 s) = Just (TElse, drop 4 s) + | "if" `isPrefixOf` s && boundary (drop 2 s) = Just (TIf, drop 2 s) + | otherwise = Nothing + boundary more = case more of + [] -> True + (c : _) -> not (isAlphaNum c) + longest cands s = case [n | n <- cands, T.unpack n `isPrefixOf` s] of + (n : _) -> Just (n, drop (T.length n) s) + [] -> Nothing + firstPrefix cands s = case [op | op <- cands, op `isPrefixOf` s] of + (op : _) -> Just op + [] -> Nothing + number s0 = + let (sign, s) = case s0 of + '-' : more0 -> ("-", more0) + _ -> ("", s0) + (body0, more) = + span (\ch -> isDigit ch || ch == '.' || ch == 'e' || ch == 'E' || ch == '-') s + body = sign ++ body0 + -- A trailing '-' belongs to the next token, not an exponent. + (num, more') = + if not (null body) && last body == '-' + then (init body, '-' : more) + else (body, more) + in case reads num :: [(Double, String)] of + [(d, "")] -> (TNum d (any (`elem` (".eE" :: String)) num) :) <$> go more' + _ -> Left ("bad number " ++ num) + +-- Precedence climbing over the printer's precedences. +type P a = [Tok] -> Either String (a, [Tok]) + +parseExpr :: P Ast +parseExpr = parseOr + +parseOr + , parseAnd + , parseCmp + , parseAdd + , parseMul + , parseUnary + , parsePrimary :: + P Ast +parseOr ts = do + (l, rest) <- parseAnd ts + case rest of + TOp ".||." : more -> do + (r, rest') <- parseOr more + pure (ACall "or" [l, r], rest') + _ -> pure (l, rest) +parseAnd ts = do + (l, rest) <- parseCmp ts + case rest of + TOp ".&&." : more -> do + (r, rest') <- parseAnd more + pure (ACall "and" [l, r], rest') + _ -> pure (l, rest) +parseCmp ts = do + (l, rest) <- parseAdd ts + case rest of + TOp op : more | Just name <- lookup op cmpNames -> do + (r, rest') <- parseAdd more + pure (ACall name [l, r], rest') + _ -> pure (l, rest) + where + cmpNames = + [ (".<=.", "leq") + , (".>=.", "geq") + , (".==.", "eq") + , ("./=.", "neq") + , (".<.", "lt") + , (".>.", "gt") + , ("<=", "leq") + , (">=", "geq") + , ("==", "eq") + , ("/=", "neq") + , ("<", "lt") + , (">", "gt") + ] +parseAdd ts = parseMul ts >>= uncurry loop + where + loop l (TOp "+" : more) = parseMul more >>= \(r, rest) -> loop (ACall "add" [l, r]) rest + loop l (TOp "-" : more) = parseMul more >>= \(r, rest) -> loop (ACall "sub" [l, r]) rest + loop l rest = pure (l, rest) +parseMul ts = parseUnary ts >>= uncurry loop + where + loop l (TOp "*" : more) = parseUnary more >>= \(r, rest) -> loop (ACall "mult" [l, r]) rest + loop l (TOp "/" : more) = parseUnary more >>= \(r, rest) -> loop (ACall "divide" [l, r]) rest + loop l (TOp "//" : more) = parseUnary more >>= \(r, rest) -> loop (ACall "div" [l, r]) rest + loop l rest = pure (l, rest) +parseUnary (TOp "~" : more) = do + (e, rest) <- parseUnary more + pure (ACall "not" [e], rest) +parseUnary (TOp "-" : more) = do + (e, rest) <- parseUnary more + pure (ACall "negate" [e], rest) +parseUnary ts = parsePrimary ts +parsePrimary ts = case ts of + TNum v isD : rest -> pure (ANum v isD, rest) + TStr s : rest -> pure (AStr s, rest) + TName n : rest -> pure (ACol n Nothing, rest) + TFun f : TLParen : rest -> do + (args, rest') <- parseArgs rest [] + pure (ACall f args, rest') + TLParen : rest -> do + (e, rest') <- parseExpr rest + case rest' of + TRParen : more -> pure (e, more) + _ -> Left "expected )" + TIf : rest -> parseIf rest + _ -> Left ("unexpected token " ++ show (take 3 ts)) + where + parseArgs toks acc = do + (e, rest) <- parseExpr toks + case rest of + TComma : more -> parseArgs more (e : acc) + TRParen : more -> pure (reverse (e : acc), more) + _ -> Left "expected , or ) in call" + +-- if C T else if C2 T2 else E +parseIf :: P Ast +parseIf ts = do + (c, rest) <- parseExpr ts + (t, rest') <- parseExpr rest + case rest' of + TElse : TIf : more -> do + (e, rest'') <- parseIf more + pure (AIf c t e, rest'') + TElse : more -> do + (e, rest'') <- parseExpr more + pure (AIf c t e, rest'') + _ -> Left "expected else" + +-- | Parse the output of 'prettyPrint', given the frame's column types. +parsePretty :: Schema -> String -> Either String SomeExpr +parsePretty schema s = do + toks <- lexPretty schema s + (ast, rest) <- parseExpr toks + if null rest then pure () else Left ("trailing tokens: " ++ show (take 5 rest)) + elab schema ast + +parseShownAs :: forall a. (Columnable a) => String -> Either String (Expr a) +parseShownAs s = parseShown s >>= asType + +parsePrettyAs :: + forall a. (Columnable a) => Schema -> String -> Either String (Expr a) +parsePrettyAs schema s = parsePretty schema s >>= asType + +asType :: forall a. (Columnable a) => SomeExpr -> Either String (Expr a) +asType u = + maybe + (Left "expression has a different type from the one requested") + pure + (fromSomeExpr u) diff --git a/src/DataFrame/Functions.hs b/src/DataFrame/Functions.hs --- a/src/DataFrame/Functions.hs +++ b/src/DataFrame/Functions.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE DisambiguateRecordFields #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} @@ -11,7 +12,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 @@ -25,20 +34,23 @@ import qualified Data.List as L import qualified Data.Map as M import qualified Data.Maybe as Maybe +import qualified Data.Set as S 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 @@ -63,12 +75,7 @@ } ) -{- | Lift a unary function over a nullable or non-nullable column expression. -When the input is @Maybe a@, 'Nothing' short-circuits (like 'fmap'). -When the input is plain @a@, the function is applied directly. - -The return type is inferred via 'NullLift1Result': no annotation needed. --} +-- | Lift a unary function over a nullable or non-nullable column expression. nullLift :: (NullLift1Op a r (NullLift1Result a r), Columnable (NullLift1Result a r)) => (BaseType a -> r) -> @@ -78,11 +85,7 @@ Unary (MkUnaryOp{unaryFn = applyNull1 f, unaryName = "nullLift", unarySymbol = Nothing}) -{- | Lift a binary function over nullable or non-nullable column expressions. -Any 'Nothing' operand short-circuits to 'Nothing' in the result. - -The return type is inferred via 'NullLift2Result': no annotation needed. --} +-- | Lift a binary function over nullable or non-nullable column expressions. nullLift2 :: (NullLift2Op a b r (NullLift2Result a b r), Columnable (NullLift2Result a b r)) => (BaseType a -> BaseType b -> r) -> @@ -100,38 +103,21 @@ } ) -{- | Lenient numeric \/ text coercion returning @Maybe a@. Looks up column -@name@ and coerces its values to @a@. Values that cannot be converted -(parse failures, type mismatches) become 'Nothing'; successfully converted -values are wrapped in 'Just'. Existing 'Nothing' in optional source columns -stays as 'Nothing'. --} +-- | Lenient coercion of a column to @Maybe a@; failures become 'Nothing'. cast :: forall a. (Columnable a, Read a) => T.Text -> Expr (Maybe a) cast colName = CastWith colName "cast" (either (const Nothing) Just) -{- | Lenient coercion that substitutes a default for unconvertible values. -Looks up column @name@, coerces its values to @a@, and uses @def@ wherever -conversion fails or the source value is 'Nothing'. --} +-- | Lenient coercion that substitutes a default for unconvertible values. castWithDefault :: forall a. (Columnable a, Read a) => a -> T.Text -> Expr a castWithDefault def colName = CastWith colName ("castWithDefault:" <> T.pack (show def)) (fromRight def) -{- | Lenient coercion returning @Either T.Text a@. Successfully converted -values are 'Right'; values that cannot be parsed are kept as 'Left' with -their original string representation, so the caller can inspect or handle -them downstream. Existing 'Nothing' in optional source columns becomes -@Left \"null\"@. --} +-- | Lenient coercion to @Either T.Text a@, keeping unparsed values as 'Left'. castEither :: forall a. (Columnable a, Read a) => T.Text -> Expr (Either T.Text a) castEither colName = CastWith colName "castEither" (either (Left . T.pack) Right) -{- | Lenient coercion for assertedly non-nullable columns. -Substitutes @error@ for @Nothing@, so it will crash at evaluation time if -any @Nothing@ is actually encountered. For non-nullable and -fully-populated nullable columns no cost is paid. --} +-- | Lenient coercion that crashes at evaluation if a @Nothing@ is found. unsafeCast :: forall a. (Columnable a, Read a) => T.Text -> Expr a unsafeCast colName = CastWith @@ -226,6 +212,25 @@ {-# SPECIALIZE count :: Expr Int64 -> Expr Int #-} {-# INLINEABLE count #-} +-- | Number of distinct values, the equivalent of SQL's @COUNT(DISTINCT x)@. +countDistinct :: forall a. (Columnable a, Ord a) => Expr a -> Expr Int +countDistinct = + Agg + ( MergeAgg + "countDistinct" + (S.empty :: S.Set a) + (flip S.insert) + S.union + S.size + ) +{-# SPECIALIZE countDistinct :: Expr Double -> Expr Int #-} +{-# SPECIALIZE countDistinct :: Expr Float -> Expr Int #-} +{-# SPECIALIZE countDistinct :: Expr Int -> Expr Int #-} +{-# SPECIALIZE countDistinct :: Expr Int8 -> Expr Int #-} +{-# SPECIALIZE countDistinct :: Expr Int16 -> Expr Int #-} +{-# SPECIALIZE countDistinct :: Expr Int32 -> Expr Int #-} +{-# SPECIALIZE countDistinct :: Expr Int64 -> Expr Int #-} + -- | Row count, the equivalent of SQL's @COUNT(*)@. countAll :: Expr Int countAll = count (Lit (0 :: Int)) @@ -307,10 +312,10 @@ Agg ( MergeAgg "mean" - (MeanAcc 0.0 0) - (\(MeanAcc s c) x -> MeanAcc (s + realToFrac x) (c + 1)) - (\(MeanAcc s1 c1) (MeanAcc s2 c2) -> MeanAcc (s1 + s2) (c1 + c2)) - (\(MeanAcc s c) -> if c == 0 then 0 / 0 else s / fromIntegral c) + ((0.0, 0) :: (Double, Int)) + (\(!s, !c) x -> (s + realToFrac x, c + 1)) + (\(!s1, !c1) (!s2, !c2) -> (s1 + s2, c1 + c2)) + (\(s, c) -> if c == 0 then 0 / 0 else s / fromIntegral c) ) {-# SPECIALIZE mean :: Expr Double -> Expr Double #-} {-# SPECIALIZE mean :: Expr Float -> Expr Double #-} @@ -321,6 +326,45 @@ {-# 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, dropping NaNs. +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 #-} @@ -406,7 +450,7 @@ zScore c = (c - mean c) / stddev c pow :: (Columnable a, Num a) => Expr a -> Int -> Expr a -pow expr i = lift2Decorated (^) "pow" (Just "^") True 8 expr (Lit i) +pow expr i = lift2Decorated (^) "pow" (Just "^") False 8 expr (Lit i) {-# SPECIALIZE pow :: Expr Double -> Int -> Expr Double #-} {-# SPECIALIZE pow :: Expr Float -> Int -> Expr Float #-} {-# SPECIALIZE pow :: Expr Int -> Int -> Expr Int #-} @@ -477,6 +521,16 @@ {-# 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 {-# SPECIALIZE fromJust :: Expr (Maybe Double) -> Expr Double #-} @@ -550,6 +604,15 @@ lastOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a) lastOrNothing = liftDecorated (Maybe.listToMaybe . reverse) "lastOrNothing" Nothing +-- | Whether the text contains the literal substring. +contains :: T.Text -> Expr T.Text -> Expr Bool +contains needle expr = + lift2Decorated (flip T.isInfixOf) "contains" Nothing False 4 expr (Lit needle) + +-- | Character count. +length :: Expr T.Text -> Expr Int +length = liftDecorated T.length "textLength" Nothing + splitOn :: T.Text -> Expr T.Text -> Expr [T.Text] splitOn delim = liftDecorated (T.splitOn delim) "splitOn" Nothing @@ -581,7 +644,7 @@ (\d1 d2 -> fromIntegral (diffDays d1 d2)) "daysBetween" Nothing - True + False 2 bind :: @@ -592,24 +655,11 @@ Expr (m b) bind f = liftDecorated (>>= f) "bind" Nothing -{- | Window function: evaluate an expression partitioned by the given columns. - -Each partition computes the inner expression independently, and the result -is broadcast back to every row in that partition. This is analogous to -Polars' @.over()@ or SQL @OVER (PARTITION BY ...)@. - -@ --- Per-country median, broadcast to every row: -F.over [\"country\"] (F.median (F.col \@Double \"amount\")) - --- Deviation from group mean: -F.col \@Double \"amount\" - F.over [\"group\"] (F.mean (F.col \@Double \"amount\")) -@ --} +-- | Window function: evaluate an expression partitioned by the given columns. over :: (Columnable a) => [T.Text] -> Expr a -> Expr a over = Over --- See Section 2.4 of the Haskell Report https://www.haskell.org/definition/haskell2010.pdf +-- Haskell Report 2.4: https://www.haskell.org/definition/haskell2010.pdf isReservedId :: T.Text -> Bool isReservedId t = case t of "case" -> True @@ -638,11 +688,7 @@ isVarId :: T.Text -> Bool isVarId t = case T.uncons t of - -- We might want to check c == '_' || Char.isLower c - -- since the haskell report considers '_' a lowercase character - -- However, to prevent an edge case where a user may have a - -- "Name" and an "_Name_" in the same scope, wherein we'd end up - -- with duplicate "_Name_"s, we eschew the check for '_' here. + -- '_' is skipped on purpose so "Name" and "_Name_" cannot collide. Just (c, _) -> Char.isLower c && Char.isAlpha c Nothing -> False diff --git a/src/DataFrame/Internal/Statistics.hs b/src/DataFrame/Internal/Statistics.hs deleted file mode 100644 --- a/src/DataFrame/Internal/Statistics.hs +++ /dev/null @@ -1,285 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} - -module DataFrame.Internal.Statistics where - -import qualified Data.Vector as V -import qualified Data.Vector.Algorithms.Intro as VA -import qualified Data.Vector.Mutable as VM -import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Unboxed.Mutable as VUM - -import Control.Exception (throw) -import Control.Monad.ST (runST) -import DataFrame.Errors (DataFrameException (..)) - -mean' :: (Real a, VU.Unbox a) => VU.Vector a -> Double -mean' samp - | VU.null samp = throw $ EmptyDataSetException "mean" - | otherwise = rtf (VU.sum samp) / fromIntegral (VU.length samp) -{-# INLINE [0] mean' #-} - -meanDouble' :: VU.Vector Double -> Double -meanDouble' samp - | VU.null samp = throw $ EmptyDataSetException "mean" - | otherwise = VU.sum samp / fromIntegral (VU.length samp) -{-# INLINE meanDouble' #-} - -meanInt' :: VU.Vector Int -> Double -meanInt' samp - | VU.null samp = throw $ EmptyDataSetException "mean" - | otherwise = fromIntegral (VU.sum samp) / fromIntegral (VU.length samp) -{-# INLINE meanInt' #-} - -{-# RULES -"mean'/Double" [1] forall (xs :: VU.Vector Double). - mean' xs = - meanDouble' xs -"mean'/Int" [1] forall (xs :: VU.Vector Int). - mean' xs = - meanInt' xs - #-} - -median' :: (Real a, VU.Unbox a) => VU.Vector a -> Double -median' samp - | VU.null samp = throw $ EmptyDataSetException "median" - | otherwise = runST $ do - mutableSamp <- VU.thaw samp - VA.sort mutableSamp - let len = VU.length samp - middleIndex = len `div` 2 - middleElement <- VUM.read mutableSamp middleIndex - if odd len - then pure (rtf middleElement) - else do - prev <- VUM.read mutableSamp (middleIndex - 1) - pure (rtf (middleElement + prev) / 2) -{-# INLINE median' #-} - --- accumulator: count, mean, m2 -data VarAcc - = VarAcc {-# UNPACK #-} !Int {-# UNPACK #-} !Double {-# UNPACK #-} !Double - deriving (Show) - -varianceStep :: VarAcc -> Double -> VarAcc -varianceStep (VarAcc !n !meanVal !m2) !x = - let !n' = n + 1 - !delta = x - meanVal - !meanVal' = meanVal + delta / fromIntegral n' - !m2' = m2 + delta * (x - meanVal') - in VarAcc n' meanVal' m2' -{-# INLINE varianceStep #-} - -computeVariance :: VarAcc -> Double -computeVariance (VarAcc !n _ !m2) - | n < 2 = 0 -- or error "variance of <2 samples" - | otherwise = m2 / fromIntegral (n - 1) -{-# INLINE computeVariance #-} - -variance' :: (Real a, VU.Unbox a) => VU.Vector a -> Double -variance' = computeVariance . VU.foldl' varianceStep (VarAcc 0 0 0) . VU.map rtf -{-# INLINE variance' #-} - -varianceDouble' :: VU.Vector Double -> Double -varianceDouble' = computeVariance . VU.foldl' varianceStep (VarAcc 0 0 0) -{-# INLINE varianceDouble' #-} - --- accumulator: count, mean, m2, m3 -data SkewAcc = SkewAcc !Int !Double !Double !Double deriving (Show) - -skewnessStep :: (VU.Unbox a, Num a, Real a) => SkewAcc -> a -> SkewAcc -skewnessStep (SkewAcc !n !meanVal !m2 !m3) !x' = - let !n' = n + 1 - x = rtf x' - !k = fromIntegral n' - !delta = x - meanVal - !meanVal' = meanVal + delta / k - !m2' = m2 + (delta ^ (2 :: Int) * (k - 1)) / k - !m3' = - m3 - + (delta ^ (3 :: Int) * (k - 1) * (k - 2)) / k ^ (2 :: Int) - - (3 * delta * m2) / k - in SkewAcc n' meanVal' m2' m3' -{-# INLINE skewnessStep #-} - -computeSkewness :: SkewAcc -> Double -computeSkewness (SkewAcc n _ m2 m3) - | n < 3 = 0 -- or error "skewness of <3 samples" - | otherwise = (sqrt (fromIntegral n - 1) * m3) / sqrt (m2 ^ (3 :: Int)) -{-# INLINE computeSkewness #-} - -skewness' :: (VU.Unbox a, Real a, Num a) => VU.Vector a -> Double -skewness' = computeSkewness . VU.foldl' skewnessStep (SkewAcc 0 0 0 0) -{-# INLINE skewness' #-} - -data CorrelationStats - = CorrelationStats - {-# UNPACK #-} !Double - {-# UNPACK #-} !Double - {-# UNPACK #-} !Double - {-# UNPACK #-} !Double - {-# UNPACK #-} !Double - -correlation' :: VU.Vector Double -> VU.Vector Double -> Maybe Double -correlation' xs ys - | n < 2 = Nothing - | VU.length xs /= VU.length ys = Nothing - | otherwise = - let nf = fromIntegral n - initial = CorrelationStats 0 0 0 0 0 - (CorrelationStats sumX sumY sumXX sumYY sumXY) = VU.ifoldl' step initial xs - - !num = nf * sumXY - sumX * sumY - !den = sqrt ((nf * sumXX - sumX * sumX) * (nf * sumYY - sumY * sumY)) - in Just (num / den) - where - n = VU.length xs - step (CorrelationStats sx sy sxx syy sxy) i x = - let !y = VU.unsafeIndex ys i - in CorrelationStats (sx + x) (sy + y) (sxx + x * x) (syy + y * y) (sxy + x * y) -{-# INLINE correlation' #-} - -quantiles' :: - (VU.Unbox a, Num a, Real a) => - VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double -quantiles' qs q samp - | VU.null samp = throw $ EmptyDataSetException "quantiles" - | q < 2 = throw $ WrongQuantileNumberException q - | VU.any (\i -> i < 0 || i > q) qs = throw $ WrongQuantileIndexException qs q - | otherwise = runST $ do - let !n = VU.length samp - mutableSamp <- VU.thaw samp - VA.sort mutableSamp - VU.mapM - ( \i -> do - let !p = fromIntegral i / fromIntegral q - !position = p * fromIntegral (n - 1) :: Double - !index = floor position :: Int - !f = position - fromIntegral index - x <- fmap rtf (VUM.read mutableSamp index) - if f == 0 - then return x - else do - y <- fmap rtf (VUM.read mutableSamp (index + 1)) - return $ (1 - f) * x + f * y - ) - qs -{-# INLINE quantiles' #-} - -percentile' :: (VU.Unbox a, Num a, Real a) => Int -> VU.Vector a -> Double -percentile' n = VU.head . quantiles' (VU.fromList [n]) 100 - -quantilesOrd' :: - (Ord a, Eq a) => - VU.Vector Int -> Int -> V.Vector a -> V.Vector a -quantilesOrd' qs q samp - | V.null samp = throw $ EmptyDataSetException "quantiles" - | q < 2 = throw $ WrongQuantileNumberException q - | VU.any (\i -> i < 0 || i > q) qs = throw $ WrongQuantileIndexException qs q - | otherwise = runST $ do - let !n = V.length samp - mutableSamp <- V.thaw samp - VA.sort mutableSamp - V.mapM - ( \i -> do - let !p = fromIntegral i / fromIntegral q :: Double - !position = p * fromIntegral (n - 1) - !index = floor position :: Int - -- This is not exact for Ord instances. - -- Figure out how to make it so. - VM.read mutableSamp index - ) - (V.convert qs) - -percentileOrd' :: (Ord a, Eq a) => Int -> V.Vector a -> a -percentileOrd' n = V.head . quantilesOrd' (VU.fromList [n]) 100 - -interQuartileRange' :: (VU.Unbox a, Num a, Real a) => VU.Vector a -> Double -interQuartileRange' samp = - let quartiles = quantiles' (VU.fromList [1, 3]) 4 samp - in quartiles VU.! 1 - quartiles VU.! 0 -{-# INLINE interQuartileRange' #-} - -meanSquaredError :: VU.Vector Double -> VU.Vector Double -> Maybe Double -meanSquaredError target prediction = - let - squareDiff = VU.ifoldl' (\sq i e -> (e - target VU.! i) ^ (2 :: Int) + sq) 0 prediction - in - Just $ squareDiff / fromIntegral (max (VU.length target) (VU.length prediction)) -{-# INLINE meanSquaredError #-} - -mutualInformationBinned :: - Int -> VU.Vector Double -> VU.Vector Double -> Maybe Double -mutualInformationBinned k xs ys - | VU.length xs /= VU.length ys = Nothing - | VU.null xs = Nothing - | k < 2 = Nothing - | rx <= 0 || ry <= 0 = Just 0 - | otherwise = - let bx = VU.map (binIndex xmin xmax k) xs - by = VU.map (binIndex ymin ymax k) ys - n = fromIntegral (VU.length xs) :: Double - mx = bincount k bx - my = bincount k by - mxy = jointBincount k bx by - in Just $ - sum - [ let !cxy = fromIntegral c - !pxy = cxy / n - !px = fromIntegral (mx VU.! i) / n - !py = fromIntegral (my VU.! j) / n - in if c == 0 then 0 else pxy * logBase 2 (pxy / (px * py)) - | i <- [0 .. k - 1] - , j <- [0 .. k - 1] - , let !c = mxy VU.! (i * k + j) - ] - where - (xmin, xmax) = (VU.minimum xs, VU.maximum xs) - (ymin, ymax) = (VU.minimum ys, VU.maximum ys) - rx = xmax - xmin - ry = ymax - ymin - -binIndex :: Double -> Double -> Int -> Double -> Int -binIndex lo hi k x - | hi == lo = 0 - | otherwise = - let !t = (x - lo) / (hi - lo) - !ix = floor (fromIntegral k * t) :: Int - in max 0 (min (k - 1) ix) -{-# INLINE binIndex #-} - -bincount :: Int -> VU.Vector Int -> VU.Vector Int -bincount k bs = VU.create $ do - mv <- VU.thaw (VU.replicate k 0) - VU.forM_ bs $ \b -> do - let i - | b < 0 = 0 - | b >= k = k - 1 - | otherwise = b - x <- VUM.read mv i - VUM.write mv i (x + 1) - pure mv -{-# INLINE bincount #-} - -jointBincount :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int -jointBincount k bx by = VU.create $ do - mv <- VU.thaw (VU.replicate (k * k) 0) - VU.forM_ (VU.zip bx by) $ \(i, j) -> do - let ii = clamp i 0 (k - 1) - jj = clamp j 0 (k - 1) - ix = ii * k + jj - x <- VUM.read mv ix - VUM.write mv ix (x + 1) - pure mv - where - clamp z a b = max a (min b z) -{-# INLINE jointBincount #-} - -rtf :: (Real a) => a -> Double -rtf = realToFrac -{-# NOINLINE [1] rtf #-} - -{-# RULES -"rtf/Double" [2] forall (x :: Double). rtf x = x - #-} diff --git a/src/DataFrame/Monad.hs b/src/DataFrame/Monad.hs --- a/src/DataFrame/Monad.hs +++ b/src/DataFrame/Monad.hs @@ -1,18 +1,54 @@ {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE InstanceSigs #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TupleSections #-} -module DataFrame.Monad where +module DataFrame.Monad ( + -- * The frame-state monad + FrameM, + runFrameM, + evalFrameM, + execFrameM, + modifyM, + inspectM, + -- * Frame verbs + deriveM, + insertM, + renameM, + filterWhereM, + sampleM, + takeM, + dropM, + selectM, + excludeM, + sortByM, + SortOrder (..), + columnAsListM, + filterJustM, + imputeM, + + -- * Column-derivation pipelines + Pipeline, + letAs, + letExpr, + pipelineSteps, + toFrameM, + runPipeline, +) where + +import Control.Monad (void) import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.DataFrame (DataFrame) -import DataFrame.Internal.Expression (Expr (..)) -import DataFrame.Internal.Nullable (BaseType) +import DataFrame.Internal.Expression (Expr (..), SomeExpr (..), 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 @@ -21,9 +57,7 @@ import Data.Tuple (swap) import System.Random --- A re-implementation of the state monad. --- `mtl` might be too heavy a dependency just to get --- a single monad instance. +-- A re-implementation of the state monad, avoiding an `mtl` dependency. newtype FrameM a = FrameM {runFrameM_ :: DataFrame -> (DataFrame, a)} instance Functor FrameM where @@ -81,6 +115,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 +150,53 @@ execFrameM :: DataFrame -> FrameM a -> DataFrame execFrameM df m = snd (runFrameM df m) + +newtype Pipeline a = Pipeline {unPipeline :: Int -> (Int, [(T.Text, SomeExpr)], 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, SomeExpr 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, SomeExpr e)], Col nm) + +instance Show (Pipeline (Expr a)) where + show p = + let (_, steps, res) = unPipeline p 0 + in concatMap + (\(nm, SomeExpr 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, SomeExpr e) -> void (deriveM nm e)) steps >> pure res + +-- | Run a pipeline over a frame, returning the result expression and frame. +runPipeline :: DataFrame -> Pipeline (Expr a) -> (Expr a, DataFrame) +runPipeline df = runFrameM df . toFrameM diff --git a/src/DataFrame/Operations/AggregateScatter.hs b/src/DataFrame/Operations/AggregateScatter.hs deleted file mode 100644 --- a/src/DataFrame/Operations/AggregateScatter.hs +++ /dev/null @@ -1,224 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ExplicitNamespaces #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - -{- | Execute a recognised aggregation plan ('AggPlan') through the vectorized -scatter kernel, producing one result column (length @nGroups@, canonical group -order). The scatter reductions live in 'DataFrame.Internal.AggKernel' (sequential) -and 'DataFrame.Internal.AggKernelPar' (parallel by disjoint group range); this -module handles the compound @max - min@ combine and the holistic grouped median. -A plan only reaches here once 'planAgg' verified the value columns are clean -unboxed Int/Double, so the @error@ branches are unreachable. - -Every reduction takes the Round-5 grouping layout @(valueIndices, offsets)@ so -the parallel kernel can split the group-id range across capabilities with no -cross-worker merge. Each group's rows stay in original-row order within one -worker's range, so results are byte-identical to the sequential path at any @-N@. --} -module DataFrame.Operations.AggregateScatter (runPlan, runMomentPlan) where - -import qualified Data.Text as T -import qualified Data.Vector.Algorithms.Intro as VA -import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Unboxed.Mutable as VUM - -import Control.Concurrent (forkIO, getNumCapabilities) -import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar) -import Control.Exception (SomeException, throwIO, try) -import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) -import DataFrame.Internal.AggKernel (Reduction (..), scatterColumnToDouble) -import DataFrame.Internal.AggKernelDirect (directReduce, directThreshold) -import DataFrame.Internal.AggKernelPar (momentScatterPar, scatterReducePar) -import DataFrame.Internal.AggPlan (AggPlan (..), MomentPlan (..), Moments (..)) -import DataFrame.Internal.Column (Column (..), fromUnboxedVector) -import DataFrame.Internal.DataFrame (GroupedDataFrame (..), getColumn) -import System.IO.Unsafe (unsafePerformIO) -import Type.Reflection (typeRep) - -runPlan :: GroupedDataFrame -> VU.Vector Int -> Int -> AggPlan -> Column -runPlan gdf rtg nGroups plan = case plan of - PlanScatter red name -> scatterColumn red name - PlanMaxMinusMin a b -> maxMinusMin vis offs nGroups (col a) (col b) - PlanMedian name -> groupedMedian vis offs nGroups (col name) - where - vis = valueIndices gdf - offs = offsets gdf - {- The low-cardinality DIRECT-INDEXED fast path: for a small dense domain the - grouping layer's @rowToGroup@ already maps row -> group, so we scatter - straight off it (no @valueIndices@ gather). 'directReduce' only admits - order-independent reductions (so the merged parallel result is byte-identical - to -N1); anything it rejects keeps the order-preserving group-range kernel. -} - scatterColumn red name = - let c = col name - direct - | nGroups <= directThreshold = directReduce red rtg nGroups c - | otherwise = Nothing - in case direct of - Just out -> out - Nothing -> case scatterReducePar red vis offs nGroups c of - Just out -> out - Nothing -> error "runPlan: scatterReducePar rejected a planned column" - col name = case getColumn name (fullDataframe gdf) of - Just c -> c - Nothing -> error ("runPlan: planned column missing: " ++ T.unpack name) - -{- | Run a recognised moment (Q9 regression) plan as one fused scatter over the -two base columns, returning each output name bound to its moment field. The six -sufficient statistics (count, Sx, Sy, Sxx, Syy, Sxy) come out of a single pass, -replacing the three derive passes and six independent scatters of the -per-expression path. Byte-identical to the sequential kernel at any @-N@. --} -runMomentPlan :: - GroupedDataFrame -> Int -> MomentPlan -> Maybe [(T.Text, Column)] -runMomentPlan gdf nGroups mp = do - ms <- momentScatterPar vis offs nGroups (col (mpColX mp)) (col (mpColY mp)) - pure - [ (mpNName mp, mN ms) - , (mpSxName mp, mSx ms) - , (mpSyName mp, mSy ms) - , (mpSxxName mp, mSxx ms) - , (mpSyyName mp, mSyy ms) - , (mpSxyName mp, mSxy ms) - ] - where - vis = valueIndices gdf - offs = offsets gdf - col name = case getColumn name (fullDataframe gdf) of - Just c -> c - Nothing -> error ("runMomentPlan: planned column missing: " ++ T.unpack name) - -{- | @max a - min b@ on the small @nGroups@ arrays. Preserves the Int element -type of the source columns (matching the interpreter), falling back to a Double -combine otherwise. --} -maxMinusMin :: - VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -> Column -maxMinusMin vis offs nGroups ca cb = - case (ca, cb) of - ( UnboxedColumn Nothing (_ :: VU.Vector x) - , UnboxedColumn Nothing (_ :: VU.Vector y) - ) - | Just Refl <- testEquality (typeRep @x) (typeRep @Int) - , Just Refl <- testEquality (typeRep @y) (typeRep @Int) -> - let mx = scatterExtremaInt RMax vis offs nGroups ca - mn = scatterExtremaInt RMin vis offs nGroups cb - in fromUnboxedVector (VU.zipWith (-) mx mn) - _ -> - let mx = scatterExtremaDbl RMax vis offs nGroups ca - mn = scatterExtremaDbl RMin vis offs nGroups cb - in fromUnboxedVector (VU.zipWith (-) mx mn) - -scatterExtremaInt :: - Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> VU.Vector Int -scatterExtremaInt red vis offs nGroups c = case scatterReducePar red vis offs nGroups c of - Just (UnboxedColumn _ (v :: VU.Vector a)) - | Just Refl <- testEquality (typeRep @a) (typeRep @Int) -> v - _ -> error "scatterExtremaInt" - -scatterExtremaDbl :: - Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> VU.Vector Double -scatterExtremaDbl red vis offs nGroups c = - case scatterReducePar red vis offs nGroups c of - Just (UnboxedColumn _ (v :: VU.Vector a)) - | Just Refl <- testEquality (typeRep @a) (typeRep @Double) -> v - | Just Refl <- testEquality (typeRep @a) (typeRep @Int) -> VU.map fromIntegral v - _ -> error "scatterExtremaDbl" - -------------------------------------------------------------------------------- --- Parallel holistic median -------------------------------------------------------------------------------- - -{- | Holistic per-group median over a single unboxed Int/Double column. The -@valueIndices@/@offsets@ layout already places each group's rows in a contiguous -run, so we copy each group's values into a scratch buffer at its own offset and -sort that slice in place — each group's slice is independent, so the per-group -sorts split across capabilities by group range with no merge. Empty groups never -occur, so the result is total. --} -groupedMedian :: VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -groupedMedian vis offs nGroups c = case scatterColumnToDouble c of - Nothing -> error "groupedMedian: non-numeric planned column" - Just vals -> fromUnboxedVector (medianByGroup vis offs nGroups vals) - -medianByGroup :: - VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double -medianByGroup vis offs nGroups vals = unsafePerformIO $ do - let !n = VU.length vis - buf <- VUM.new (max 1 n) - out <- VUM.new (max 1 nGroups) - caps <- getNumCapabilities - let !bounds = groupRangeBounds offs nGroups caps - -- Each worker fills+sorts the buffer slices of its own group range, then - -- writes that range's medians. Disjoint ranges => safe to parallelise. - forEachRange bounds caps $ \gs ge -> - let grp !g - | g >= ge = pure () - | otherwise = do - let !s = VU.unsafeIndex offs g - !e = VU.unsafeIndex offs (g + 1) - !len = e - s - fill !pos - | pos >= e = pure () - | otherwise = do - VUM.unsafeWrite buf pos (VU.unsafeIndex vals (VU.unsafeIndex vis pos)) - fill (pos + 1) - fill s - let slice = VUM.unsafeSlice s len buf - VA.sort slice - let mid = s + len `div` 2 - med <- - if odd len - then VUM.unsafeRead buf mid - else do - hi <- VUM.unsafeRead buf mid - lo <- VUM.unsafeRead buf (mid - 1) - pure ((hi + lo) / 2) - VUM.unsafeWrite out g med - grp (g + 1) - in grp gs - VU.unsafeFreeze (VUM.unsafeSlice 0 nGroups out) -{-# NOINLINE medianByGroup #-} - -------------------------------------------------------------------------------- --- Group-range partitioning (shared with the median path) -------------------------------------------------------------------------------- - -{- | Split @[0, nGroups)@ into @caps@ contiguous group ranges balanced by row -count. Identical policy to 'DataFrame.Internal.AggKernelPar.groupRangeBounds'. --} -groupRangeBounds :: VU.Vector Int -> Int -> Int -> VU.Vector Int -groupRangeBounds offs nGroups caps = VU.create $ do - b <- VUM.new (caps + 1) - let !nRows = VU.unsafeIndex offs nGroups - !per = max 1 ((nRows + caps - 1) `div` caps) - adv !target !gg - | gg >= nGroups = nGroups - | VU.unsafeIndex offs gg >= target = gg - | otherwise = adv target (gg + 1) - go !w !prev - | w >= caps = VUM.unsafeWrite b caps nGroups - | otherwise = do - let !target = min nRows (w * per) - !g = adv target prev - VUM.unsafeWrite b w g - go (w + 1) g - VUM.unsafeWrite b 0 0 - go 1 0 - pure b - -forEachRange :: VU.Vector Int -> Int -> (Int -> Int -> IO ()) -> IO () -forEachRange bounds caps act - | caps <= 1 = act (VU.unsafeIndex bounds 0) (VU.unsafeIndex bounds caps) - | otherwise = do - vars <- mapM spawn [0 .. caps - 1] - results <- mapM takeMVar vars - mapM_ (either (throwIO :: SomeException -> IO ()) pure) results - where - spawn w = do - var <- newEmptyMVar - let !s = VU.unsafeIndex bounds w - !e = VU.unsafeIndex bounds (w + 1) - _ <- forkIO (try (act s e) >>= putMVar var) - pure var diff --git a/src/DataFrame/Operations/Aggregation.hs b/src/DataFrame/Operations/Aggregation.hs --- a/src/DataFrame/Operations/Aggregation.hs +++ b/src/DataFrame/Operations/Aggregation.hs @@ -14,38 +14,55 @@ 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) -{- | Per-row key hash over the selected key columns. Delegates to the shared -'computeRowHashesIO' kernel, which forks over contiguous row ranges for large -frames (the hashing of a wide 1e7-row text/factor join key dominates that join) -and is bit-for-bit identical to a single sequential pass at any capability count. --} +-- | Per-row key hash over the selected key columns, computed in parallel. computeRowHashes :: [Int] -> DataFrame -> VU.Vector Int computeRowHashes indices df = let n = fst (dimensions df) @@ -53,33 +70,104 @@ in unsafePerformIO (computeRowHashesIO n selectedCols) {-# NOINLINE computeRowHashes #-} -{- | Aggregate a grouped dataframe using the expressions given. -All ungrouped columns will be dropped. --} +-- | Aggregate a grouped dataframe using the expressions given. 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: one fused parallel gather over the representative rows. 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) + -- Under -XStrict this must stay cheap when the moment path covers it. + 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 -- Even a single candidate uses the gather path. + if not (null cands) + then + MS.fromList + ( zip + (map fst cands) + ( runGatherAggs + (valueIndices gdf) + offs + nGroups + (map snd cands) + ) + ) + else MS.empty + + -- Median and std/var on one column share a single gather. + 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 - -- of two base columns) becomes one scatter over the base columns, - -- dropping the derived product columns and the six separate folds. - -- 'planMoments' returns 'Nothing' on any other set, falling back below. + -- Fused fast path: the Q9 regression family becomes one scatter. fusedMoments = do mp <- planMoments gdf aggs :: Maybe MomentPlan runMomentPlan gdf nGroups mp @@ -90,7 +178,7 @@ -- | The fall-back path: evaluate one named aggregation via the interpreter. interpretNamed :: GroupedDataFrame -> NamedExpr -> Column -interpretNamed gdf (_, UExpr (expr :: Expr a)) = +interpretNamed gdf (_, SomeExpr (expr :: Expr a)) = case interpretAggregation @a gdf expr of Left e -> throw e Right (UnAggregated _) -> throw $ UnaggregatedException (T.pack $ show expr) @@ -107,4 +195,5 @@ distinct :: DataFrame -> DataFrame distinct df = selectIndices (VU.map (indices VU.!) (VU.init os)) df where - (Grouped _ _ indices os _rtg) = groupBy (columnNames df) df + -- Wildcard last field: a named pattern would force the deferred thunk. + (Grouped _ _ indices os _) = groupBy (columnNames df) df diff --git a/src/DataFrame/Operations/Aggregation/Run.hs b/src/DataFrame/Operations/Aggregation/Run.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Aggregation/Run.hs @@ -0,0 +1,288 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Execute a recognised aggregation plan, producing one result column. +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.Distinct (countDistinctGroups) +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 kernel is chosen. +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 the fused kernels are exact. + 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) + PlanCountDistinct name -> case countDistinctGroups rtg nGroups (col name) of + Just counts -> fromUnboxedVector counts + Nothing -> error "runPlan: countDistinctGroups rejected a planned column" + where + vis = valueIndices gdf + offs = offsets gdf + -- Dense fast path: scatter straight off the eager @rowToGroup@. + scatterColumn red name = + let c = col name + dense + | nGroups <= denseThreshold = 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 plan as one fused scatter over two columns. +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. + 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@ over the small per-group arrays. +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 grouped median and variance family over the same column. +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'. + 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 #-} + +-- | Holistic per-group median over a single unboxed Int or Double column. +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 sorts and takes medians of its own disjoint group range. + 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, unordered. + 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 #-} + +-- | Split the groups into contiguous ranges balanced by row count. +groupRangeBounds :: VU.Vector Int -> Int -> Int -> VU.Vector Int +groupRangeBounds offs nGroups caps = VU.create $ do + b <- VUM.new (caps + 1) + let !nRows = VU.unsafeIndex offs nGroups + !per = max 1 ((nRows + caps - 1) `div` caps) + adv !target !gg + | gg >= nGroups = nGroups + | VU.unsafeIndex offs gg >= target = gg + | otherwise = adv target (gg + 1) + go !w !prev + | w >= caps = VUM.unsafeWrite b caps nGroups + | otherwise = do + let !target = min nRows (w * per) + !g = adv target prev + VUM.unsafeWrite b w g + go (w + 1) g + VUM.unsafeWrite b 0 0 + go 1 0 + pure b diff --git a/src/DataFrame/Operations/Core.hs b/src/DataFrame/Operations/Core.hs --- a/src/DataFrame/Operations/Core.hs +++ b/src/DataFrame/Operations/Core.hs @@ -6,8 +6,47 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -module DataFrame.Operations.Core where +module DataFrame.Operations.Core ( + -- * Dimensions + dimensions, + nRows, + nColumns, + -- * Construction + fromUnnamedColumns, + fromRows, + + -- * Insertion + insert, + insertVector, + insertUnboxedVector, + insertWithDefault, + insertVectorWithDefault, + + -- * Column management + cloneColumn, + rename, + renameMany, + + -- * Inspection + describeColumns, + valueCounts, + valueProportions, + showDerivedExpressions, + + -- * Folds & matrix/vector conversions + fold, + toFloatMatrix, + toDoubleMatrix, + toIntMatrix, + columnAsVector, + columnAsList, + columnAsIntVector, + columnAsDoubleVector, + columnAsFloatVector, + columnAsUnboxedVector, +) where + import qualified Data.List as L import qualified Data.Map as M import qualified Data.Map.Strict as MS @@ -32,7 +71,9 @@ columnTypeString, fromList, fromVector, + materializeMerged, materializePacked, + mergedHead, toDoubleVector, toFloatVector, toIntVector, @@ -45,7 +86,7 @@ columnNames, derivingExpressions, empty, - fromNamedColumns, + fromColumns, getColumn, insertColumn, null, @@ -57,96 +98,22 @@ import Type.Reflection import Prelude hiding (null) -{- | O(1) Get DataFrame dimensions i.e. (rows, columns) - -==== __Example__ -@ ->>> :set -XOverloadedStrings ->>> import qualified DataFrame as D ->>> df = D.fromNamedColumns [("a", D.fromList [1..100]), ("b", D.fromList [1..100]), ("c", D.fromList [1..100])] ->>> D.dimensions df - -(100, 3) -@ --} +-- | O(1) Get DataFrame dimensions, i.e. (rows, columns). dimensions :: DataFrame -> (Int, Int) dimensions = dataframeDimensions {-# INLINE dimensions #-} -{- | O(1) Get number of rows in a dataframe. - -==== __Example__ -@ ->>> :set -XOverloadedStrings ->>> import qualified DataFrame as D ->>> df = D.fromNamedColumns [("a", D.fromList [1..100]), ("b", D.fromList [1..100]), ("c", D.fromList [1..100])] ->>> D.nRows df -100 -@ --} +-- | O(1) Get the number of rows in a dataframe. nRows :: DataFrame -> Int nRows = fst . dataframeDimensions -{- | O(1) Get number of columns in a dataframe. - -==== __Example__ -@ ->>> :set -XOverloadedStrings ->>> import qualified DataFrame as D ->>> df = D.fromNamedColumns [("a", D.fromList [1..100]), ("b", D.fromList [1..100]), ("c", D.fromList [1..100])] ->>> D.nColumns df -3 -@ --} +-- | O(1) Get the number of columns in a dataframe. nColumns :: DataFrame -> Int nColumns = snd . dataframeDimensions -{- | O(k) Get column names of the DataFrame in order of insertion. - -==== __Example__ -@ ->>> :set -XOverloadedStrings ->>> import qualified DataFrame as D ->>> df = D.fromNamedColumns [("a", D.fromList [1..100]), ("b", D.fromList [1..100]), ("c", D.fromList [1..100])] ->>> D.columnNames df - -["a", "b", "c"] -@ --} - --- 'columnNames' is now defined in "DataFrame.Internal.DataFrame" and --- re-exported from "DataFrame" at the top level. - -{- | Adds a vector to the dataframe. If the vector has less elements than the dataframe and the dataframe is not empty -the vector is converted to type `Maybe a` filled with `Nothing` to match the size of the dataframe. Similarly, -if the vector has more elements than what's currently in the dataframe, the other columns in the dataframe are -change to `Maybe ` and filled with `Nothing`. - -==== __Example__ -@ ->>> :set -XOverloadedStrings ->>> import qualified DataFrame as D ->>> import qualified Data.Vector as V ->>> D.insertVector "numbers" (V.fromList [(1 :: Int)..10]) D.empty - --------- - numbers --------- - Int --------- - 1 - 2 - 3 - 4 - 5 - 6 - 7 - 8 - 9 - 10 +-- 'columnNames' lives in "DataFrame.Internal.DataFrame" (re-exported). -@ --} +-- | Add a vector as a named column; size mismatches pad with Nothing. insertVector :: forall a. (Columnable a) => @@ -160,38 +127,7 @@ insertVector name xs = insertColumn name (fromVector xs) {-# INLINE insertVector #-} -{- | Adds a foldable collection to the dataframe. If the collection has less elements than the -dataframe and the dataframe is not empty -the collection is converted to type `Maybe a` filled with `Nothing` to match the size of the dataframe. Similarly, -if the collection has more elements than what's currently in the dataframe, the other columns in the dataframe are -change to `Maybe ` and filled with `Nothing`. - -Be careful not to insert infinite collections with this function as that will crash the program. - -==== __Example__ -@ ->>> :set -XOverloadedStrings ->>> import qualified DataFrame as D ->>> D.insert "numbers" [(1 :: Int)..10] D.empty - --------- - numbers --------- - Int --------- - 1 - 2 - 3 - 4 - 5 - 6 - 7 - 8 - 9 - 10 - -@ --} +-- | Add a foldable as a named column; do not pass infinite collections. insert :: forall a t. (Columnable a, Foldable t) => @@ -202,37 +138,10 @@ -- | 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. - -==== __Example__ -@ ->>> :set -XOverloadedStrings ->>> import qualified Data.Vector as V ->>> import qualified DataFrame as D ->>> df = D.fromNamedColumns [("x", D.fromList [(1 :: Int)..10])] ->>> D.insertVectorWithDefault 0 "numbers" (V.fromList [(1 :: Int),2,3]) df - -------------- - x | numbers -----|-------- -Int | Int -----|-------- -1 | 1 -2 | 2 -3 | 3 -4 | 0 -5 | 0 -6 | 0 -7 | 0 -8 | 0 -9 | 0 -10 | 0 - -@ --} +-- | Add a vector, padding with a default if it has fewer elements than rows. insertVectorWithDefault :: forall a. (Columnable a) => @@ -250,33 +159,7 @@ values = xs V.++ V.replicate (rows - V.length xs) defaultValue in insertColumn name (fromVector values) d -{- | Adds a list to the dataframe and pads it with a default value if it has less elements than the number of rows. - -==== __Example__ -@ ->>> :set -XOverloadedStrings ->>> import qualified DataFrame as D ->>> df = D.fromNamedColumns [("x", D.fromList [(1 :: Int)..10])] ->>> D.insertWithDefault 0 "numbers" [(1 :: Int),2,3] df - -------------- - x | numbers -----|-------- -Int | Int -----|-------- -1 | 1 -2 | 2 -3 | 3 -4 | 0 -5 | 0 -6 | 0 -7 | 0 -8 | 0 -9 | 0 -10 | 0 - -@ --} +-- | Add a list, padding with a default if it has fewer elements than rows. insertWithDefault :: forall a t. (Columnable a, Foldable t) => @@ -295,11 +178,7 @@ 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. insertUnboxedVector :: forall a. (Columnable a, VU.Unbox a) => @@ -312,62 +191,9 @@ DataFrame insertUnboxedVector name xs = insertColumn name (UnboxedColumn Nothing xs) -{- | /O(n)/ Add a column to the dataframe. - -==== __Example__ -@ ->>> :set -XOverloadedStrings ->>> import qualified DataFrame as D ->>> D.insertColumn "numbers" (D.fromList [(1 :: Int)..10]) D.empty - --------- - numbers --------- - Int --------- - 1 - 2 - 3 - 4 - 5 - 6 - 7 - 8 - 9 - 10 - -@ --} - -- 'insertColumn' is now defined in "DataFrame.Internal.DataFrame". -{- | /O(n)/ Clones a column and places it under a new name in the dataframe. - -==== __Example__ -@ ->>> :set -XOverloadedStrings ->>> import qualified Data.Vector as V ->>> df = insertVector "numbers" (V.fromList [1..10]) D.empty ->>> D.cloneColumn "numbers" "others" df - ------------------ - numbers | others ----------|------- - Int | Int ----------|------- - 1 | 1 - 2 | 2 - 3 | 3 - 4 | 4 - 5 | 5 - 6 | 6 - 7 | 7 - 8 | 8 - 9 | 9 - 10 | 10 - -@ --} +-- | /O(n)/ Clone a column and place it under a new name. cloneColumn :: T.Text -> T.Text -> DataFrame -> DataFrame cloneColumn original new df | null df = throw (EmptyDataSetException "cloneColumn") @@ -379,83 +205,11 @@ column <- getColumn original df return $ insertColumn new column df -{- | /O(n)/ Renames a single column. - -==== __Example__ -@ ->>> :set -XOverloadedStrings ->>> import qualified DataFrame as D ->>> import qualified Data.Vector as V ->>> df = insertVector "numbers" (V.fromList [1..10]) D.empty ->>> D.rename "numbers" "others" df - -------- - others -------- - Int -------- - 1 - 2 - 3 - 4 - 5 - 6 - 7 - 8 - 9 - 10 - -@ --} +-- | /O(n)/ Rename a single column. rename :: T.Text -> T.Text -> DataFrame -> DataFrame rename orig new df = either throw id (renameSafe orig new df) -{- | /O(n)/ Renames many columns. - -==== __Example__ -@ ->>> :set -XOverloadedStrings ->>> import qualified DataFrame as D ->>> import qualified Data.Vector as V ->>> df = D.insertVector "others" (V.fromList [11..20]) (D.insertVector "numbers" (V.fromList [1..10]) D.empty) ->>> df - ------------------ - numbers | others ----------|------- - Int | Int ----------|------- - 1 | 11 - 2 | 12 - 3 | 13 - 4 | 14 - 5 | 15 - 6 | 16 - 7 | 17 - 8 | 18 - 9 | 19 - 10 | 20 - ->>> D.renameMany [("numbers", "first_10"), ("others", "next_10")] df - -------------------- - first_10 | next_10 -----------|-------- - Int | Int -----------|-------- - 1 | 11 - 2 | 12 - 3 | 13 - 4 | 14 - 5 | 15 - 6 | 16 - 7 | 17 - 8 | 18 - 9 | 19 - 10 | 20 - -@ --} +-- | /O(n)/ Rename many columns. renameMany :: [(T.Text, T.Text)] -> DataFrame -> DataFrame renameMany = fold (uncurry rename) @@ -478,24 +232,7 @@ , typeOfColumn :: !T.Text } -{- | O(n * k ^ 2) Returns the number of non-null columns in the dataframe and the type associated with each column. - -==== __Example__ -@ ->>> import qualified Data.Vector as V ->>> df = D.insertVector "others" (V.fromList [11..20]) (D.insertVector "numbers" (V.fromList [1..10]) D.empty) ->>> D.describeColumns df - --------------------------------------------------------- - Column Name | # Non-null Values | # Null Values | Type --------------|-------------------|---------------|----- - Text | Int | Int | Text --------------|-------------------|---------------|----- - others | 10 | 0 | Int - numbers | 10 | 0 | Int - -@ --} +-- | O(n * k ^ 2) Count non-null values and the type of each column. describeColumns :: DataFrame -> DataFrame describeColumns df = empty @@ -540,10 +277,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 @@ -557,100 +294,21 @@ nulls c@(PackedText _ _) = nulls (materializePacked c) nulls _ = 0 -{- | Creates a dataframe from a list of tuples with name and column. - -==== __Example__ -@ ->>> df = D.fromNamedColumns [("numbers", D.fromList [1..10]), ("others", D.fromList [11..20])] ->>> df ------------------ - numbers | others ----------|------- - Int | Int ----------|------- - 1 | 11 - 2 | 12 - 3 | 13 - 4 | 14 - 5 | 15 - 6 | 16 - 7 | 17 - 8 | 18 - 9 | 19 - 10 | 20 - -@ --} - --- 'fromNamedColumns' is now defined in "DataFrame.Internal.DataFrame". - -{- | Create a dataframe from a list of columns. The column names are "0", "1"... etc. -Useful for quick exploration but you should probably always rename the columns after -or drop the ones you don't want. - -==== __Example__ -@ ->>> df = D.fromUnnamedColumns [D.fromList [1..10], D.fromList [11..20]] ->>> df ------------------ - 0 | 1 ------|---- - Int | Int ------|---- - 1 | 11 - 2 | 12 - 3 | 13 - 4 | 14 - 5 | 15 - 6 | 16 - 7 | 17 - 8 | 18 - 9 | 19 - 10 | 20 +-- 'fromColumns' is now defined in "DataFrame.Internal.DataFrame". -@ --} +-- | Create a dataframe from a list of columns named "0", "1", and so on. fromUnnamedColumns :: [Column] -> DataFrame -fromUnnamedColumns = fromNamedColumns . zip (map (T.pack . show) [(0 :: Int) ..]) - -{- | Create a dataframe from a list of column names and rows. - -==== __Example__ -@ ->>> df = D.fromRows ["A", "B"] [[D.toAny 1, D.toAny 11], [D.toAny 2, D.toAny 12], [D.toAny 3, D.toAny 13]] - ->>> df - ----------- - A | B ------|---- - Int | Int ------|---- - 1 | 11 - 2 | 12 - 3 | 13 +fromUnnamedColumns = fromColumns . zip (map (T.pack . show) [(0 :: Int) ..]) -@ --} +-- | Create a dataframe from a list of column names and rows. 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] - -{- | O (k * n) Counts the occurences of each value in a given column. - -==== __Example__ -@ ->>> df = D.fromUnnamedColumns [D.fromList [1..10], D.fromList [11..20]] - ->>> D.valueCounts @Int "0" df - -[(1,1),(2,1),(3,1),(4,1),(5,1),(6,1),(7,1),(8,1),(9,1),(10,1)] + (zip [0 ..] names) -@ --} +-- | O (k * n) Count the occurrences of each value in a given column. valueCounts :: forall a. (Ord a, Columnable a) => Expr a -> DataFrame -> [(a, Int)] valueCounts expr df @@ -663,18 +321,7 @@ in M.toAscList column -{- | O (k * n) Shows the proportions of each value in a given column. - -==== __Example__ -@ ->>> df = D.fromUnnamedColumns [D.fromList [1..10], D.fromList [11..20]] - ->>> D.valueCounts @Int "0" df - -[(1,0.1),(2,0.1),(3,0.1),(4,0.1),(5,0.1),(6,0.1),(7,0.1),(8,0.1),(9,0.1),(10,0.1)] - -@ --} +-- | O (k * n) Show the proportion of each value in a given column. valueProportions :: forall a. (Ord a, Columnable a) => Expr a -> DataFrame -> [(a, Double)] valueProportions expr df @@ -690,56 +337,11 @@ in map (fmap ((/ total) . fromIntegral)) counts -{- | A left fold for dataframes that takes the dataframe as the last object. -This makes it easier to chain operations. - -==== __Example__ -@ ->>> df = D.fromNamedColumns [("x", D.fromList [1..100]), ("y", D.fromList [11..110])] ->>> D.fold D.dropLast [1..5] df - ---------- - x | y -----|---- -Int | Int -----|---- -1 | 11 -2 | 12 -3 | 13 -4 | 14 -5 | 15 -6 | 16 -7 | 17 -8 | 18 -9 | 19 -10 | 20 -11 | 21 -12 | 22 -13 | 23 -14 | 24 -15 | 25 -16 | 26 -17 | 27 -18 | 28 -19 | 29 -20 | 30 - -Showing 20 rows out of 85 - -@ --} +-- | A left fold for dataframes that takes the dataframe as the last argument. 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 float matrix, 'Left' if conversion fails. toFloatMatrix :: DataFrame -> Either DataFrameException (V.Vector (VU.Vector Float)) toFloatMatrix df = case V.foldl' @@ -753,15 +355,7 @@ (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 double matrix, 'Left' if conversion fails. toDoubleMatrix :: DataFrame -> Either DataFrameException (V.Vector (VU.Vector Double)) toDoubleMatrix df = case V.foldl' @@ -775,15 +369,7 @@ (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 int matrix, 'Left' if conversion fails. toIntMatrix :: DataFrame -> Either DataFrameException (V.Vector (VU.Vector Int)) toIntMatrix df = case V.foldl' (\acc c -> V.snoc <$> acc <*> toIntVector c) @@ -806,18 +392,7 @@ orFallback Nothing = Just (T.unpack name) withColumnName _ result = result -{- | Get a specific column as a vector. - -You must specify the type via type applications. - -==== __Examples__ - ->>> columnAsVector (F.col @Int "age") df -Right [25, 30, 35, ...] - ->>> columnAsVector (F.col @Text "name") df -Right ["Alice", "Bob", "Charlie", ...] --} +-- | Get a specific column as a vector; the type comes from type application. columnAsVector :: forall a. (Columnable a) => Expr a -> DataFrame -> Either DataFrameException (V.Vector a) @@ -833,11 +408,7 @@ 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', 'Left' if not convertible. columnAsIntVector :: (Columnable a, Num a) => Expr a -> DataFrame -> Either DataFrameException (VU.Vector Int) @@ -850,11 +421,7 @@ 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', 'Left' if not convertible. columnAsDoubleVector :: (Columnable a, Num a) => Expr a -> DataFrame -> Either DataFrameException (VU.Vector Double) @@ -870,11 +437,7 @@ 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', 'Left' if not convertible. columnAsFloatVector :: (Columnable a, Num a) => Expr a -> DataFrame -> Either DataFrameException (VU.Vector Float) @@ -910,30 +473,10 @@ #-} {-# INLINE columnAsUnboxedVector #-} -{- | Get a specific column as a list. - -You must specify the type via type applications. - -==== __Examples__ - ->>> columnAsList @Int "age" df -[25, 30, 35, ...] - ->>> columnAsList @Text "name" df -["Alice", "Bob", "Charlie", ...] - -==== __Throws__ - -* 'error' - if the column type doesn't match the requested type --} +-- | Get a specific column as a list; the type comes from type application. columnAsList :: forall a. (Columnable a) => Expr a -> DataFrame -> [a] columnAsList expr df = either throw V.toList (columnAsVector expr df) -{- | Returns the provenance of all columns in the DataFrame as a list of -@(name, expression)@ pairs. Derived columns show their expression; -raw columns show an identity @col \@type name@ expression. --} - -- TODO: mchavinda - Expand out these expressions if possible. showDerivedExpressions :: DataFrame -> [NamedExpr] showDerivedExpressions df = @@ -941,14 +484,21 @@ names = columnNames df toNamedExpr name = case M.lookup name exprs of Just uexpr -> (name, uexpr) - Nothing -> (name, identityUExpr name) + Nothing -> (name, identityExpr name) in map toNamedExpr names where - identityUExpr name = case getColumn name df of - Just (BoxedColumn (Just _) (_ :: V.Vector a)) -> UExpr (Col @(Maybe a) name) - Just (BoxedColumn Nothing (_ :: V.Vector a)) -> UExpr (Col @a name) - Just (UnboxedColumn (Just _) (_ :: VU.Vector a)) -> UExpr (Col @(Maybe a) name) - 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) + identityExpr name = case getColumn name df of + Just (BoxedColumn (Just _) (_ :: V.Vector a)) -> SomeExpr (Col @(Maybe a) name) + Just (BoxedColumn Nothing (_ :: V.Vector a)) -> SomeExpr (Col @a name) + Just (UnboxedColumn (Just _) (_ :: VU.Vector a)) -> SomeExpr (Col @(Maybe a) name) + Just (UnboxedColumn Nothing (_ :: VU.Vector a)) -> SomeExpr (Col @a name) + Just (PackedText (Just _) _) -> SomeExpr (Col @(Maybe T.Text) name) + Just (PackedText Nothing _) -> SomeExpr (Col @T.Text name) + Just c@(MergedColumn _ _) -> case mergedHead c of + BoxedColumn (Just _) (_ :: V.Vector a) -> SomeExpr (Col @(Maybe a) name) + BoxedColumn Nothing (_ :: V.Vector a) -> SomeExpr (Col @a name) + _ -> + error $ + "showDerivedExpressions: merged column did not materialize boxed: " + ++ T.unpack name Nothing -> error $ "showDerivedExpressions: column not found: " ++ T.unpack name diff --git a/src/DataFrame/Operations/Inference.hs b/src/DataFrame/Operations/Inference.hs --- a/src/DataFrame/Operations/Inference.hs +++ b/src/DataFrame/Operations/Inference.hs @@ -1,13 +1,6 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE ScopedTypeVariables #-} -{- | Single-pass sampled inference lattice (Round-2 S2). One walk over -the sampled cells maintains a candidate mask {Bool, Int, Double, Date} -cleared by attempting the WS-B byte parsers; the assumption is the -highest-priority surviving candidate. Shared by every reader (audit T1), -together with the Int -> Double prefix promotion that replaces the -full-column retry chain (audit T2). --} module DataFrame.Operations.Inference ( DateFormat, ParsingAssumption (..), @@ -51,35 +44,23 @@ parseTimeOpt :: DateFormat -> T.Text -> Maybe Day parseTimeOpt dateFormat s = - parseTimeM {- Accept leading/trailing whitespace -} + parseTimeM True defaultTimeLocale 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) -> @@ -159,9 +127,7 @@ (parseD . V.unsafeIndex vec) {-# INLINE promoteIntColumn #-} -{- | 'promoteIntColumn' over row indices instead of a materialized cell -vector (used by readers that keep cells as buffer offsets). --} +-- | 'promoteIntColumn' over row indices instead of a materialized cell vector. promoteIntColumnIndexed :: Int -> (Int -> Bool) -> @@ -188,8 +154,7 @@ Just dv -> promote i anyNull dv Nothing -> return Nothing - -- Switch to Double at row i: convert the Int prefix, then - -- continue parsing the remaining cells as Double. + -- Switch to Double at row i: convert the Int prefix, then keep going. promote !i !anyNull !dv = do dvalues <- VUM.unsafeNew n let copy !j diff --git a/src/DataFrame/Operations/Join.hs b/src/DataFrame/Operations/Join.hs --- a/src/DataFrame/Operations/Join.hs +++ b/src/DataFrame/Operations/Join.hs @@ -8,13 +8,37 @@ {-# 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 tests; not a stable public API. + 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 +49,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,66 +91,73 @@ 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 hash join switches strategy. 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. 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. , ciMask :: {-# UNPACK #-} !Int } -{- | Look up a hash in the open-addressing table. -Returns @(start, len)@ of the matching run, or @(-1, 0)@ on a miss. --} +-- | 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: top bits of a Fibonacci multiply, not low bits. +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, returning the run's (start, len) 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 #-} -{- | Smallest power of two strictly greater than @n@, at least 2. -Used to size the open-addressing table so the load factor stays below ~0.5. --} +-- | Smallest power of two strictly greater than @n@, at least 2. nextPow2Above :: Int -> Int nextPow2Above n = go 2 where @@ -121,45 +165,40 @@ | 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 row hashes by sorting and scanning for runs. 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,10 +208,7 @@ | 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. sortWithIndices :: VU.Vector Int -> (VU.Vector Int, VU.Vector Int) sortWithIndices hashes = runST $ do let n = VU.length hashes @@ -225,35 +261,13 @@ [] -> M.elems $ M.restrictKeys columnIdxs csSet _ -> throw (ColumnsNotFoundException missingKeys callPoint (M.keys columnIdxs)) --- ============================================================ --- Inner Join --- ============================================================ - -{- | Performs an inner join on two dataframes using the specified key columns. -Returns only rows where the key values exist in both dataframes. - -==== __Example__ -@ -ghci> df = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2", "K3"]), ("A", D.fromList ["A0", "A1", "A2", "A3"])] -ghci> other = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2"]), ("B", D.fromList ["B0", "B1", "B2"])] -ghci> D.innerJoin ["key"] df other - ------------------ - key | A | B -------|-----|---- - Text | Text| Text -------|-----|---- - K0 | A0 | B0 - K1 | A1 | B1 - K2 | A2 | B2 - -@ --} +-- | Inner join on two dataframes using the specified key columns. innerJoin :: [T.Text] -> DataFrame -> DataFrame -> DataFrame innerJoin cs left right | 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. innerJoinNonEmpty :: [T.Text] -> DataFrame -> DataFrame -> DataFrame innerJoinNonEmpty cs left right = let @@ -267,16 +281,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,40 +289,40 @@ | 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 assembleInner csSet left right leftIxs rightIxs -{- | Select the inner-join kernel: parallel chunked-probe hash join when the -caller decided the probe side is large enough and multi-core, otherwise the -sequential single-threaded hash kernel. --} +-- | Select the inner-join kernel: parallel probe when large, else sequential. innerKernel :: Bool -> VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) innerKernel True = parInnerKernel innerKernel False = hashInnerKernel {-# 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). --} +-- | Parallel inner-join kernel: build the index once, probe in parallel. 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' for the parallel probe kernels. +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,11 +330,7 @@ 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 a batch of rows against a pre-built 'CompactIndex'. hashProbeKernel :: -- | Built once from the full right\/build side. CompactIndex -> @@ -393,17 +394,7 @@ <*> 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). --} - -{- | 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, guarding cross-product explosion. maxJoinOutputRows :: Int maxJoinOutputRows = 500_000_000 @@ -472,16 +463,9 @@ (,) <$> 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 returning left and right expanded indices. sortMergeInnerKernel :: VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) sortMergeInnerKernel leftHashes rightHashes = @@ -565,11 +549,9 @@ (,) <$> 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. +-- | Assemble the inner join result DataFrame from expanded index vectors. assembleInner :: S.Set T.Text -> DataFrame -> @@ -582,7 +564,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 +575,6 @@ idx <- M.lookup name (D.columnIndices right) return (expandedRightCols `VB.unsafeIndex` idx) - -- Base DataFrame: all left columns, expanded baseDf = left { columns = expandedLeftCols @@ -607,46 +587,21 @@ 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 baseDf --- ============================================================ --- Left Join --- ============================================================ - -{- | Performs a left join on two dataframes using the specified key columns. -Returns all rows from the left dataframe, with matching rows from the right dataframe. -Non-matching rows will have Nothing/null values for columns from the right dataframe. - -==== __Example__ -@ -ghci> df = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2", "K3"]), ("A", D.fromList ["A0", "A1", "A2", "A3"])] -ghci> other = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2"]), ("B", D.fromList ["B0", "B1", "B2"])] -ghci> D.leftJoin ["key"] df other - ------------------------- - key | A | B -------|-----|---------- - Text | Text| Maybe Text -------|-----|---------- - K0 | A0 | Just "B0" - K1 | A1 | Just "B1" - K2 | A2 | Just "B2" - K3 | A3 | Nothing - -@ --} +-- | Left join on two dataframes, keeping every row from the left dataframe. leftJoin :: [T.Text] -> DataFrame -> DataFrame -> DataFrame leftJoin = leftJoinWithCallPoint "leftJoin" @@ -657,6 +612,7 @@ | D.null left || D.nRows left == 0 = D.empty | otherwise = leftJoinNonEmpty callPoint cs left right +-- The right side is the build side; unmatched right rows get a @-1@ sentinel. leftJoinNonEmpty :: T.Text -> [T.Text] -> DataFrame -> DataFrame -> DataFrame leftJoinNonEmpty callPoint cs left right = let @@ -669,10 +625,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,26 +636,17 @@ | 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, -then probe @leftHashes@ in parallel. Bit-for-bit identical output to -'hashLeftKernel' (left-row order preserved, @-1@ sentinel for misses). --} +-- | Parallel left-join kernel: build the index once, probe in parallel. parLeftKernel :: 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 with @-1@ marking unmatched right rows. hashLeftKernel :: VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) hashLeftKernel leftHashes rightHashes = runST $ do @@ -769,11 +712,7 @@ <$> 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 with a @-1@ sentinel for unmatched rows. sortMergeLeftKernel :: VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) sortMergeLeftKernel leftHashes rightHashes = runST $ do @@ -873,9 +812,7 @@ <$> VU.unsafeFreeze (VUM.slice 0 total lv) <*> VU.unsafeFreeze (VUM.slice 0 total rv) -{- | Assemble the result DataFrame for a left join. -Right index vectors use @-1@ sentinel, gathered via 'gatherWithSentinel'. --} +-- | Assemble the result DataFrame for a left join. assembleLeft :: S.Set T.Text -> DataFrame -> @@ -917,33 +854,14 @@ then insertIfPresent name - (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name) + (D.mkMergedColumns <$> getExpandedLeft name <*> getExpandedRight name) df else insertIfPresent name (getExpandedRight name) df ) rightColNames baseDf -{- | Performs a right join on two dataframes using the specified key columns. -Returns all rows from the right dataframe, with matching rows from the left dataframe. -Non-matching rows will have Nothing/null values for columns from the left dataframe. - -==== __Example__ -@ -ghci> df = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2", "K3"]), ("A", D.fromList ["A0", "A1", "A2", "A3"])] -ghci> other = D.fromNamedColumns [("key", D.fromList ["K0", "K1"]), ("B", D.fromList ["B0", "B1"])] -ghci> D.rightJoin ["key"] df other - ------------------ - key | A | B -------|-----|---- - Text | Text| Text -------|-----|---- - K0 | A0 | B0 - K1 | A1 | B1 - -@ --} +-- | Right join on two dataframes, keeping every row from the right dataframe. rightJoin :: [T.Text] -> DataFrame -> DataFrame -> DataFrame rightJoin cs left right = leftJoinWithCallPoint "rightJoin" cs right left @@ -967,20 +885,15 @@ 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. -Builds compact indices on both sides. -Returns @(leftExpandedIndices, rightExpandedIndices)@ with @-1@ sentinels. --} +-- | Hash-based full outer join kernel with @-1@ sentinels. hashFullOuterKernel :: VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) hashFullOuterKernel leftHashes rightHashes = runST $ do @@ -989,36 +902,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 +938,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 +975,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) @@ -1094,9 +1003,7 @@ (,) <$> VU.unsafeFreeze lv <*> VU.unsafeFreeze rv -{- | Sort-merge full outer join kernel. -Returns @(leftExpandedIndices, rightExpandedIndices)@ with @-1@ sentinels. --} +-- | Sort-merge full outer join kernel with @-1@ sentinels. sortMergeFullOuterKernel :: VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) sortMergeFullOuterKernel leftHashes rightHashes = runST $ do @@ -1105,7 +1012,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 +1027,6 @@ !rh = rightSH `VU.unsafeIndex` ri !totalRows = countLoop 0 0 0 - -- Pass 2: fill lv <- VUM.unsafeNew totalRows rv <- VUM.unsafeNew totalRows @@ -1164,10 +1069,7 @@ 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. assembleFullOuter :: S.Set T.Text -> DataFrame -> @@ -1201,8 +1103,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 +1152,15 @@ in D.fold ( \name df -> if S.member name csSet - then -- Key column: coalesce left and right - case (getExpandedLeft name, getExpandedRight name) of - (Just lc, Just rc) -> D.insertColumn name (coalesceKeyColumn lc rc) df - _ -> df + then case (getExpandedLeft name, getExpandedRight name) of + (Just lc, Just rc) -> D.insertColumn name (coalesceKeyColumn lc rc) df + _ -> df else if S.member name leftColSet then insertIfPresent name - (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name) + (D.mkMergedColumns <$> getExpandedLeft name <*> getExpandedRight name) df else insertIfPresent name (getExpandedRight name) df ) diff --git a/src/DataFrame/Operations/Join/Parallel.hs b/src/DataFrame/Operations/Join/Parallel.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Join/Parallel.hs @@ -0,0 +1,166 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Parallel chunked-probe join kernels over a shared 'CompactIndex'. +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 overhead outweighs parallel probing. +parJoinThreshold :: Int +parJoinThreshold = 200000 + +-- | Below this many build rows the sequential hash probe beats forking. +parBuildThreshold :: Int +parBuildThreshold = 500000 + +-- | Whether a join should take the parallel probe path. +shouldParallelizeJoin :: Int -> Int -> Bool +shouldParallelizeJoin probeRows buildRows = + probeRows >= parJoinThreshold + && buildRows >= parBuildThreshold + && capabilities > 1 +{-# NOINLINE shouldParallelizeJoin #-} + +-- | Probe-row count above which parallel probing wins even for small builds. +parProbeThreshold :: Int +parProbeThreshold = 1000000 + +-- | Whether a small-build join should take the parallel probe path. +shouldParallelizeSmallBuildProbe :: Int -> Bool +shouldParallelizeSmallBuildProbe probeRows = + probeRows >= parProbeThreshold + && capabilities > 1 +{-# NOINLINE shouldParallelizeSmallBuildProbe #-} + +-- | A read-only view of the build-side index fields needed by the probe. +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 returning (probeIxs, buildIxs). +parInnerProbe :: + ProbeTable -> + VU.Vector Int -> + IO (VU.Vector Int, VU.Vector Int) +parInnerProbe = runProbe False + +-- | Parallel left-join probe; unmatched rows carry a @-1@ build sentinel. +parLeftProbe :: + ProbeTable -> + VU.Vector Int -> + IO (VU.Vector Int, VU.Vector Int) +parLeftProbe = runProbe True + +-- | Shared two-pass parallel probe covering inner and left join semantics. +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. + 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 gives per-chunk 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 each chunk on its own forked thread; rethrow failures. +forkRanges :: Int -> (Int -> IO ()) -> IO () +forkRanges nChunks body = forkJoin_ (map body [0 .. nChunks - 1]) diff --git a/src/DataFrame/Operations/JoinPar.hs b/src/DataFrame/Operations/JoinPar.hs deleted file mode 100644 --- a/src/DataFrame/Operations/JoinPar.hs +++ /dev/null @@ -1,215 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ScopedTypeVariables #-} - -{- | -Parallel chunked-probe join kernels. The build side is indexed once into a -shared, read-only 'CompactIndex' (open-addressing, from -"DataFrame.Operations.Join"); the probe side is split into @caps@ /contiguous/ -row ranges and probed in parallel by 'forkIO' workers (no sparks). Each worker -makes two passes over its range — a count pass to size its slice, then a fill -pass — and writes into the single shared output buffers at a precomputed -prefix-sum offset. Because ranges are contiguous and laid out in range order, -the produced @(probeIxs, buildIxs)@ vectors are /bit-for-bit identical/ to the -sequential 'hashInnerKernel' \/ 'hashLeftKernel': probe rows appear in original -order and, within a probe row, build matches in @ciSortedIndices@ order. - -This is the parallel==sequential correctness gate (see -@tests/Operations/ParallelJoin.hs@). A sequential fallback is used when there is -a single capability or the probe side is below 'parJoinThreshold'; the caller -('innerJoin' \/ 'leftJoin') decides via 'shouldParallelizeJoin'. --} -module DataFrame.Operations.JoinPar ( - parInnerProbe, - parLeftProbe, - shouldParallelizeJoin, - shouldParallelizeSmallBuildProbe, - parJoinThreshold, - parBuildThreshold, - parProbeThreshold, -) where - -import Control.Concurrent (forkIO, getNumCapabilities) -import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar) -import Control.Exception (SomeException, throwIO, try) -import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Unboxed.Mutable as VUM -import System.IO.Unsafe (unsafePerformIO) - -{- | Below this many probe rows the fork/coordination overhead is not worth it; -the caller uses its sequential 'ST' kernel instead. --} -parJoinThreshold :: Int -parJoinThreshold = 200000 - -{- | Below this many build rows the shared 'CompactIndex' is small and hot, so -the sequential hash probe is already memory-bound-fast and the fork overhead -loses (measured: a 1e4-row Text-key build probed by 1e7 rows is /slower/ in -parallel). Parallelism only pays once the build index is large enough to spill -cache — exactly the regime where sort-merge used to be chosen. --} -parBuildThreshold :: Int -parBuildThreshold = 500000 - -{- | Whether a join should take the parallel probe path: more than one -capability, a probe side of at least 'parJoinThreshold' rows, and a build side -of at least 'parBuildThreshold' rows (a small/hot index is faster probed -sequentially). --} -shouldParallelizeJoin :: Int -> Int -> Bool -shouldParallelizeJoin probeRows buildRows = - probeRows >= parJoinThreshold - && buildRows >= parBuildThreshold - && capabilities > 1 -{-# NOINLINE shouldParallelizeJoin #-} - -{- | Above this many probe rows the probe-side row hashing and table lookups -dominate the join, so partitioning the probe across cores wins even when the -build side is small and cache-resident (the regime 'shouldParallelizeJoin' -deliberately leaves sequential). Sized at 1e6: below it the per-question gain is -swamped by 'forkIO'/coordination overhead, so small/medium-inner joins stay -sequential (measured). This is the small-build large-probe lever closing the -medium-factor 1e7 join (1e7 probe x ~1e4 build). --} -parProbeThreshold :: Int -parProbeThreshold = 1000000 - -{- | Whether a /small-build/ join (build below 'parBuildThreshold', so radix -partitioning / sort-merge is not used) should take the parallel probe path: a -very large probe side (at least 'parProbeThreshold') and more than one -capability. The shared build index is read-only across threads, so probing it in -parallel needs no synchronization. Independent of build size on purpose: the -build is already tiny; the cost is the 1e7-row probe hashing, which parallelizes -cleanly. --} -shouldParallelizeSmallBuildProbe :: Int -> Bool -shouldParallelizeSmallBuildProbe probeRows = - probeRows >= parProbeThreshold - && capabilities > 1 -{-# NOINLINE shouldParallelizeSmallBuildProbe #-} - -capabilities :: Int -capabilities = unsafePerformIO getNumCapabilities -{-# NOINLINE capabilities #-} - -{- | A read-only view of the build-side index needed by the probe: the lookup -returns @(start, len)@ of the matching run in @sortedIndices@, or @(-1, 0)@ on a -miss. Passed in by the caller so this module need not depend on the -'CompactIndex' record directly. --} -data ProbeIndex = ProbeIndex - { piSorted :: !(VU.Vector Int) - , piLookup :: !(Int -> (Int, Int)) - } - -{- | Parallel inner-join probe. @parInnerProbe sortedIdxs lookup probeHashes@ -returns @(probeIxs, buildIxs)@ identical to a sequential probe of the same -index. The build index must already be constructed from the build side. --} -parInnerProbe :: - VU.Vector Int -> - (Int -> (Int, Int)) -> - VU.Vector Int -> - IO (VU.Vector Int, VU.Vector Int) -parInnerProbe sortedIdxs lookupFn = - runProbe False (ProbeIndex sortedIdxs lookupFn) - -{- | Parallel left-join probe. Like 'parInnerProbe' but every probe row emits at -least one output row; unmatched rows carry a @-1@ sentinel in the build column. --} -parLeftProbe :: - VU.Vector Int -> - (Int -> (Int, Int)) -> - VU.Vector Int -> - IO (VU.Vector Int, VU.Vector Int) -parLeftProbe sortedIdxs lookupFn = - runProbe True (ProbeIndex sortedIdxs lookupFn) - -{- | Shared two-pass parallel probe. @keepUnmatched@ selects left- vs -inner-join semantics. Splits @[0, probeN)@ into @caps@ contiguous ranges, counts -each range's output, prefix-sums to global offsets, then fills the single output -buffers in parallel. --} -runProbe :: - Bool -> - ProbeIndex -> - VU.Vector Int -> - IO (VU.Vector Int, VU.Vector Int) -runProbe keepUnmatched pidx probeHashes = do - caps <- getNumCapabilities - let !probeN = VU.length probeHashes - !nChunks = max 1 (min caps probeN) - !sorted = piSorted pidx - !lookupFn = piLookup pidx - chunkBounds k = (lo, hi) - where - !lo = (probeN * k) `div` nChunks - !hi = (probeN * (k + 1)) `div` nChunks - -- Count pass: output rows produced by probe range [lo, hi). - countRange !lo !hi = - let go !i !acc - | i >= hi = acc - | otherwise = - let (!start, !len) = lookupFn (VU.unsafeIndex probeHashes i) - in if start < 0 - then go (i + 1) (if keepUnmatched then acc + 1 else acc) - else go (i + 1) (acc + len) - in go lo 0 - chunkCounts <- VUM.new (nChunks + 1) - forkRanges nChunks $ \k -> - let (lo, hi) = chunkBounds k - in VUM.unsafeWrite chunkCounts k (countRange lo hi) - -- Exclusive prefix sum -> per-chunk global start offsets; total at [nChunks]. - let scan !k !acc - | k > nChunks = pure acc - | otherwise = do - c <- if k < nChunks then VUM.unsafeRead chunkCounts k else pure 0 - VUM.unsafeWrite chunkCounts k acc - scan (k + 1) (acc + c) - !total <- scan 0 0 - pv <- VUM.unsafeNew (max 1 total) - bv <- VUM.unsafeNew (max 1 total) - -- Fill pass: each chunk writes from its prefix-sum offset. - offs <- VU.unsafeFreeze chunkCounts - forkRanges nChunks $ \k -> do - let (lo, hi) = chunkBounds k - !base = VU.unsafeIndex offs k - fill !i !p - | i >= hi = pure () - | otherwise = do - let (!start, !len) = lookupFn (VU.unsafeIndex probeHashes i) - if start < 0 - then - if keepUnmatched - then do - VUM.unsafeWrite pv p i - VUM.unsafeWrite bv p (-1) - fill (i + 1) (p + 1) - else fill (i + 1) p - else do - let writeMatch !j !q - | j >= len = pure () - | otherwise = do - VUM.unsafeWrite pv q i - VUM.unsafeWrite bv q (VU.unsafeIndex sorted (start + j)) - writeMatch (j + 1) (q + 1) - writeMatch 0 p - fill (i + 1) (p + len) - fill lo base - pf <- VU.unsafeFreeze (VUM.slice 0 total pv) - bf <- VU.unsafeFreeze (VUM.slice 0 total bv) - pure (pf, bf) - -{- | Run @body k@ for @k@ in @[0, nChunks)@, one chunk per task, on @nChunks@ -forked threads; rethrow the first failure. Chunk @k@ is owned by exactly one -thread, so concurrent writes to disjoint output regions are race-free. --} -forkRanges :: Int -> (Int -> IO ()) -> IO () -forkRanges nChunks body = do - vars <- mapM spawn [0 .. nChunks - 1] - results <- mapM takeMVar vars - mapM_ (either (throwIO :: SomeException -> IO ()) pure) results - where - spawn k = do - var <- newEmptyMVar - _ <- forkIO (try (body k) >>= putMVar var) - pure var diff --git a/src/DataFrame/Operations/Merge.hs b/src/DataFrame/Operations/Merge.hs --- a/src/DataFrame/Operations/Merge.hs +++ b/src/DataFrame/Operations/Merge.hs @@ -1,7 +1,11 @@ {-# LANGUAGE InstanceSigs #-} {-# OPTIONS_GHC -Wno-orphans #-} -module DataFrame.Operations.Merge where +module DataFrame.Operations.Merge ( + -- * Horizontal (side-by-side) merge + (|||), + -- Also exports the orphan @Semigroup@/@Monoid DataFrame@ instances. +) where import qualified Data.List as L import qualified Data.Text as T @@ -11,9 +15,7 @@ import Data.Maybe -{- | Vertically merge two dataframes using shared columns. -Columns that exist in only one dataframe are padded with Nothing. --} +-- | Vertically merge two dataframes, padding unshared columns with Nothing. instance Semigroup D.DataFrame where (<>) :: D.DataFrame -> D.DataFrame -> D.DataFrame (<>) a b = @@ -38,9 +40,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 +47,7 @@ Nothing -> D.insertColumn name (D.leftExpandColumn sumRows b'') df Just a'' -> - let concatedColumns = D.concatColumnsEither a'' b'' + let concatedColumns = D.mappendColumnsEither a'' b'' in D.insertColumn name concatedColumns df result = L.foldl' (addColumns a b) D.empty (D.columnNames a `L.union` D.columnNames b) in diff --git a/src/DataFrame/Operations/Permutation.hs b/src/DataFrame/Operations/Permutation.hs --- a/src/DataFrame/Operations/Permutation.hs +++ b/src/DataFrame/Operations/Permutation.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE InstanceSigs #-} @@ -5,31 +6,61 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -module DataFrame.Operations.Permutation where +module DataFrame.Operations.Permutation ( + -- * Sorting + SortOrder (..), + sortBy, + topBy, + -- * Shuffling + shuffle, + shuffledIndices, +) where + import qualified Data.List as L import qualified Data.Text as T import qualified Data.Vector as V import qualified Data.Vector.Algorithms.Merge as VA +import qualified Data.Vector.Algorithms.Radix as Radix import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM import Control.Exception (throw) +import Control.Monad (forM_, when) import Control.Monad.ST (runST) +import Data.Bits ( + complement, + countLeadingZeros, + testBit, + unsafeShiftR, + xor, + (.&.), + (.|.), + ) +import Data.Int (Int16, Int32, Int64, Int8) import Data.Type.Equality (testEquality, (:~:) (Refl)) +import Data.Typeable (Typeable) import Data.Vector.Internal.Check (HasCallStack) +import Data.Word (Word16, Word32, Word64, Word8) import DataFrame.Errors (DataFrameException (..)) -import DataFrame.Internal.Column (Column (..), Columnable, atIndicesStable) +import DataFrame.Internal.Column ( + Column (..), + Columnable, + atIndicesStable, + materializeMerged, + sliceColumn, + ) +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 GHC.Float (castDoubleToWord64) import System.Random (Random (randomR), RandomGen) import Type.Reflection (typeRep) @@ -52,10 +83,7 @@ mustFlipCompare (Asc _) = True mustFlipCompare (Desc _) = False -{- | Materialize any compound sort expressions into synthetic columns on -a working dataframe, returning rewritten 'SortOrder's that reference -those columns by name. --} +-- | Materialize compound sort expressions into synthetic columns. prepareSortColumns :: [SortOrder] -> DataFrame -> ([SortOrder], DataFrame) prepareSortColumns = go 0 where @@ -78,10 +106,7 @@ syntheticName :: Int -> T.Text syntheticName i = "__sortBy_synthetic_" <> T.pack (show i) <> "__" -{- | O(k log n) Sorts the dataframe by a given row. - -> sortBy Ascending ["Age"] df --} +-- | O(k log n) Sort the dataframe by the given sort orders. sortBy :: [SortOrder] -> DataFrame -> @@ -99,16 +124,88 @@ comparators = map (`sortOrderComparator` df') sortOrds' compositeCompare i j = mconcat [c i j | c <- comparators] nRows = fst (dataframeDimensions df') - indexes = sortIndices compositeCompare nRows + indexes = case sortOrds' of + [ord] | Just ixs <- radixSortIndices ord df' -> ixs + _ -> sortIndices compositeCompare nRows in df{columns = V.map (atIndicesStable indexes) (columns df)} where referenced = L.nub (concatMap sortOrderColumns sortOrds) missing = referenced L.\\ columnNames df -{- | Build a row-index comparator from a SortOrder and a DataFrame. -The Ord dictionary is recovered from the SortOrder GADT. --} +-- | The rows a 'sortBy' would place at positions @[offset, offset + count)@. +topBy :: Int -> Int -> [SortOrder] -> DataFrame -> DataFrame +topBy offset count sortOrds df + | not (null missing) = + throw $ + ColumnsNotFoundException + missing + "topBy" + (columnNames df) + | k >= nRows || k <= 0 = + let sorted = sortBy sortOrds df + in sorted + { columns = V.map (sliceColumn offset' (k - offset')) (columns sorted) + , dataframeDimensions = (k - offset', snd (dataframeDimensions df)) + } + | otherwise = + let + (sortOrds', df') = prepareSortColumns sortOrds df + comparators = map (`sortOrderComparator` df') sortOrds' + -- Ties keep their original order, as in the stable full sort. + cmp i j = mconcat [c i j | c <- comparators] <> compare i j + leading = smallestIndices cmp k nRows + window = VU.slice offset' (k - offset') (sortIndicesOf cmp leading) + in + df + { columns = V.map (atIndicesStable window) (columns df) + , dataframeDimensions = (VU.length window, snd (dataframeDimensions df)) + } + where + nRows = fst (dataframeDimensions df) + k = min nRows (max 0 offset + max 0 count) + offset' = min k (max 0 offset) + referenced = L.nub (concatMap sortOrderColumns sortOrds) + missing = referenced L.\\ columnNames df + +-- | The @k@ smallest of @[0, n)@ under @cmp@, unordered (a bounded max-heap). +smallestIndices :: (Int -> Int -> Ordering) -> Int -> Int -> VU.Vector Int +smallestIndices cmp k n = runST $ do + heap <- VUM.generate k id + let siftDown !i = do + let l = 2 * i + 1 + r = l + 1 + if l >= k + then pure () + else do + vi <- VUM.unsafeRead heap i + vl <- VUM.unsafeRead heap l + (c, vc) <- + if r < k + then do + vr <- VUM.unsafeRead heap r + pure (if cmp vr vl == GT then (r, vr) else (l, vl)) + else pure (l, vl) + when (cmp vc vi == GT) $ do + VUM.unsafeWrite heap i vc + VUM.unsafeWrite heap c vi + siftDown c + forM_ [k `div` 2 - 1, k `div` 2 - 2 .. 0] siftDown + forM_ [k .. n - 1] $ \i -> do + top <- VUM.unsafeRead heap 0 + when (cmp i top == LT) $ do + VUM.unsafeWrite heap 0 i + siftDown 0 + VU.unsafeFreeze heap + +-- | 'sortIndices' over an explicit index set. +sortIndicesOf :: (Int -> Int -> Ordering) -> VU.Vector Int -> VU.Vector Int +sortIndicesOf cmp ixs = runST $ do + mv <- VG.thaw (VU.convert ixs :: V.Vector Int) + VA.sortBy cmp mv + VU.convert <$> VG.unsafeFreeze mv + +-- | Build a row-index comparator from a SortOrder and a DataFrame. sortOrderComparator :: SortOrder -> DataFrame -> Int -> Int -> Ordering sortOrderComparator (Asc (Col name :: Expr a)) df = case unsafeGetColumn name df of @@ -124,6 +221,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,9 +240,71 @@ (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. + +-- | Stable LSD radix sort by one clean numeric column, else 'Nothing'. +radixSortIndices :: SortOrder -> DataFrame -> Maybe (VU.Vector Int) +radixSortIndices ord df = case ord of + Asc (Col name) -> keysOf name False + Desc (Col name) -> keysOf name True + _ -> Nothing + where + keysOf name desc = case unsafeGetColumn name df of + UnboxedColumn Nothing (v :: VU.Vector b) -> do + key <- orderKey @b + -- NaN has no 'compare' order; leave it to the comparison sort. + case testEquality (typeRep @b) (typeRep @Double) of + Just Refl | VU.any isNaN v -> Nothing + _ -> Just () + let flipIf k = if desc then complement k else k + Just (radixByKey (VU.map (flipIf . key) v)) + _ -> Nothing + +-- | An unsigned key whose order matches the type's 'compare' order. +orderKey :: forall b. (Typeable b) => Maybe (b -> Word64) +orderKey + | Just Refl <- testEquality (typeRep @b) (typeRep @Int) = Just signed + | Just Refl <- testEquality (typeRep @b) (typeRep @Int64) = Just signed + | Just Refl <- testEquality (typeRep @b) (typeRep @Int32) = Just signed + | Just Refl <- testEquality (typeRep @b) (typeRep @Int16) = Just signed + | Just Refl <- testEquality (typeRep @b) (typeRep @Int8) = Just signed + | Just Refl <- testEquality (typeRep @b) (typeRep @Word64) = Just fromIntegral + | Just Refl <- testEquality (typeRep @b) (typeRep @Word32) = Just fromIntegral + | Just Refl <- testEquality (typeRep @b) (typeRep @Word16) = Just fromIntegral + | Just Refl <- testEquality (typeRep @b) (typeRep @Word8) = Just fromIntegral + | Just Refl <- testEquality (typeRep @b) (typeRep @Double) = Just double + | otherwise = Nothing + where + signed :: (Integral a) => a -> Word64 + signed x = fromIntegral (fromIntegral x :: Int64) `xor` 0x8000000000000000 + -- -0.0 maps to 0.0's key, as 'compare' treats them equal. + double d = + let w = castDoubleToWord64 (if d == 0 then 0 else d) + in if testBit w 63 then complement w else w `xor` 0x8000000000000000 + +-- | Row indices ordered by key, ties in row order; skips constant bytes. +radixByKey :: VU.Vector Word64 -> VU.Vector Int +radixByKey keys + | VU.null keys = VU.empty + | otherwise = runST $ do + let !k0 = VU.unsafeHead keys + !varying = VU.foldl' (\acc k -> acc .|. (k `xor` k0)) 0 keys + !passes = (64 - countLeadingZeros varying + 7) `div` 8 + pairs <- VU.thaw (VU.imap (\i k -> (k, i)) keys) + Radix.sortBy + passes + 256 + (\p (k, _) -> fromIntegral ((k `unsafeShiftR` (8 * p)) .&. 255)) + pairs + VU.map snd <$> VU.unsafeFreeze pairs + sortIndices :: (Int -> Int -> Ordering) -> Int -> VU.Vector Int sortIndices cmp nRows = runST $ do withIndexes <- VG.thaw (V.generate nRows id :: V.Vector Int) @@ -168,14 +332,12 @@ shuffleVec :: (RandomGen g) => g -> VU.Vector Int shuffleVec g = runST $ do vm <- VUM.generate k id - let (n, nGen) = randomR (1, k - 1) g - go vm n nGen + go vm (k - 1) g VU.unsafeFreeze vm - go _v (-1) _ = pure () - go _v 0 _ = pure () - go v maxInd gen = + go _v i _ | i <= 0 = pure () + go v i gen = let - (n, nextGen) = randomR (1, maxInd) gen + (j, nextGen) = randomR (0, i) gen in - VUM.swap v 0 n *> go (VUM.tail v) (maxInd - 1) nextGen + VUM.swap v i j *> go v (i - 1) nextGen diff --git a/src/DataFrame/Operations/SetOps.hs b/src/DataFrame/Operations/SetOps.hs --- a/src/DataFrame/Operations/SetOps.hs +++ b/src/DataFrame/Operations/SetOps.hs @@ -1,19 +1,6 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ScopedTypeVariables #-} -{- | -Module : DataFrame.Operations.SetOps -Description : Set-theoretic ("topos") row operations. - -These treat a 'DataFrame' as a /set/ of rows and implement the subobject -lattice from relational algebra: 'union', 'intersect', 'difference', and -'symmetricDifference'. Every result is deduplicated, so each operation has the -schema-preserving shape @DataFrame -> DataFrame -> DataFrame@. - -Row equality is the same hash-based notion used by 'distinct' (see -"DataFrame.Operations.Aggregation"), so these operations and 'distinct' agree -on what "the same row" means. Both inputs are expected to share a schema; the -typed layer ('DataFrame.Typed') enforces that statically. --} module DataFrame.Operations.SetOps ( union, intersect, @@ -33,41 +20,23 @@ import DataFrame.Operations.Aggregation (distinct, groupBy, selectIndices) import DataFrame.Operations.Merge () -{- | All rows that appear in either dataframe, deduplicated. - -@union a b@ is @distinct (a <> b)@: the set union of the two row sets. --} +-- | All rows that appear in either dataframe, deduplicated. union :: DataFrame -> DataFrame -> DataFrame union a b = distinct (a <> b) -{- | Rows that appear in both dataframes, deduplicated. - -A row survives iff an equal row is present in each input. --} +-- | Rows that appear in both dataframes, deduplicated. intersect :: DataFrame -> DataFrame -> DataFrame intersect = setOp (&&) -{- | Rows present in the left dataframe but absent from the right, deduplicated -(the relational @EXCEPT@; the subobject complement of @a@ by @b@). --} +-- | Rows in the left dataframe but not the right, deduplicated. difference :: DataFrame -> DataFrame -> DataFrame difference = setOp (\inLeft inRight -> inLeft && not inRight) -{- | Rows present in exactly one of the two dataframes, deduplicated. - -@symmetricDifference a b@ is @union (difference a b) (difference b a)@. --} +-- | Rows present in exactly one of the two dataframes, deduplicated. symmetricDifference :: DataFrame -> DataFrame -> DataFrame symmetricDifference a b = difference a b `union` difference b a -{- | Core engine for 'intersect' and 'difference'. - -Concatenate the inputs, group by every column, and decide each group from -whether it has a member on the left side (original-row index @< nRows a@) and/or -the right side. The first member of a qualifying group is emitted as the -representative row, which keeps the result deduplicated and, because group -members come out in ascending row order, prefers the left dataframe's row. --} +-- | Core engine for 'intersect' and 'difference'. setOp :: (Bool -> Bool -> Bool) -> DataFrame -> DataFrame -> DataFrame setOp keep a b = selectIndices (VU.fromList chosen) combined @@ -79,8 +48,8 @@ chosen = [ VU.head members | k <- [0 .. nGroups - 1] - , let s = VU.unsafeIndex offs k - e = VU.unsafeIndex offs (k + 1) + , let !s = VU.unsafeIndex offs k + !e = VU.unsafeIndex offs (k + 1) members = VU.slice s (e - s) vis inLeft = VU.any (< leftRows) members inRight = VU.any (>= leftRows) members diff --git a/src/DataFrame/Operations/Statistics.hs b/src/DataFrame/Operations/Statistics.hs --- a/src/DataFrame/Operations/Statistics.hs +++ b/src/DataFrame/Operations/Statistics.hs @@ -9,8 +9,30 @@ {-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_GHC -Wno-orphans #-} -module DataFrame.Operations.Statistics where +module DataFrame.Operations.Statistics ( + -- * Summaries + summarize, + frequencies, + -- * Aggregate statistics + mean, + meanMaybe, + median, + medianMaybe, + percentile, + genericPercentile, + standardDeviation, + skewness, + variance, + interQuartileRange, + correlation, + sum, + + -- * Imputation + imputeWith, + -- Also exports the orphan @ImputeOp (Maybe b)@ instance. +) where + import qualified Data.List as L import qualified Data.Map as M import qualified Data.Text as T @@ -18,7 +40,8 @@ import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU -import Prelude as P +import Prelude hiding (sum) +import qualified Prelude as P import Control.Exception (throw) import Data.Function ((&)) @@ -30,39 +53,21 @@ DataFrame (..), columnNames, empty, - fromNamedColumns, + fromColumns, 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) import Text.Printf (printf) import Type.Reflection (typeRep) -{- | Show a frequency table for a categorical feaure. - -__Examples:__ - -@ -ghci> df <- D.readCsv ".\/data\/housing.csv" - -ghci> D.frequencies "ocean_proximity" df - ---------------------------------------------------------------------- - Statistic | <1H OCEAN | INLAND | ISLAND | NEAR BAY | NEAR OCEAN -----------------|-----------|--------|--------|----------|----------- - Text | Any | Any | Any | Any | Any -----------------|-----------|--------|--------|----------|----------- - Count | 9136 | 6551 | 5 | 2290 | 2658 - Percentage (%) | 44.26% | 31.74% | 0.02% | 11.09% | 12.88% -@ --} +-- | Show a frequency table for a categorical feature. frequencies :: forall a. (Columnable a, Ord a) => Expr a -> DataFrame -> DataFrame frequencies expr df = @@ -122,7 +127,7 @@ Left e -> throw e Right xs -> median' xs --- | Calculates the median of a given column (containing optional values) as a standalone value. +-- | The median of a column of optional values, as a standalone value. medianMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> DataFrame -> Double medianMaybe (Col name) df = @@ -196,7 +201,7 @@ Left e -> throw e Right xs -> variance' xs --- | Calculates the inter-quartile range of a given column as a standalone value. +-- | The inter-quartile range of a column as a standalone value. interQuartileRange :: forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double interQuartileRange (Col name) df = case columnAsUnboxedVector (Col @a name) df of @@ -208,7 +213,7 @@ Left e -> throw e Right xs -> interQuartileRange' xs --- | Calculates the Pearson's correlation coefficient between two given columns as a standalone value. +-- | Pearson's correlation coefficient of two columns, as a standalone value. correlation :: T.Text -> T.Text -> DataFrame -> Maybe Double correlation first second df = do f <- _getColumnAsDouble first df @@ -227,7 +232,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,49 +254,14 @@ Just Refl -> VG.sum column Nothing -> 0 Just (PackedText _ _) -> 0 + Just (MergedColumn _ _) -> 0 -- matches the old eager These column sum expr df = case interpret df expr of Left e -> throw e Right (TColumn xs) -> case toVector @a @V.Vector xs of Left e -> throw e Right xs' -> VG.sum xs' -{- | /O(n)/ Impute missing values in a column using a derived scalar. - -Given - -* an expression @f :: 'Expr' b -> 'Expr' b@ that, when interpreted over a - non-nullable column, produces the same value in every row (for example a - mean, median, or other aggregate), and -* a nullable column @'Expr' ('Maybe' b)@ - -this function: - -1. Drops all @Nothing@ values from the target column. -2. Interprets @f@ on the remaining non-null values. -3. Checks that the resulting column contains a single repeated value. -4. Uses that value to impute all @Nothing@s in the original column. - -==== __Throws__ - -* 'DataFrameException' - if the column does not exist, is empty, - -==== __Example__ -@ ->>> :set -XOverloadedStrings ->>> import qualified DataFrame as D ->>> let df = -... D.fromNamedColumns -... [ ("age", D.fromList [Just 10, Nothing, Just 20 :: Maybe Int]) ] ->>> ->>> -- Impute missing ages with the mean of the observed ages ->>> D.imputeWith F.mean "age" df --- age --- ---- --- 10 --- 15 --- 20 -@ --} +-- | /O(n)/ Impute missing values in a column using a derived scalar. instance {-# OVERLAPPING #-} (Columnable b) => ImputeOp (Maybe b) where runImpute = imputeCore @@ -304,7 +274,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. @@ -339,7 +309,7 @@ fold columnStats (columnNames df) - ( fromNamedColumns + ( fromColumns [ ( "Statistic" , fromList diff --git a/src/DataFrame/Operations/Subset.hs b/src/DataFrame/Operations/Subset.hs --- a/src/DataFrame/Operations/Subset.hs +++ b/src/DataFrame/Operations/Subset.hs @@ -5,12 +5,53 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -module DataFrame.Operations.Subset where +module DataFrame.Operations.Subset ( + -- * Row slicing + take, + takeLast, + drop, + dropLast, + range, + cube, + selectRows, + -- * Filtering + filter, + filterBy, + filterWhere, + filterJust, + filterNothing, + filterAllJust, + filterAllNothing, + + -- * Column selection + select, + selectBy, + exclude, + SelectionCriteria, + byName, + byProperty, + byNameProperty, + byNameRange, + byIndexRange, + + -- * Sampling & splitting + SplittableGen, + sample, + randomSplit, + kFolds, + stratifiedSample, + stratifiedSplit, + + -- * Label helpers (exported for satellite packages) + rowsAtIndices, +) where + import qualified Data.List as L import qualified Data.Map as M import qualified Data.Set as S @@ -33,26 +74,55 @@ DataFrameException (..), TypeErrorContext (..), ) -import DataFrame.Internal.Column +import DataFrame.Expression.Operators ( + col, + name, + (.<.), + (.<=.), + (.>.), + (.>=.), + ) +import DataFrame.Internal.Column ( + Column (..), + Columnable, + TypedColumn (TColumn), + atIndicesStable, + columnToTextVec, + findIndices, + hasMissing, + maskIndices, + 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,20 +177,20 @@ 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 first: an unclamped very negative end wraps. + end' = clip end start' r + n' = end' - start' clip :: Int -> Int -> Int -> Int clip n left right = min right $ max n left -{- | O(n * k) Filter rows by a given condition. - -> filter "x" even df --} +-- | O(n * k) Filter rows by a given condition. filter :: forall a. (Columnable a) => @@ -137,8 +207,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 @@ -202,26 +273,23 @@ , dataframeDimensions = (VG.length ixs, snd (dataframeDimensions df)) } -{- | O(k) a version of filter where the predicate comes first. - -> filterBy even "x" df --} +-- | O(k) A version of filter where the predicate comes first. filterBy :: (Columnable a) => (a -> Bool) -> Expr a -> DataFrame -> DataFrame filterBy = flip filter -{- | O(k) filters the dataframe with a boolean expression. - -> filterWhere (F.col @Int x + F.col y F.> 5) df --} +-- | O(k) Filter the dataframe with a boolean expression. filterWhere :: Expr Bool -> DataFrame -> DataFrame filterWhere expr df = let (TColumn col') = case interpret @Bool df (normalize expr) of Left e -> throw e Right c -> c - indexes = case findIndices id col' of - Right ixs -> ixs - Left e -> throw e + indexes = case col' of + UnboxedColumn Nothing (mask :: VU.Vector b) + | Just Refl <- testEquality (typeRep @b) (typeRep @Bool) -> maskIndices mask + _ -> case findIndices id col' of + Right ixs -> ixs + Left e -> throw e c' = snd $ dataframeDimensions df in df @@ -229,10 +297,7 @@ , dataframeDimensions = (VU.length indexes, c') } -{- | O(k) removes all rows with `Nothing` in a given column from the dataframe. - -> filterJust "col" df --} +-- | O(k) Remove all rows with Nothing in a given column. filterJust :: T.Text -> DataFrame -> DataFrame filterJust colName df = case getColumn colName df of Nothing -> @@ -248,10 +313,7 @@ _ -> df Just _ -> df -{- | O(k) returns all rows with `Nothing` in a give column. - -> filterNothing "col" df --} +-- | O(k) Return all rows with Nothing in a given column. filterNothing :: T.Text -> DataFrame -> DataFrame filterNothing colName df = case getColumn colName df of Nothing -> @@ -265,32 +327,19 @@ _ -> df _ -> df -{- | O(n * k) removes all rows with `Nothing` from the dataframe. - -> filterAllJust df --} +-- | O(n * k) Remove all rows with Nothing from the dataframe. filterAllJust :: DataFrame -> DataFrame filterAllJust df = foldr filterJust df (columnNames df) -{- | O(n * k) keeps any row with a null value. - -> filterAllNothing df --} +-- | O(n * k) Keep any row with a null value. filterAllNothing :: DataFrame -> DataFrame filterAllNothing df = foldr filterNothing df (columnNames df) -{- | O(k) cuts the dataframe in a cube of size (a, b) where - a is the length and b is the width. - -> cube (10, 5) df --} +-- | O(k) Cut the dataframe to a cube of size (length, width). cube :: (Int, Int) -> DataFrame -> DataFrame cube (len, width) = take len . selectBy [ColumnIndexRange (0, width - 1)] -{- | O(n) Selects a number of columns in a given dataframe. - -> select ["name", "age"] df --} +-- | O(n) Select a number of columns in a given dataframe. select :: [T.Text] -> DataFrame -> @@ -319,42 +368,23 @@ | ColumnIndexRange (Int, Int) | ColumnName T.Text -{- | Criteria for selecting a column by name. - -> selectBy [byName "Age"] df - -equivalent to: - -> select ["Age"] df --} +-- | Criteria for selecting a column by name. byName :: T.Text -> SelectionCriteria byName = ColumnName -{- | Criteria for selecting columns whose property satisfies given predicate. - -> selectBy [byProperty isNumeric] df --} +-- | Criteria for selecting columns whose property satisfies a predicate. byProperty :: (Column -> Bool) -> SelectionCriteria byProperty = ColumnProperty -{- | Criteria for selecting columns whose name satisfies given predicate. - -> selectBy [byNameProperty (T.isPrefixOf "weight")] df --} +-- | Criteria for selecting columns whose name satisfies a predicate. byNameProperty :: (T.Text -> Bool) -> SelectionCriteria byNameProperty = ColumnNameProperty -{- | Criteria for selecting columns whose names are in the given lexicographic range (inclusive). - -> selectBy [byNameRange ("a", "c")] df --} +-- | Criteria for selecting columns whose names are in a lexicographic range. byNameRange :: (T.Text, T.Text) -> SelectionCriteria byNameRange = ColumnTextRange -{- | Criteria for selecting columns whose indices are in the given (inclusive) range. - -> selectBy [byIndexRange (0, 5)] df --} +-- | Criteria for selecting columns whose indices are in an inclusive range. byIndexRange :: (Int, Int) -> SelectionCriteria byIndexRange = ColumnIndexRange @@ -377,10 +407,7 @@ where ixs = V.ifoldl' (\acc' i c -> if f c then i : acc' else acc') [] (columns df) -{- | O(k * n) select rows by index - -> selectRows [0, 2, 4] df --} +-- | O(k * n) Select rows by index. selectRows :: [Int] -> DataFrame -> DataFrame selectRows ixs df = df @@ -390,10 +417,7 @@ where ixs' = VU.fromList ixs -{- | O(n) inverse of select - -> exclude ["Name"] df --} +-- | O(n) Inverse of select. exclude :: [T.Text] -> DataFrame -> @@ -402,15 +426,7 @@ let keysToKeep = columnNames df L.\\ cs in select keysToKeep df -{- | Sample a dataframe. The double parameter must be between 0 and 1 (inclusive). - -==== __Example__ -@ -ghci> import System.Random -ghci> D.sample (mkStdGen 137) 0.1 df - -@ --} +-- | Sample a dataframe; the fraction must be between 0 and 1 inclusive. sample :: (RandomGen g) => g -> Double -> DataFrame -> DataFrame sample pureGen p df = let @@ -422,15 +438,7 @@ & filterWhere (cRand .>=. Lit (1 - p)) & exclude [name cRand] -{- | Split a dataset into two. The first in the tuple gets a sample of p (0 <= p <= 1) and the second gets (1 - p). This is useful for creating test and train splits. - -==== __Example__ -@ -ghci> import System.Random -ghci> D.randomSplit (mkStdGen 137) 0.9 df - -@ --} +-- | Split a dataset in two: a sample of fraction p and the remainder. randomSplit :: (RandomGen g) => g -> Double -> DataFrame -> (DataFrame, DataFrame) randomSplit pureGen p df = @@ -448,15 +456,7 @@ & exclude [name cRand] ) -{- | Creates n folds of a dataframe. - -==== __Example__ -@ -ghci> import System.Random -ghci> D.kFolds (mkStdGen 137) 5 df - -@ --} +-- | Create n folds of a dataframe. kFolds :: (RandomGen g) => g -> Int -> DataFrame -> [DataFrame] kFolds pureGen folds df = let @@ -476,26 +476,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' = @@ -515,14 +495,7 @@ , dataframeDimensions = (VU.length ixs, snd (dataframeDimensions df)) } -{- | Sample a dataframe, preserving per-stratum proportions. - -==== __Example__ -@ -ghci> import System.Random -ghci> D.stratifiedSample (mkStdGen 42) 0.8 "label" df -@ --} +-- | Sample a dataframe, preserving per-stratum proportions. stratifiedSample :: forall a g. (SplittableGen g, Columnable a) => @@ -530,7 +503,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) = @@ -539,14 +514,7 @@ in sample g1 p stratum <> go g2 rest in go gen groups -{- | Split a dataframe into two, preserving per-stratum proportions. - -==== __Example__ -@ -ghci> import System.Random -ghci> D.stratifiedSplit (mkStdGen 42) 0.8 "label" df -@ --} +-- | Split a dataframe into two, preserving per-stratum proportions. stratifiedSplit :: forall a g. (SplittableGen g, Columnable a) => @@ -554,7 +522,9 @@ stratifiedSplit gen p strataCol df = let col' = case strataCol of Col colName -> unsafeGetColumn colName df - _ -> unwrapTypedColumn (either throw id (interpret @a df strataCol)) + _ -> either throw id $ do + v <- eval (FlatCtx df) strataCol + pure $ materialize @a (fst (dataframeDimensions df)) v groups = M.elems (groupByIndices col') go _ [] = (mempty, mempty) go g (ixs : rest) = diff --git a/src/DataFrame/Operations/Transformations.hs b/src/DataFrame/Operations/Transformations.hs --- a/src/DataFrame/Operations/Transformations.hs +++ b/src/DataFrame/Operations/Transformations.hs @@ -1,5 +1,5 @@ {-# LANGUAGE ConstrainedClassMethods #-} -{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} @@ -7,11 +7,31 @@ {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE UndecidableSuperClasses #-} -module DataFrame.Operations.Transformations where +module DataFrame.Operations.Transformations ( + -- * Apply + apply, + safeApply, + applyMany, + applyInt, + applyDouble, + applyWhere, + applyAtIndex, + -- * Derive + derive, + deriveWithExpr, + deriveMany, + + -- * Impute + ImputeOp (..), + impute, + imputeCore, +) where + import qualified Data.List as L import qualified Data.Map as M import qualified Data.Text as T @@ -21,8 +41,10 @@ import Data.Maybe import DataFrame.Errors (DataFrameException (..), TypeErrorContext (..)) import DataFrame.Internal.Column ( + Column, Columnable, TypedColumn (..), + columnTypeString, hasMissing, ifoldrColumn, imapColumn, @@ -30,9 +52,11 @@ ) import DataFrame.Internal.DataFrame (DataFrame (..), getColumn, insertColumn) import DataFrame.Internal.Expression +import DataFrame.Internal.Expression.Operators.Nullable (BaseType) import DataFrame.Internal.Interpreter -import DataFrame.Internal.Nullable (BaseType) import DataFrame.Operations.Core +import GHC.TypeLits (ErrorMessage (..), TypeError) +import Type.Reflection (typeRep) -- | O(k) Apply a function to a given column in a dataframe. apply :: @@ -51,7 +75,7 @@ Left exception -> throw exception Right df -> df --- | O(k) Safe version of the apply function. Returns (instead of throwing) the error. +-- | O(k) Safe version of 'apply'. Returns the error instead of throwing. safeApply :: forall b c. (Columnable b, Columnable c) => @@ -69,25 +93,16 @@ column' <- mapColumn f column pure $ insertColumn columnName column' d -{- | O(k) Apply a function to an expression in a dataframe and -add the result into `alias` column. --} +-- | O(k) Derive a new column named alias from an expression. derive :: forall a. (Columnable a) => T.Text -> Expr a -> DataFrame -> DataFrame derive name expr df = case interpret @a df (normalize expr) of Left e -> throw e Right (TColumn value) -> (insertColumn name value df) - { derivingExpressions = M.insert name (UExpr expr) (derivingExpressions df) + { derivingExpressions = M.insert name (SomeExpr expr) (derivingExpressions df) } -{- | O(k) Apply a function to an expression in a dataframe and -add the result into `alias` column but - -==== __Examples__ - ->>> (z, df') = deriveWithExpr "z" (F.col @Int "x" + F.col "y") df ->>> filterWhere (z .>= 50) --} +-- | O(k) Like 'derive' but also returns a reference to the new column. deriveWithExpr :: forall a. (Columnable a) => T.Text -> Expr a -> DataFrame -> (Expr a, DataFrame) deriveWithExpr name expr df = case interpret @a df (normalize expr) of @@ -95,14 +110,14 @@ Right (TColumn value) -> ( Col name , (insertColumn name value df) - { derivingExpressions = M.insert name (UExpr expr) (derivingExpressions df) + { derivingExpressions = M.insert name (SomeExpr expr) (derivingExpressions df) } ) deriveMany :: [NamedExpr] -> DataFrame -> DataFrame deriveMany exprs df = let - f (name, UExpr (expr :: Expr a)) d = + f (name, SomeExpr (expr :: Expr a)) d = case interpret @a df expr of Left e -> throw e Right (TColumn value) -> insertColumn name value d @@ -142,11 +157,7 @@ DataFrame applyDouble = apply -{- | O(k * n) Apply a function to a column only if there is another column -value that matches the given criterion. - -> applyWhere (<20) "Age" (const "Gen-Z") "Generation" df --} +-- | O(k * n) Apply a function to a column where another column matches. applyWhere :: forall a b. (Columnable a, Columnable b) => @@ -199,7 +210,7 @@ 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 for nullable columns; a column with no nulls is an error. imputeCore :: forall b. (Columnable b) => @@ -215,9 +226,22 @@ 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 with no null bitmap, so nothing to 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,13 +251,23 @@ 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. --} +-- | Replace all instances of Nothing in a column with the given value. impute :: forall a. (ImputeOp a) => diff --git a/src/DataFrame/Operations/Typing.hs b/src/DataFrame/Operations/Typing.hs --- a/src/DataFrame/Operations/Typing.hs +++ b/src/DataFrame/Operations/Typing.hs @@ -6,9 +6,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} --- The inference lattice (sample classification, 'ParsingAssumption', --- Int -> Double promotion) lives in "DataFrame.Operations.Inference" --- and is re-exported here for backwards compatibility. +-- The inference lattice lives in "DataFrame.Operations.Inference", re-exported. module DataFrame.Operations.Typing ( module DataFrame.Operations.Typing, module DataFrame.Operations.Inference, @@ -29,34 +27,27 @@ 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: throw, 'Nothing', or 'Left' raw text. data SafeReadMode = NoSafeRead | MaybeRead @@ -68,26 +59,16 @@ { missingValues :: [T.Text] -- ^ Values to treat as @Nothing@ when the effective mode is 'MaybeRead'. , sampleSize :: Int - -- ^ Number of rows to inspect when inferring a column's type (0 = all rows). + -- ^ 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 with no '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'. , parseDateFormat :: DateFormat - -- ^ Date format string as accepted by "Data.Time.Format" (e.g. @\"%Y-%m-%d\"@). + -- ^ Date format as accepted by "Data.Time.Format" (e.g. @\"%Y-%m-%d\"@). } -{- | Sensible out-of-the-box parse options: infer from the first 100 rows, - treat common nullish strings as missing, and expect ISO 8601 dates. --} +-- | Sensible out-of-the-box parse options: infer from 100 rows, ISO dates. defaultParseOptions :: ParseOptions defaultParseOptions = ParseOptions @@ -98,9 +79,7 @@ , parseDateFormat = "%Y-%m-%d" } -{- | Resolve a column's effective 'SafeReadMode': the override if present, -otherwise the default. --} +-- | Resolve a column's effective 'SafeReadMode': override if present. effectiveSafeRead :: SafeReadMode -> [(T.Text, SafeReadMode)] -> T.Text -> SafeReadMode effectiveSafeRead def overrides name = fromMaybe def (lookup name overrides) @@ -108,7 +87,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 +122,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,11 +139,7 @@ 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 into an @Either Text a@ column. handleEitherAssumption :: DateFormat -> ParsingAssumption -> V.Vector T.Text -> Column handleEitherAssumption dfmt assumption raw = case assumption of @@ -177,9 +147,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,11 +201,7 @@ (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 to Double promotion. handleIntAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column handleIntAssumption isNull cols = case promoteIntColumn (\_ t -> isNull t) readIntStrict readDouble cols of @@ -251,10 +214,7 @@ (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. handleTextAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column handleTextAssumption isNull cols | V.any isNull cols = @@ -262,19 +222,12 @@ (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 columns: one boxed parse pass, falling back to Text on failure. 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 +257,10 @@ Nothing -> return Nothing loop 0 False +-- Reached only for an all-null sample: try each concrete type, else Text. 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,20 +271,11 @@ 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. isNullishOrMissing :: [T.Text] -> T.Text -> Bool isNullishOrMissing missing v = isNullish v || v `elem` missing @@ -344,7 +287,7 @@ unsafeParseTime :: DateFormat -> T.Text -> Day unsafeParseTime dateFormat s = - parseTimeOrError {- Accept leading/trailing whitespace -} + parseTimeOrError True defaultTimeLocale dateFormat @@ -361,11 +304,7 @@ 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. parseWithTypes :: (T.Text -> SafeReadMode) -> M.Map T.Text SchemaType -> @@ -379,9 +318,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,9 +328,10 @@ 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 st@(SType (_ :: P.Proxy a)) c@(PackedText _ _) = + case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> c + Nothing -> 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 Just HRefl -> case testEquality (typeRep @a) (typeRep @b) of @@ -434,9 +371,7 @@ where s = if null v then "\"\"" else v -{- | Try 'readMaybe'; on failure return @Left raw@ where @raw@ is the original -input text. Used by 'parseWithTypes' under 'EitherRead'. --} +-- | Try 'readMaybe'; on failure return @Left raw@ with the original text. readEitherRaw :: forall a. (Read a) => String -> Either T.Text a readEitherRaw s = case readMaybe s of Just v -> Right v diff --git a/src/DataFrame/Typed/Access.hs b/src/DataFrame/Typed/Access.hs --- a/src/DataFrame/Typed/Access.hs +++ b/src/DataFrame/Typed/Access.hs @@ -10,23 +10,38 @@ -- * 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 -the schema. The column must exist (enforced at compile time). --} +-- | Retrieve a column as a boxed 'Vector', typed by the schema. columnAsVector :: forall name cols a. ( KnownSymbol name @@ -53,3 +68,86 @@ D.columnAsList (Col @a colName) df where colName = T.pack (symbolVal (Proxy @name)) + +-- | Retrieve a column as an unboxed 'Int' vector, checked at compile time. +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 as 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 as 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. +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' as a row-major matrix; total. +toDoubleMatrix :: + (AllColumnsReal "toDoubleMatrix" cols) => + TypedDataFrame cols -> V.Vector (VU.Vector Double) +toDoubleMatrix (TDF df) = either throw id (D.toDoubleMatrix df) + +-- | Every column as 'Float', row-major. See 'toDoubleMatrix'. +toFloatMatrix :: + (AllColumnsReal "toFloatMatrix" cols) => + TypedDataFrame cols -> V.Vector (VU.Vector Float) +toFloatMatrix (TDF df) = either throw id (D.toFloatMatrix df) + +-- | Every column as 'Int', row-major. See 'toDoubleMatrix'. +toIntMatrix :: + (AllColumnsReal "toIntMatrix" cols) => + TypedDataFrame cols -> V.Vector (VU.Vector Int) +toIntMatrix (TDF df) = either throw id (D.toIntMatrix df) diff --git a/src/DataFrame/Typed/Aggregate.hs b/src/DataFrame/Typed/Aggregate.hs --- a/src/DataFrame/Typed/Aggregate.hs +++ b/src/DataFrame/Typed/Aggregate.hs @@ -34,77 +34,23 @@ import DataFrame.Typed.Schema import DataFrame.Typed.Types -{- | Group a typed DataFrame by one or more key columns. - -@ -grouped = groupBy \@'[\"department\"] employees -@ --} +-- | Group a typed DataFrame by one or more key columns. groupBy :: forall (keys :: [Symbol]) cols. (AllKnownSymbol keys, AssertAllPresent keys cols) => TypedDataFrame cols -> TypedGrouped keys cols groupBy (TDF df) = TGD (DA.groupBy (symbolVals @keys) df) -{- | Build a named aggregation entry. The result column name is supplied via -@TypeApplications@; the underlying expression is validated against the -source schema at compile time. - -@as@ produces a /transformer/ on the aggregation chain — entries compose -with plain @(.)@ from Prelude (or via @(|>)@ for SQL-like postfix -reading). 'aggregate' applies the composed transformer to the empty chain -internally, so no terminator is needed. - -==== __Prefix form__ - -@ -result = grouped |> aggregate - ( as \@\"total\" (sum (col \@\"amount\")) - . as \@\"orders\" (count (col \@\"order_id\")) - . as \@\"avg\" (mean (col \@\"amount\")) - ) -@ - -==== __Postfix form (SQL-like)__ - -@ -result = grouped |> aggregate - ( (sum (col \@\"amount\") |> as \@\"total\") - . (count (col \@\"order_id\") |> as \@\"orders\") - . (mean (col \@\"amount\") |> as \@\"avg\") - ) -@ - -Per-entry parentheses are required in the postfix form because -@(.)@ binds tighter than @(|>)@. --} +-- | Build a named aggregation entry, validated against the schema. as :: forall name a keys cols aggs. (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. - -The first argument is a chain of 'as' entries composed with @(.)@. The -empty composition (@id@) yields just the group keys. The result schema is -the group-key columns followed by the aggregation columns in declaration -order. - -@ -result = grouped |> aggregate - ( as \@\"total\" (sum (col \@\"amount\")) - . as \@\"orders\" (count (col \@\"order_id\")) - ) --- result :: TypedDataFrame --- '[ Column \"region\" Text --- , Column \"total\" Double --- , Column \"orders\" Int --- ] -@ --} +-- | Run a typed aggregation against a grouped DataFrame. aggregate :: forall keys cols aggs. (TAgg keys cols '[] -> TAgg keys cols aggs) -> diff --git a/src/DataFrame/Typed/Apply.hs b/src/DataFrame/Typed/Apply.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Apply.hs @@ -0,0 +1,192 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} + +-- | Typed column transformations that track schema changes in the types. +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@. +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. +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 on rows passing a condition. +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. +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 with disjoint schemas side by side. +(|||) :: + (AssertDisjoint left right) => + TypedDataFrame left -> + TypedDataFrame right -> + TypedDataFrame (Append left right) +(TDF a) ||| (TDF b) = unsafeFreeze (a D.||| b) diff --git a/src/DataFrame/Typed/Expr.hs b/src/DataFrame/Typed/Expr.hs --- a/src/DataFrame/Typed/Expr.hs +++ b/src/DataFrame/Typed/Expr.hs @@ -14,32 +14,7 @@ {-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_GHC -Wno-orphans #-} -{- | Type-safe expression construction for typed DataFrames. - -Unlike the untyped @Expr a@ where column references are unchecked strings, -'TExpr' ensures at compile time that: - -* Referenced columns exist in the schema -* Column types match the expression type - -== Example - -@ -type Schema = '[Column \"age\" Int, Column \"salary\" Double] - --- This compiles: -goodExpr :: TExpr Schema Double -goodExpr = col \@\"salary\" - --- This gives a compile-time error (column not found): -badExpr :: TExpr Schema Double -badExpr = col \@\"nonexistent\" - --- This gives a compile-time error (type mismatch): -wrongType :: TExpr Schema Int -wrongType = col \@\"salary\" -- salary is Double, not Int -@ --} +-- | Type-safe expression construction for typed DataFrames. module DataFrame.Typed.Expr ( -- * Core typed expression type (re-exported from Types) TExpr (..), @@ -52,6 +27,8 @@ -- * Conditional ifThenElse, + caseWhen, + orElse, -- * Unary / binary lifting lift, @@ -123,6 +100,9 @@ -- * Sort helpers asc, desc, + + -- * Additional expression functions + module DataFrame.Typed.Expr.Extra, ) where import Data.Either (fromRight) @@ -134,12 +114,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 +136,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, @@ -168,16 +149,7 @@ import DataFrame.Typed.Types (TExpr (..), TSortOrder (..)) import Prelude hiding (maximum, minimum, sum) -{- | Create a typed column reference. This is the key type-safety entry point. - -The column name must exist in @cols@ and its type must match @a@. -Both checks happen at compile time via type families. - -@ -salary :: TExpr '[Column \"salary\" Double] Double -salary = col \@\"salary\" -@ --} +-- | Create a typed column reference, checked against the schema. col :: forall (name :: Symbol) cols a. ( KnownSymbol name @@ -188,13 +160,7 @@ TExpr cols a col = TExpr (Col (T.pack (symbolVal (Proxy @name)))) -{- | Use a column name as an @OverloadedLabels@ label: @#age@ is sugar for -@col \@\"age\"@. Enable @OverloadedLabels@ at the use site. - -@ -adults = filterWhere (#age .>=. lit 18) people -@ --} +-- | Use a column name as an @OverloadedLabels@ label such as @#age@. instance ( KnownSymbol name , a ~ SafeLookup name cols @@ -205,9 +171,7 @@ where fromLabel = col @name -{- | Create a literal expression. Valid for any schema since it -references no columns. --} +-- | Create a literal expression, valid for any schema. lit :: (Columnable a) => a -> TExpr cols a lit = TExpr . Lit @@ -217,10 +181,16 @@ 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) -------------------------------------------------------------------------------- +-- | One branch of a chained conditional; the fallback goes last. +caseWhen :: + (Columnable a) => + TExpr cols Bool -> TExpr cols a -> TExpr cols a -> TExpr cols a +caseWhen (TExpr c) (TExpr t) (TExpr e) = TExpr (If c t e) +-- | Closes a 'caseWhen' chain with its fallback value. +orElse :: TExpr cols a -> TExpr cols a +orElse = id + 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 +226,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 @@ -271,10 +237,7 @@ (a -> b -> c) -> TExpr cols a -> TExpr cols b -> TExpr cols c lift2 f (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp f "binaryUdf" Nothing False 0) a b) -{- | Typed 'nullLift': lift a unary function with nullable propagation. -When the input is @Maybe a@, 'Nothing' short-circuits; when plain @a@, applies directly. -The return type is inferred via 'NullLift1Result': no annotation needed. --} +-- | Typed 'nullLift': lift a unary function with nullable propagation. nullLift :: (NullLift1Op a r (NullLift1Result a r), Columnable (NullLift1Result a r)) => (BaseType a -> r) -> @@ -282,10 +245,7 @@ TExpr cols (NullLift1Result a r) nullLift f (TExpr e) = TExpr (Unary (MkUnaryOp (applyNull1 f) "nullLift" Nothing) e) -{- | Typed 'nullLift2': lift a binary function with nullable propagation. -Any 'Nothing' operand short-circuits to 'Nothing' in the result. -The return type is inferred via 'NullLift2Result': no annotation needed. --} +-- | Typed 'nullLift2': lift a binary function with nullable propagation. nullLift2 :: (NullLift2Op a b r (NullLift2Result a b r), Columnable (NullLift2Result a b r)) => (BaseType a -> BaseType b -> r) -> @@ -359,7 +319,7 @@ (.||.) :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool (.||.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (||) "or" (Just ".||.") True 2) a b) --- | Nullable-aware logical AND. Returns @Maybe Bool@ when either operand is nullable. +-- | Nullable-aware logical AND; @Maybe Bool@ when either operand is nullable. (.&&) :: (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) => TExpr cols a -> @@ -368,7 +328,7 @@ (.&&) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (nullCmpOp (&&)) "nulland" (Just ".&&") True 3) a b) --- | Nullable-aware logical OR. Returns @Maybe Bool@ when either operand is nullable. +-- | Nullable-aware logical OR; @Maybe Bool@ when either operand is nullable. (.||) :: (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) => TExpr cols a -> @@ -377,16 +337,10 @@ (.||) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (nullCmpOp (||)) "nullor" (Just ".||") True 2) a b) -------------------------------------------------------------------------------- --- Nullable-aware arithmetic operators -------------------------------------------------------------------------------- - infixl 6 .+, .- infixl 7 .*, ./ -{- | Nullable-aware addition. Works for all combinations of nullable\/non-nullable operands. -@col \@\"x\" '.+' col \@\"y\" -- :: TExpr cols (Maybe Int) when y :: Maybe Int@ --} +-- | Nullable-aware addition for nullable and non-nullable operands. (.+) :: ( NumericWidenOp (BaseType a) (BaseType b) , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b) @@ -482,14 +436,7 @@ (.^) (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@ -when either operand is nullable. --} +-- | Nullable-aware equality that widens numeric operands to a common type. (.==) :: ( NumericWidenOp (BaseType a) (BaseType b) , NullLift2Op a b Bool (NullCmpResult a b) @@ -528,7 +475,7 @@ TExpr (Binary (MkBinaryOp (applyNull2 (widenCmpOp (<))) "lt" (Just "<") False 4) a b) --- | Nullable-aware less-than-or-equal. Widens numeric operands to their common type. +-- | Nullable-aware <=. Widens numeric operands to their common type. (.<=) :: ( NumericWidenOp (BaseType a) (BaseType b) , NullLift2Op a b Bool (NullCmpResult a b) @@ -541,7 +488,7 @@ TExpr (Binary (MkBinaryOp (applyNull2 (widenCmpOp (<=))) "leq" (Just "<=") False 4) a b) --- | Nullable-aware greater-than-or-equal. Widens numeric operands to their common type. +-- | Nullable-aware >=. Widens numeric operands to their common type. (.>=) :: ( NumericWidenOp (BaseType a) (BaseType b) , NullLift2Op a b Bool (NullCmpResult a b) @@ -570,10 +517,6 @@ not :: TExpr cols Bool -> TExpr cols Bool not (TExpr e) = TExpr (Unary (MkUnaryOp Prelude.not "not" (Just "!")) e) -------------------------------------------------------------------------------- --- Aggregation combinators -------------------------------------------------------------------------------- - sum :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a sum (TExpr e) = TExpr (F.sum e) @@ -605,10 +548,6 @@ (Columnable a, AllKnownSymbol names, AssertAllPresent names cols) => TExpr cols a -> TExpr cols a over (TExpr e) = TExpr{unTExpr = F.over (symbolVals @names) e} - -------------------------------------------------------------------------------- --- Cast / coercion expressions -------------------------------------------------------------------------------- castExpr :: forall b cols src. diff --git a/src/DataFrame/Typed/Expr/Extra.hs b/src/DataFrame/Typed/Expr/Extra.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Expr/Extra.hs @@ -0,0 +1,217 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Typed counterparts of the remaining "DataFrame.Functions" combinators. +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. +reduce :: + (Columnable a, Columnable b) => + TExpr cols b -> a -> (a -> b -> a) -> TExpr cols a +reduce (TExpr e) start f = TExpr (F.reduce e start f) + +-- | Wrap each value in 'Just'. +toMaybe :: (Columnable a) => TExpr cols a -> TExpr cols (Maybe a) +toMaybe (TExpr e) = TExpr (F.toMaybe e) + +-- | Replace 'Nothing' with a default. +fromMaybe :: (Columnable a) => a -> TExpr cols (Maybe a) -> TExpr cols a +fromMaybe d (TExpr e) = TExpr (F.fromMaybe d e) + +-- | True where the value is 'Just'. +isJust :: (Columnable a) => TExpr cols (Maybe a) -> TExpr cols Bool +isJust (TExpr e) = TExpr (F.isJust e) + +-- | True where the value is 'Nothing'. +isNothing :: (Columnable a) => TExpr cols (Maybe a) -> TExpr cols Bool +isNothing (TExpr e) = TExpr (F.isNothing e) + +-- | Unwrap a 'Just', erroring on 'Nothing'. +fromJust :: (Columnable a) => TExpr cols (Maybe a) -> TExpr cols a +fromJust (TExpr e) = TExpr (F.fromJust e) + +-- | Apply a function only where the value is present. +whenPresent :: + (Columnable a, Columnable b) => + (a -> b) -> TExpr cols (Maybe a) -> TExpr cols (Maybe b) +whenPresent f (TExpr e) = TExpr (F.whenPresent f e) + +-- | Apply a binary function only where both values are present. +whenBothPresent :: + (Columnable a, Columnable b, Columnable c) => + (a -> b -> c) -> + TExpr cols (Maybe a) -> + TExpr cols (Maybe b) -> + TExpr cols (Maybe c) +whenBothPresent f (TExpr a) (TExpr b) = TExpr (F.whenBothPresent f a b) + +-- | Map values through a lookup table, yielding 'Nothing' for misses. +recode :: + (Columnable a, Columnable b, Show a, Show b, Show (a, b)) => + [(a, b)] -> TExpr cols a -> TExpr cols (Maybe b) +recode mapping (TExpr e) = TExpr (F.recode mapping e) + +-- | Pick the first value whose condition holds, else a fallback. +recodeWithCondition :: + (Columnable a, Columnable b) => + TExpr cols b -> + [(TExpr cols a -> TExpr cols Bool, b)] -> + TExpr cols a -> + TExpr cols b +recodeWithCondition (TExpr fallback) conds (TExpr e) = + TExpr (F.recodeWithCondition fallback (map untype conds) e) + where + untype (p, v) = (unTExpr . p . TExpr, v) + +-- | Map values through a lookup table, with a default for misses. +recodeWithDefault :: + (Columnable a, Columnable b, Show (a, b)) => + b -> [(a, b)] -> TExpr cols a -> TExpr cols b +recodeWithDefault d mapping (TExpr e) = TExpr (F.recodeWithDefault d mapping e) + +-- | First element of a list column, or 'Nothing'. +firstOrNothing :: (Columnable a) => TExpr cols [a] -> TExpr cols (Maybe a) +firstOrNothing (TExpr e) = TExpr (F.firstOrNothing e) + +-- | Last element of a list column, or 'Nothing'. +lastOrNothing :: (Columnable a) => TExpr cols [a] -> TExpr cols (Maybe a) +lastOrNothing (TExpr e) = TExpr (F.lastOrNothing e) + +-- | Split text on a delimiter. +splitOn :: T.Text -> TExpr cols T.Text -> TExpr cols [T.Text] +splitOn delim (TExpr e) = TExpr (F.splitOn delim e) + +-- | First regex match, or 'Nothing'. +match :: T.Text -> TExpr cols T.Text -> TExpr cols (Maybe T.Text) +match regex (TExpr e) = TExpr (F.match regex e) + +-- | All regex matches. +matchAll :: T.Text -> TExpr cols T.Text -> TExpr cols [T.Text] +matchAll regex (TExpr e) = TExpr (F.matchAll regex e) + +-- | Parse text into a time value with the given format. +parseDate :: + (ParseTime t, Columnable t) => + T.Text -> TExpr cols T.Text -> TExpr cols (Maybe t) +parseDate format (TExpr e) = TExpr (F.parseDate format e) + +-- | Number of days between two dates. +daysBetween :: TExpr cols Day -> TExpr cols Day -> TExpr cols Int +daysBetween (TExpr a) (TExpr b) = TExpr (F.daysBetween a b) + +-- | Monadic bind over a column of monadic values. +bind :: + (Columnable a, Columnable (m a), Monad m, Columnable b, Columnable (m b)) => + (a -> m b) -> TExpr cols (m a) -> TExpr cols (m b) +bind f (TExpr e) = TExpr (F.bind f e) diff --git a/src/DataFrame/Typed/Operations.hs b/src/DataFrame/Typed/Operations.hs --- a/src/DataFrame/Typed/Operations.hs +++ b/src/DataFrame/Typed/Operations.hs @@ -17,6 +17,7 @@ filterAllJust, filterJust, filterNothing, + filterAllNothing, sortBy, take, takeLast, @@ -42,6 +43,10 @@ dropColumn, replaceColumn, + -- * Frequencies + valueCounts, + valueProportions, + -- * Metadata dimensions, nRows, @@ -80,15 +85,8 @@ 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. --} +-- | Filter rows where a boolean expression evaluates to True. filterWhere :: TExpr cols Bool -> TypedDataFrame cols -> TypedDataFrame cols filterWhere (TExpr expr) (TDF df) = TDF (D.filterWhere expr df) @@ -104,24 +102,11 @@ (a -> Bool) -> TExpr cols a -> TypedDataFrame cols -> TypedDataFrame cols filterBy pred' (TExpr expr) (TDF df) = TDF (D.filterBy pred' expr df) -{- | Keep only rows where ALL Optional columns have Just values. -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] -@ --} +-- | Keep only rows where all Optional columns have Just values. filterAllJust :: TypedDataFrame cols -> TypedDataFrame (StripAllMaybe cols) filterAllJust (TDF df) = unsafeFreeze (D.filterAllJust df) -{- | Keep only rows where the named column has Just values. -Strips 'Maybe' from that column's type in the result schema. - -@ -filterJust \@\"x\" df -@ --} +-- | Keep only rows where the named column has Just values. filterJust :: forall name cols. ( KnownSymbol name @@ -132,9 +117,7 @@ where colName = T.pack (symbolVal (Proxy @name)) -{- | Keep only rows where the named column has Nothing. -Schema is preserved (column types unchanged, just fewer rows). --} +-- | Keep only rows where the named column has Nothing. filterNothing :: forall name cols. ( KnownSymbol name @@ -145,9 +128,11 @@ where colName = T.pack (symbolVal (Proxy @name)) -{- | Sort by the given typed sort orders. -Sort orders reference columns that are validated against the schema. --} +-- | Keep only rows where every nullable column has Nothing. +filterAllNothing :: TypedDataFrame cols -> TypedDataFrame cols +filterAllNothing (TDF df) = TDF (D.filterAllNothing df) + +-- | Sort by the given typed sort orders, validated against the schema. sortBy :: [TSortOrder cols] -> TypedDataFrame cols -> TypedDataFrame cols sortBy ords (TDF df) = TDF (D.sortBy (map toUntypedSort ords) df) where @@ -192,19 +177,7 @@ 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) -@ --} +-- | Derive a new column from a typed expression; its name must be absent. derive :: forall name a cols. ( KnownSymbol name @@ -213,7 +186,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 +248,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 +260,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 +272,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 +285,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)) @@ -329,9 +302,7 @@ where colName = T.pack (symbolVal (Proxy @name)) -{- | Replace an existing column with new values derived from a typed expression. -The column must already exist and the new type must match. --} +-- | Replace an existing column with values from a typed expression. replaceColumn :: forall name a cols. ( KnownSymbol name @@ -348,13 +319,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) @@ -363,9 +327,7 @@ intersect :: TypedDataFrame cols -> TypedDataFrame cols -> TypedDataFrame cols intersect (TDF a) (TDF b) = TDF (DS.intersect a b) -{- | Rows in the left DataFrame but not the right, deduplicated -(relational @EXCEPT@; the subobject complement). --} +-- | Rows in the left DataFrame but not the right, deduplicated. difference :: TypedDataFrame cols -> TypedDataFrame cols -> TypedDataFrame cols difference (TDF a) (TDF b) = TDF (DS.difference a b) @@ -374,10 +336,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 +348,15 @@ columnNames :: TypedDataFrame cols -> [T.Text] columnNames (TDF df) = D.columnNames df -------------------------------------------------------------------------------- --- Internal helpers -------------------------------------------------------------------------------- +-- | Count occurrences of each distinct value in a column. +valueCounts :: + (Ord a, Columnable a) => TExpr cols a -> TypedDataFrame cols -> [(a, Int)] +valueCounts (TExpr e) (TDF df) = D.valueCounts e df + +-- | Proportion of each distinct value in a column. +valueProportions :: + (Ord a, Columnable a) => TExpr cols a -> TypedDataFrame cols -> [(a, Double)] +valueProportions (TExpr e) (TDF df) = D.valueProportions e df -- | Helper class for extracting [(Text, Text)] from type-level pairs. class AllKnownPairs (pairs :: [(Symbol, Symbol)]) where diff --git a/src/DataFrame/Typed/Sampling.hs b/src/DataFrame/Typed/Sampling.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Sampling.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE FlexibleContexts #-} + +-- | Typed sampling and splitting; all operations preserve the schema. +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; may fail if the indices are out of bounds. +selectRows :: [Int] -> TypedDataFrame cols -> TypedDataFrame cols +selectRows ixs (TDF df) = TDF (D.selectRows ixs df) + +-- | Sample a fraction of rows, preserving the distribution of a strata column. +stratifiedSample :: + (SplittableGen g, Columnable a) => + g -> Double -> TExpr cols a -> TypedDataFrame cols -> TypedDataFrame cols +stratifiedSample g p (TExpr e) (TDF df) = TDF (D.stratifiedSample g p e df) + +-- | Split rows by a fraction, preserving the distribution of a strata column. +stratifiedSplit :: + (SplittableGen g, Columnable a) => + g -> + Double -> + TExpr cols a -> + TypedDataFrame cols -> + (TypedDataFrame cols, TypedDataFrame cols) +stratifiedSplit g p (TExpr e) (TDF df) = + let (a, b) = D.stratifiedSplit g p e df in (TDF a, TDF b) diff --git a/src/DataFrame/Typed/Statistics.hs b/src/DataFrame/Typed/Statistics.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Statistics.hs @@ -0,0 +1,152 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} + +-- | Typed statistical reducers over a 'TypedDataFrame'; import qualified. +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 named numeric columns. +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, returned as an untyped 'D.DataFrame'. +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. +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 as an untyped 'D.DataFrame'. +summarize :: TypedDataFrame cols -> D.DataFrame +summarize (TDF df) = Stats.summarize df + +-- | Per-column summary of counts, unique values and type, untyped. +describeColumns :: TypedDataFrame cols -> D.DataFrame +describeColumns (TDF df) = Core.describeColumns df