diff --git a/dataframe-operations.cabal b/dataframe-operations.cabal --- a/dataframe-operations.cabal +++ b/dataframe-operations.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: dataframe-operations -version: 2.5.0.1 +version: 2.7.0.0 synopsis: Column operations, expression DSL, and statistics for the dataframe ecosystem. description: Untyped column operations (select, filter, sort, join, groupBy, @@ -31,6 +31,7 @@ import: warnings exposed-modules: DataFrame.Internal.Statistics + DataFrame.Expression.Parse DataFrame.Functions DataFrame.Monad DataFrame.Operations.Aggregation.Run @@ -51,6 +52,7 @@ DataFrame.Typed.Apply DataFrame.Typed.Expr DataFrame.Typed.Expr.Extra + DataFrame.Typed.Functions DataFrame.Typed.Join DataFrame.Typed.Operations DataFrame.Typed.Sampling @@ -58,7 +60,7 @@ build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.14, containers >= 0.6.7 && < 0.10, - dataframe-core >= 2.5 && < 2.6, + dataframe-core >= 2.7 && < 2.8, dataframe-parsing >= 2.2 && < 2.3, random >= 1.2 && < 2, regex-tdfa >= 1.3.0 && < 2, diff --git a/src-internal/DataFrame/Internal/Statistics.hs b/src-internal/DataFrame/Internal/Statistics.hs --- a/src-internal/DataFrame/Internal/Statistics.hs +++ b/src-internal/DataFrame/Internal/Statistics.hs @@ -14,6 +14,10 @@ import Control.Monad.ST (runST) import DataFrame.Errors (DataFrameException (..)) +-- | The present values of a nullable vector, as doubles, in reverse order. +optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double +optionalToDoubleVector = VU.fromList . V.foldl' (\acc e -> maybe acc ((: acc) . realToFrac) e) [] + mean' :: (Real a, VU.Unbox a) => VU.Vector a -> Double mean' samp | VU.null samp = throw $ EmptyDataSetException "mean" @@ -186,8 +190,7 @@ 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. + -- Not exact for Ord instances; TODO make it so. VM.read mutableSamp index ) (V.convert qs) 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.int c) = Just TInt + | ok (F.double c) = Just TDouble + | ok (F.text c) = Just TText + | ok (F.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.int c) + TDouble -> SomeExpr (F.double c) + TText -> SomeExpr (F.text c) + TBool -> SomeExpr (F.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 @@ -27,13 +27,13 @@ import DataFrame.Internal.Statistics import Control.Applicative -import qualified Data.Char as Char import Data.Either import Data.Function (on) import Data.Int 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)) @@ -55,7 +55,8 @@ lift :: (Columnable a, Columnable b) => (a -> b) -> Expr a -> Expr b lift f = - Unary (MkUnaryOp{unaryFn = f, unaryName = "unaryUdf", unarySymbol = Nothing}) + Unary + (MkUnaryOp{unaryFn = f, unaryId = Anonymous "unaryUdf", unarySymbol = Nothing}) lift2 :: (Columnable c, Columnable b, Columnable a) => @@ -67,19 +68,14 @@ Binary ( MkBinaryOp { binaryFn = f - , binaryName = "binaryUdf" + , binaryId = Anonymous "binaryUdf" , binarySymbol = Nothing , binaryCommutative = False , binaryPrecedence = 0 } ) -{- | 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) -> @@ -87,13 +83,14 @@ Expr (NullLift1Result a r) nullLift f = 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. + ( MkUnaryOp + { unaryFn = applyNull1 f + , unaryId = Anonymous "nullLift" + , unarySymbol = Nothing + } + ) -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) -> @@ -104,50 +101,36 @@ Binary ( MkBinaryOp { binaryFn = applyNull2 f - , binaryName = "nullLift2" + , binaryId = Anonymous "nullLift2" , binarySymbol = Nothing , binaryCommutative = False , binaryPrecedence = 0 } ) -{- | 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) +cast colName = CastWith colName (Builtin CCast) (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) + CastWith + colName + (Builtin (CCastWithDefault (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) +castEither colName = CastWith colName (Builtin CCastEither) (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 colName - "unsafeCast" + (Builtin CUnsafeCast) (fromRight (error "unsafeCast: unexpected Nothing in column")) castExpr :: @@ -155,13 +138,13 @@ (Columnable b, Columnable src, Read b) => Expr src -> Expr (Maybe b) -castExpr = CastExprWith @b @(Maybe b) @src "castExpr" (either (const Nothing) Just) +castExpr = CastExprWith @b @(Maybe b) @src (Builtin CCast) (either (const Nothing) Just) castExprWithDefault :: forall b src. (Columnable b, Columnable src, Read b) => b -> Expr src -> Expr b castExprWithDefault def = CastExprWith @b @b @src - ("castExprWithDefault:" <> T.pack (show def)) + (Builtin (CCastWithDefault (T.pack (show def)))) (fromRight def) castExprEither :: @@ -171,14 +154,14 @@ Expr (Either T.Text b) castExprEither = CastExprWith @b @(Either T.Text b) @src - "castExprEither" + (Builtin CCastEither) (either (Left . T.pack) Right) unsafeCastExpr :: forall b src. (Columnable b, Columnable src, Read b) => Expr src -> Expr b unsafeCastExpr = CastExprWith @b @b @src - "unsafeCastExpr" + (Builtin CUnsafeCast) (fromRight (error "unsafeCastExpr: unexpected Nothing in column")) toDouble :: (Columnable a, Real a) => Expr a -> Expr Double @@ -186,34 +169,34 @@ Unary ( MkUnaryOp { unaryFn = realToFrac - , unaryName = "toDouble" + , unaryId = Builtin UToDouble , unarySymbol = Nothing } ) infix 8 `div` div :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a -div = lift2Decorated Prelude.div "div" (Just "//") False 7 +div = lift2Builtin BDiv Prelude.div (Just "//") False 7 mod :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a -mod = lift2Decorated Prelude.mod "mod" Nothing False 7 +mod = lift2Builtin BMod Prelude.mod Nothing False 7 eq :: (Columnable a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool -eq = lift2Decorated (==) "eq" (Just "==") True 4 +eq = lift2Builtin BEq (==) (Just "==") True 4 lt :: (Columnable a, Ord a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool -lt = lift2Decorated (<) "lt" (Just "<") False 4 +lt = lift2Builtin BLt (<) (Just "<") False 4 gt :: (Columnable a, Ord a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool -gt = lift2Decorated (>) "gt" (Just ">") False 4 +gt = lift2Builtin BGt (>) (Just ">") False 4 leq :: (Columnable a, Ord a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool -leq = lift2Decorated (<=) "leq" (Just "<=") False 4 +leq = lift2Builtin BLeq (<=) (Just "<=") False 4 geq :: (Columnable a, Ord a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool -geq = lift2Decorated (>=) "geq" (Just ">=") False 4 +geq = lift2Builtin BGeq (>=) (Just ">=") False 4 and :: Expr Bool -> Expr Bool -> Expr Bool and = (.&&) @@ -224,10 +207,10 @@ not :: Expr Bool -> Expr Bool not = Unary - (MkUnaryOp{unaryFn = Prelude.not, unaryName = "not", unarySymbol = Just "~"}) + (MkUnaryOp{unaryFn = Prelude.not, unaryId = Builtin UNot, unarySymbol = Just "~"}) count :: (Columnable a) => Expr a -> Expr Int -count = Agg (MergeAgg "count" (0 :: Int) (\c _ -> c + 1) (+) id) +count = Agg (MergeAgg (Builtin ACount) (0 :: Int) (\c _ -> c + 1) (+) id) {-# SPECIALIZE count :: Expr Double -> Expr Int #-} {-# SPECIALIZE count :: Expr Float -> Expr Int #-} {-# SPECIALIZE count :: Expr Int -> Expr Int #-} @@ -237,13 +220,32 @@ {-# 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 + (Builtin ACountDistinct) + (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)) {-# INLINE countAll #-} collect :: (Columnable a) => Expr a -> Expr [a] -collect = Agg (FoldAgg "collect" (Just []) (flip (:))) +collect = Agg (FoldAgg (Builtin ACollect) (Just []) (flip (:))) {-# SPECIALIZE collect :: Expr Double -> Expr [Double] #-} {-# SPECIALIZE collect :: Expr Float -> Expr [Float] #-} {-# SPECIALIZE collect :: Expr Int -> Expr [Int] #-} @@ -253,7 +255,7 @@ mode = Agg ( CollectAgg - "mode" + (Builtin AMode) ( fst . L.maximumBy (compare `on` snd) . M.toList @@ -270,7 +272,7 @@ {-# INLINEABLE mode #-} minimum :: (Columnable a, Ord a) => Expr a -> Expr a -minimum = Agg (FoldAgg "minimum" Nothing Prelude.min) +minimum = Agg (FoldAgg (Builtin AMinimum) Nothing Prelude.min) {-# SPECIALIZE DataFrame.Functions.minimum :: Expr Double -> Expr Double #-} {-# SPECIALIZE DataFrame.Functions.minimum :: Expr Float -> Expr Float #-} {-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int -> Expr Int #-} @@ -281,7 +283,7 @@ {-# INLINEABLE DataFrame.Functions.minimum #-} maximum :: (Columnable a, Ord a) => Expr a -> Expr a -maximum = Agg (FoldAgg "maximum" Nothing Prelude.max) +maximum = Agg (FoldAgg (Builtin AMaximum) Nothing Prelude.max) {-# SPECIALIZE DataFrame.Functions.maximum :: Expr Double -> Expr Double #-} {-# SPECIALIZE DataFrame.Functions.maximum :: Expr Float -> Expr Float #-} {-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int -> Expr Int #-} @@ -292,7 +294,7 @@ {-# INLINEABLE DataFrame.Functions.maximum #-} sum :: forall a. (Columnable a, Num a) => Expr a -> Expr a -sum = Agg (FoldAgg "sum" Nothing (+)) +sum = Agg (FoldAgg (Builtin ASum) Nothing (+)) {-# SPECIALIZE DataFrame.Functions.sum :: Expr Double -> Expr Double #-} {-# SPECIALIZE DataFrame.Functions.sum :: Expr Float -> Expr Float #-} {-# SPECIALIZE DataFrame.Functions.sum :: Expr Int -> Expr Int #-} @@ -303,7 +305,7 @@ {-# INLINEABLE DataFrame.Functions.sum #-} sumMaybe :: forall a. (Columnable a, Num a) => Expr (Maybe a) -> Expr a -sumMaybe = Agg (CollectAgg "sumMaybe" (P.sum . Maybe.catMaybes . V.toList)) +sumMaybe = Agg (CollectAgg (Builtin ASumMaybe) (P.sum . Maybe.catMaybes . V.toList)) {-# SPECIALIZE sumMaybe :: Expr (Maybe Double) -> Expr Double #-} {-# SPECIALIZE sumMaybe :: Expr (Maybe Float) -> Expr Float #-} {-# SPECIALIZE sumMaybe :: Expr (Maybe Int) -> Expr Int #-} @@ -317,7 +319,7 @@ mean = Agg ( MergeAgg - "mean" + (Builtin AMean) ((0.0, 0) :: (Double, Int)) (\(!s, !c) x -> (s + realToFrac x, c + 1)) (\(!s1, !c1) (!s2, !c2) -> (s1 + s2, c1 + c2)) @@ -334,33 +336,28 @@ -- | The k largest values per group. topK :: (Columnable a, Ord a) => Int -> Expr a -> Expr [a] -topK = kExtremes "topK" (>) +topK = kExtremes True {-# INLINEABLE topK #-} -- | The k smallest values per group. bottomK :: (Columnable a, Ord a) => Int -> Expr a -> Expr [a] -bottomK = kExtremes "bottomK" (<) +bottomK = kExtremes False {-# INLINEABLE bottomK #-} -{- | Top k values after applying an ordering function. Floating-point NaNs -are dropped: they have no place in a total order, and admitting them would -make the merge non-associative (results would depend on chunk boundaries in -the lazy executor). --} +-- | The k largest (True) or smallest (False) values, dropping NaNs. kExtremes :: - forall a. - (Columnable a, Ord a) => - T.Text -> (a -> a -> Bool) -> Int -> Expr a -> Expr [a] -kExtremes opName wins k = + forall a. (Columnable a, Ord a) => Bool -> Int -> Expr a -> Expr [a] +kExtremes largest k = Agg ( MergeAgg - (opName <> "_" <> T.pack (show k)) + (Builtin (if largest then ATopK k else ABottomK k)) [] step (L.foldl' step) id ) where + wins = if largest then (>) else (<) step acc x | isNaNLike x = acc | otherwise = insert x acc @@ -376,7 +373,7 @@ | otherwise = False meanMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double -meanMaybe = Agg (CollectAgg "meanMaybe" (mean' . optionalToDoubleVector)) +meanMaybe = Agg (CollectAgg (Builtin AMeanMaybe) (mean' . optionalToDoubleVector)) {-# SPECIALIZE meanMaybe :: Expr (Maybe Double) -> Expr Double #-} {-# SPECIALIZE meanMaybe :: Expr (Maybe Float) -> Expr Double #-} {-# SPECIALIZE meanMaybe :: Expr (Maybe Int) -> Expr Double #-} @@ -387,7 +384,7 @@ {-# INLINEABLE meanMaybe #-} variance :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double -variance = Agg (CollectAgg "variance" variance') +variance = Agg (CollectAgg (Builtin AVariance) variance') {-# SPECIALIZE variance :: Expr Double -> Expr Double #-} {-# SPECIALIZE variance :: Expr Float -> Expr Double #-} {-# SPECIALIZE variance :: Expr Int -> Expr Double #-} @@ -398,7 +395,7 @@ {-# INLINEABLE variance #-} median :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double -median = Agg (CollectAgg "median" median') +median = Agg (CollectAgg (Builtin AMedian) median') {-# SPECIALIZE median :: Expr Double -> Expr Double #-} {-# SPECIALIZE median :: Expr Float -> Expr Double #-} {-# SPECIALIZE median :: Expr Int -> Expr Double #-} @@ -409,7 +406,7 @@ {-# INLINEABLE median #-} medianMaybe :: (Columnable a, Real a) => Expr (Maybe a) -> Expr Double -medianMaybe = Agg (CollectAgg "meanMaybe" (median' . optionalToDoubleVector)) +medianMaybe = Agg (CollectAgg (Builtin AMedianMaybe) (median' . optionalToDoubleVector)) {-# SPECIALIZE medianMaybe :: Expr (Maybe Double) -> Expr Double #-} {-# SPECIALIZE medianMaybe :: Expr (Maybe Float) -> Expr Double #-} {-# SPECIALIZE medianMaybe :: Expr (Maybe Int) -> Expr Double #-} @@ -419,23 +416,16 @@ {-# SPECIALIZE medianMaybe :: Expr (Maybe Int64) -> Expr Double #-} {-# INLINEABLE medianMaybe #-} -optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double -optionalToDoubleVector = - VU.fromList - . V.foldl' - (\acc e -> if Maybe.isJust e then realToFrac (Maybe.fromMaybe 0 e) : acc else acc) - [] - percentile :: Int -> Expr Double -> Expr Double percentile n = Agg ( CollectAgg - (T.pack $ "percentile " ++ show n) + (Builtin (APercentile n)) (percentile' n) ) stddev :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double -stddev = Agg (CollectAgg "stddev" (sqrt . variance')) +stddev = Agg (CollectAgg (Builtin AStddev) (sqrt . variance')) {-# SPECIALIZE stddev :: Expr Double -> Expr Double #-} {-# SPECIALIZE stddev :: Expr Float -> Expr Double #-} {-# SPECIALIZE stddev :: Expr Int -> Expr Double #-} @@ -446,7 +436,9 @@ {-# INLINEABLE stddev #-} stddevMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double -stddevMaybe = Agg (CollectAgg "stddevMaybe" (sqrt . variance' . optionalToDoubleVector)) +stddevMaybe = + Agg + (CollectAgg (Builtin AStddevMaybe) (sqrt . variance' . optionalToDoubleVector)) {-# SPECIALIZE stddevMaybe :: Expr (Maybe Double) -> Expr Double #-} {-# SPECIALIZE stddevMaybe :: Expr (Maybe Float) -> Expr Double #-} {-# SPECIALIZE stddevMaybe :: Expr (Maybe Int) -> Expr Double #-} @@ -460,21 +452,21 @@ 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 = lift2Builtin BPow (^) (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 #-} {-# INLINEABLE pow #-} relu :: (Columnable a, Num a, Ord a) => Expr a -> Expr a -relu = liftDecorated (Prelude.max 0) "relu" Nothing +relu = liftBuiltin URelu (Prelude.max 0) Nothing {-# SPECIALIZE relu :: Expr Double -> Expr Double #-} {-# SPECIALIZE relu :: Expr Float -> Expr Float #-} {-# SPECIALIZE relu :: Expr Int -> Expr Int #-} {-# INLINEABLE relu #-} min :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr a -min = lift2Decorated Prelude.min "min" Nothing True 1 +min = lift2Builtin BMin Prelude.min Nothing True 1 {-# SPECIALIZE DataFrame.Functions.min :: Expr Double -> Expr Double -> Expr Double #-} @@ -485,7 +477,7 @@ {-# INLINEABLE DataFrame.Functions.min #-} max :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr a -max = lift2Decorated Prelude.max "max" Nothing True 1 +max = lift2Builtin BMax Prelude.max Nothing True 1 {-# SPECIALIZE DataFrame.Functions.max :: Expr Double -> Expr Double -> Expr Double #-} @@ -502,31 +494,31 @@ a -> (a -> b -> a) -> Expr a -reduce expr start f = Agg (FoldAgg "foldUdf" (Just start) f) expr +reduce expr start f = Agg (FoldAgg (Anonymous "foldUdf") (Just start) f) expr {-# INLINEABLE reduce #-} toMaybe :: (Columnable a) => Expr a -> Expr (Maybe a) -toMaybe = liftDecorated Just "toMaybe" Nothing +toMaybe = liftBuiltin UToMaybe Just Nothing {-# SPECIALIZE toMaybe :: Expr Double -> Expr (Maybe Double) #-} {-# SPECIALIZE toMaybe :: Expr Float -> Expr (Maybe Float) #-} {-# SPECIALIZE toMaybe :: Expr Int -> Expr (Maybe Int) #-} {-# INLINEABLE toMaybe #-} fromMaybe :: (Columnable a) => a -> Expr (Maybe a) -> Expr a -fromMaybe d = liftDecorated (Maybe.fromMaybe d) "fromMaybe" Nothing +fromMaybe d = liftBuiltin (UFromMaybe (T.pack (show d))) (Maybe.fromMaybe d) Nothing {-# SPECIALIZE fromMaybe :: Double -> Expr (Maybe Double) -> Expr Double #-} {-# SPECIALIZE fromMaybe :: Float -> Expr (Maybe Float) -> Expr Float #-} {-# SPECIALIZE fromMaybe :: Int -> Expr (Maybe Int) -> Expr Int #-} {-# INLINEABLE fromMaybe #-} isJust :: (Columnable a) => Expr (Maybe a) -> Expr Bool -isJust = liftDecorated Maybe.isJust "isJust" Nothing +isJust = liftBuiltin UIsJust Maybe.isJust Nothing {-# SPECIALIZE isJust :: Expr (Maybe Double) -> Expr Bool #-} {-# SPECIALIZE isJust :: Expr (Maybe Int) -> Expr Bool #-} {-# INLINEABLE isJust #-} isNothing :: (Columnable a) => Expr (Maybe a) -> Expr Bool -isNothing = liftDecorated Maybe.isNothing "isNothing" Nothing +isNothing = liftBuiltin UIsNothing Maybe.isNothing Nothing {-# SPECIALIZE isNothing :: Expr (Maybe Double) -> Expr Bool #-} {-# SPECIALIZE isNothing :: Expr (Maybe Int) -> Expr Bool #-} {-# INLINEABLE isNothing #-} @@ -542,7 +534,7 @@ {-# INLINEABLE isNotNull #-} fromJust :: (Columnable a) => Expr (Maybe a) -> Expr a -fromJust = liftDecorated Maybe.fromJust "fromJust" Nothing +fromJust = liftBuiltin UFromJust Maybe.fromJust Nothing {-# SPECIALIZE fromJust :: Expr (Maybe Double) -> Expr Double #-} {-# SPECIALIZE fromJust :: Expr (Maybe Int) -> Expr Int #-} {-# INLINEABLE fromJust #-} @@ -553,7 +545,7 @@ (a -> b) -> Expr (Maybe a) -> Expr (Maybe b) -whenPresent f = liftDecorated (fmap f) "whenPresent" Nothing +whenPresent f = Unary (MkUnaryOp (fmap f) (Anonymous "whenPresent") Nothing) {-# INLINEABLE whenPresent #-} whenBothPresent :: @@ -563,7 +555,9 @@ Expr (Maybe a) -> Expr (Maybe b) -> Expr (Maybe c) -whenBothPresent f = lift2Decorated (\l r -> f <$> l <*> r) "whenBothPresent" Nothing False 0 +whenBothPresent f = + Binary + (MkBinaryOp (\l r -> f <$> l <*> r) (Anonymous "whenBothPresent") Nothing False 0) {-# INLINEABLE whenBothPresent #-} recode :: @@ -576,7 +570,7 @@ Unary ( MkUnaryOp { unaryFn = (`lookup` mapping) - , unaryName = "recode " <> T.pack (show mapping) + , unaryId = Builtin (URecode (T.pack (show mapping))) , unarySymbol = Nothing } ) @@ -602,50 +596,59 @@ Unary ( MkUnaryOp { unaryFn = Maybe.fromMaybe d . (`lookup` mapping) - , unaryName = - "recodeWithDefault " <> T.pack (show d) <> " " <> T.pack (show mapping) + , unaryId = + Builtin (URecodeWithDefault (T.pack (show d)) (T.pack (show mapping))) , unarySymbol = Nothing } ) firstOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a) -firstOrNothing = liftDecorated Maybe.listToMaybe "firstOrNothing" Nothing +firstOrNothing = liftBuiltin UFirstOrNothing Maybe.listToMaybe Nothing lastOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a) -lastOrNothing = liftDecorated (Maybe.listToMaybe . reverse) "lastOrNothing" Nothing +lastOrNothing = liftBuiltin ULastOrNothing (Maybe.listToMaybe . reverse) Nothing +-- | Whether the text contains the literal substring. +contains :: T.Text -> Expr T.Text -> Expr Bool +contains needle expr = + lift2Builtin BContains (flip T.isInfixOf) Nothing False 4 expr (Lit needle) + +-- | Character count. +length :: Expr T.Text -> Expr Int +length = liftBuiltin UTextLength T.length Nothing + splitOn :: T.Text -> Expr T.Text -> Expr [T.Text] -splitOn delim = liftDecorated (T.splitOn delim) "splitOn" Nothing +splitOn delim = liftBuiltin (USplitOn delim) (T.splitOn delim) Nothing match :: T.Text -> Expr T.Text -> Expr (Maybe T.Text) match regex = - liftDecorated + liftBuiltin + (UMatch regex) ((\r -> if T.null r then Nothing else Just r) . (=~ regex)) - ("match " <> T.pack (show regex)) Nothing matchAll :: T.Text -> Expr T.Text -> Expr [T.Text] matchAll regex = - liftDecorated + liftBuiltin + (UMatchAll regex) (getAllTextMatches . (=~ regex)) - ("matchAll " <> T.pack (show regex)) Nothing parseDate :: (ParseTime t, Columnable t) => T.Text -> Expr T.Text -> Expr (Maybe t) parseDate format = - liftDecorated + liftBuiltin + (UParseDate format) (parseTimeM True defaultTimeLocale (T.unpack format) . T.unpack) - ("parseDate " <> format) Nothing daysBetween :: Expr Day -> Expr Day -> Expr Int daysBetween = - lift2Decorated + lift2Builtin + BDaysBetween (\d1 d2 -> fromIntegral (diffDays d1 d2)) - "daysBetween" Nothing - True + False 2 bind :: @@ -654,88 +657,8 @@ (a -> m b) -> Expr (m a) -> 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\")) +bind f = Unary (MkUnaryOp (>>= f) (Anonymous "bind") Nothing) --- 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 -isReservedId :: T.Text -> Bool -isReservedId t = case t of - "case" -> True - "class" -> True - "data" -> True - "default" -> True - "deriving" -> True - "do" -> True - "else" -> True - "foreign" -> True - "if" -> True - "import" -> True - "in" -> True - "infix" -> True - "infixl" -> True - "infixr" -> True - "instance" -> True - "let" -> True - "module" -> True - "newtype" -> True - "of" -> True - "then" -> True - "type" -> True - "where" -> True - _ -> False - -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. - Just (c, _) -> Char.isLower c && Char.isAlpha c - Nothing -> False - -isHaskellIdentifier :: T.Text -> Bool -isHaskellIdentifier t = Prelude.not (isVarId t) || isReservedId t - -sanitize :: T.Text -> T.Text -sanitize t - | isValid = t - | isHaskellIdentifier t' = "_" <> t' <> "_" - | otherwise = t' - where - isValid = - Prelude.not (isHaskellIdentifier t) - && isVarId t - && T.all Char.isAlphaNum t - t' = T.map replaceInvalidCharacters . T.filter (Prelude.not . parentheses) $ t - replaceInvalidCharacters c - | Char.isUpper c = Char.toLower c - | Char.isSpace c = '_' - | Char.isPunctuation c = '_' -- '-' will also become a '_' - | Char.isSymbol c = '_' - | Char.isAlphaNum c = c -- Blanket condition - | otherwise = '_' -- If we're unsure we'll default to an underscore - parentheses c = case c of - '(' -> True - ')' -> True - '{' -> True - '}' -> True - '[' -> True - ']' -> True - _ -> False diff --git a/src/DataFrame/Monad.hs b/src/DataFrame/Monad.hs --- a/src/DataFrame/Monad.hs +++ b/src/DataFrame/Monad.hs @@ -20,7 +20,7 @@ deriveM, insertM, renameM, - filterWhereM, + filterM, sampleM, takeM, dropM, @@ -44,7 +44,7 @@ import Control.Monad (void) import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.DataFrame (DataFrame) -import DataFrame.Internal.Expression (Expr (..), UExpr (..), prettyPrint) +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) @@ -57,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 @@ -105,8 +103,8 @@ renameM (Col oldName) newName = execWithExpr newName (D.rename oldName newName) renameM expr newName = deriveM newName expr -filterWhereM :: Expr Bool -> FrameM () -filterWhereM p = modifyM (D.filterWhere p) +filterM :: Expr Bool -> FrameM () +filterM p = modifyM (D.filter p) sampleM :: (RandomGen g) => g -> Double -> FrameM () sampleM pureGen p = modifyM (D.sample pureGen p) @@ -153,7 +151,7 @@ execFrameM :: DataFrame -> FrameM a -> DataFrame execFrameM df m = snd (runFrameM df m) -newtype Pipeline a = Pipeline {unPipeline :: Int -> (Int, [(T.Text, UExpr)], a)} +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) @@ -173,18 +171,18 @@ -- | Derive a column under an explicit name, returning a reference to it. letAs :: (Columnable a) => T.Text -> Expr a -> Pipeline (Expr a) -letAs nm e = Pipeline (,[(nm, UExpr e)],Col nm) +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, UExpr e)], Col nm) + 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, UExpr e) -> T.unpack nm ++ " = " ++ prettyPrint e ++ "\n") + (\(nm, SomeExpr e) -> T.unpack nm ++ " = " ++ prettyPrint e ++ "\n") steps ++ "return " ++ prettyPrint res @@ -197,10 +195,8 @@ toFrameM :: Pipeline (Expr a) -> FrameM (Expr a) toFrameM p = let (_, steps, res) = unPipeline p 0 - in mapM_ (\(nm, UExpr e) -> void (deriveM nm e)) steps >> pure res + in mapM_ (\(nm, SomeExpr e) -> void (deriveM nm e)) steps >> pure res -{- | Run a pipeline over a frame: derive its columns and return the result -expression (a reference to the final column) and the resulting frame. --} +-- | 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/Aggregation.hs b/src/DataFrame/Operations/Aggregation.hs --- a/src/DataFrame/Operations/Aggregation.hs +++ b/src/DataFrame/Operations/Aggregation.hs @@ -62,11 +62,7 @@ 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) @@ -74,30 +70,14 @@ in unsafePerformIO (computeRowHashesIO n selectedCols) {-# NOINLINE computeRowHashes #-} -{- | Aggregate a grouped dataframe using the expressions given. -All ungrouped columns will be dropped. - -NOTE: this function deliberately never pattern-matches or strictly binds the -'Grouped' per-row fields (this module is compiled with @-XStrict@, whose strict -patterns and bindings would force them): on direct-grouped frames BOTH -'valueIndices' (the placement permutation) and 'rowToGroup' are deferred -thunks, and each aggregation path needs at most one of them — always passed as -un-forced argument expressions. Key columns materialize through 'groupRepRows' -(one representative row per group) instead of gathering -@valueIndices[offsets[g]]@. --} +-- | Aggregate a grouped dataframe using the expressions given. aggregate :: [NamedExpr] -> GroupedDataFrame -> DataFrame aggregate aggs gdf = let df = fullDataframe gdf offs = offsets gdf - {- Key columns materialize through ONE fused parallel gather over the - representative rows ('atIndicesStableMulti'): the 1e8-group Q10 result - was six sequential latency-bound random-gather passes as per-column - 'selectIndices'. Each result column is still identical to (and as - deferred as) the per-column gather; the row-count field uses the eager - @offsets@ so nothing here forces 'groupRepRows' early. -} + -- Key columns: one fused parallel gather over the representative rows. df' = let sub = select (groupedColumns gdf) df in sub @@ -113,21 +93,7 @@ !nGroups = VU.length offs - 1 !nRows' = fst (dataframeDimensions df) - {- Fused multi-reduction fast path (Q3/Q4/Q5-shaped aggregates): every - recognised simple scatter reduction (sum/mean/count/min/max over a clean - unboxed Int/Double column) in this aggregate runs in ONE pass instead of - one full pass per expression. At or below 'streamGroupCap' groups that - pass streams over (rowToGroup, columns) with per-worker accumulators — - never touching the (possibly lazy) valueIndices; above it (necessarily a - hash-path grouping, whose valueIndices is already eager) the accumulator - arrays would thrash, so the pass gathers by disjoint group range instead - (register accumulators, bit-identical to the unfused gather kernels, - one traversal instead of one per expression). Only taken when at - least two reductions fuse; non-fusable expressions (median, var/std, - max-min, arbitrary DSL) keep their per-expression path below. -} - {- This binding is strict (-XStrict), so it must stay empty-and-cheap - whenever the moment path below already covers the aggregate — otherwise - the pass would run redundantly before the moment result is consulted. -} + -- 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 @@ -154,13 +120,7 @@ , Just ga <- [mkGatherAgg nGroups (valueIndices gdf) offs red c] ] - in {- Unlike the stream branch, a SINGLE candidate also - takes this path: the per-expression fallback is a - scatter over rowToGroup, which on a hash-path grouping - is now a deferred thunk — the gather kernel (documented - bit-identical to the unfused kernels) works off the - already-eager valueIndices instead and skips that whole - random-write pass. -} + in -- Even a single candidate uses the gather path. if not (null cands) then MS.fromList @@ -175,12 +135,7 @@ ) else MS.empty - {- Fused median + std/var over one column (the Q6 shape): both are - holistic gathers over the same values, so one shared gather serves the - Welford fold and the median selection ('runMedianVarFused', - bit-identical to the separate kernels). Only built when a median and a - std/var on the same column appear together; empty-and-cheap otherwise - (same strictness caveat as 'fusedScatterCols'). -} + -- Median and std/var on one column share a single gather. medianVarCols :: MS.Map T.Text Column medianVarCols = case fusedMoments of Just _ -> MS.empty @@ -202,9 +157,6 @@ ] -- 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 MS.lookup name medianVarCols of Just c -> c @@ -215,10 +167,7 @@ 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 @@ -229,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) @@ -246,6 +195,5 @@ distinct :: DataFrame -> DataFrame distinct df = selectIndices (VU.map (indices VU.!) (VU.init os)) df where - -- The trailing field stays a wildcard: under -XStrict a named pattern - -- variable would force the (possibly deferred) rowToGroup thunk. + -- 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 --- a/src/DataFrame/Operations/Aggregation/Run.hs +++ b/src/DataFrame/Operations/Aggregation/Run.hs @@ -4,27 +4,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | Execute a recognised aggregation plan ('AggPlan'), producing one result -column (length @nGroups@, canonical group order). - -This is the dispatch layer: it owns the policy of WHICH kernel runs — the dense -direct-indexed one ('DataFrame.Internal.Aggregation.Kernel.Dense') when the -group domain is small enough, otherwise the group-range scatter -('DataFrame.Internal.Aggregation.Kernel.Scatter') — and handles the compound -@max - min@ combine and the holistic grouped median itself. The kernels carry no -policy of their own. A plan only reaches here once 'planAgg' verified the value -columns are clean unboxed Int/Double, so the @error@ branches are unreachable. - -Every reduction takes the Round-5 grouping layout @(valueIndices, offsets)@ so -the parallel kernel can split the group-id range across capabilities with no -cross-worker merge. Each group's rows stay in original-row order within one -worker's range, so results are byte-identical to the sequential path at any @-N@. -(Two exceptions, both deterministic at a fixed @-N@: the direct streaming -Double sum/mean above the small-group cutoff, whose chunked partials change the -float summation order, and the direct var/std, which finalize from -(count, sum, sumsq) partials rather than the gather kernel's row-order Welford -recurrence; see 'DataFrame.Internal.AggKernelDirect'.) --} +-- | Execute a recognised aggregation plan, producing one result column. module DataFrame.Operations.Aggregation.Run ( runPlan, runMomentPlan, @@ -42,6 +22,7 @@ denseMaxMinusMin, denseReduce, ) +import DataFrame.Internal.Aggregation.Kernel.Distinct (countDistinctGroups) import DataFrame.Internal.Aggregation.Kernel.Moments ( Moments (..), momentScatterPar, @@ -66,11 +47,7 @@ import System.IO.Unsafe (unsafePerformIO) import Type.Reflection (typeRep) -{- | Group-domain size at or below which the dense direct-indexed kernel is -chosen; wider domains go to the group-range scatter kernel. A @2^18@-slot -accumulator is replicated per worker there, which is what bounds this. Dispatch -policy, so it lives with the dispatcher rather than inside the kernel. --} +-- | Group-domain size at or below which the dense kernel is chosen. denseThreshold :: Int denseThreshold = 262144 @@ -78,15 +55,7 @@ runPlan gdf rtg nGroups plan = case plan of PlanScatter red name -> scatterColumn red name PlanMaxMinusMin a b -> - {- min/max are order-independent, so both fused single-pass kernels - (the direct streaming one up to 'streamGroupCap', the group-range - gather one above it) are exactly the two gather extrema they replace; - anything they reject (mixed/unclean columns, small inputs) keeps the - two-pass gather path. The streaming cap extends past 'denseThreshold' - for the same reason as the fused multi-reduction pass: on a - direct-grouped frame it works off the eager @rowToGroup@ and skips the - deferred @valueIndices@ placement entirely (measured at 1e6 groups / - 1e8 rows on -N16: stream 1.1s against placement 0.7s + gather 0.7s). -} + -- min/max are order-independent, so the fused kernels are exact. let ca = col a cb = col b direct @@ -96,32 +65,17 @@ 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 - {- The low-cardinality DENSE fast path: for a small dense domain the grouping - layer's @rowToGroup@ already maps row -> group, so we scatter straight off it - (no @valueIndices@ gather). 'denseReduce' admits the order-independent - reductions (exact partial merge, byte-identical to -N1) plus the streaming - Double sum/mean/var/std variants (byte-identical sequential row order at - small group counts; deterministic chunked partials for the large-domain - sum/mean — see "DataFrame.Internal.Aggregation.Kernel.Dense"); anything it - rejects keeps the order-preserving group-range kernel. -} + -- Dense fast path: scatter straight off the eager @rowToGroup@. scatterColumn red name = let c = col name dense | nGroups <= denseThreshold = denseReduce red rtg nGroups c - {- Top-2 selection merges exactly (a multiset selection, no - float adds until finalize), so like the fused passes it streams - off @rowToGroup@ up to 'streamGroupCap': on a direct-grouped - frame that skips the deferred @valueIndices@ placement, which - costs more than the accumulator cache misses it saves - (measured at 1e6 groups / 1e8 rows on -N16: stream 1.0s - against placement 0.7s + gather 0.4s). RTop2Snd shares the - same accumulator machinery and merge-exactness. -} - | red == RTop2Sum || red == RTop2Snd - , nGroups <= streamGroupCap = - denseReduce red rtg nGroups c | otherwise = Nothing in case dense of Just out -> out @@ -132,26 +86,13 @@ Just c -> c Nothing -> error ("runPlan: planned column missing: " ++ T.unpack name) -{- | Run a recognised moment (Q9 regression) plan as one fused scatter over the -two base columns, returning each output name bound to its moment field. The six -sufficient statistics (count, Sx, Sy, Sxx, Syy, Sxy) come out of a single pass, -replacing the three derive passes and six independent scatters of the -per-expression path. The streaming kernel's count is exact; its five Double -sums accumulate per worker chunk in original row order and merge in fixed -worker order — deterministic at a fixed @-N@, float summation order chunk-major -rather than per-group (see 'momentStreamPar'). The gather fallback remains -byte-identical to the sequential kernel at any @-N@. --} +-- | 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 and - the TYPED base columns (no sequential Int->Double materialization, no - valueIndices gather, so a direct-grouped frame never runs its placement - pass). Falls back to the gather kernel above 'streamGroupCap' or on - unclean columns. -} + -- Preferred: the streaming kernel, one fused pass over rowToGroup. streamed = momentStreamPar (rowToGroup gdf) nGroups cx cy ms <- case streamed of Just m -> Just m @@ -171,10 +112,7 @@ 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. --} +-- | @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 = @@ -208,24 +146,7 @@ | Just Refl <- testEquality (typeRep @a) (typeRep @Int) -> VU.map fromIntegral v _ -> error "scatterExtremaDbl" -------------------------------------------------------------------------------- --- Fused holistic median + var/std over one shared gather -------------------------------------------------------------------------------- - -{- | Fused grouped median and variance family over the SAME column: one gather -into the shared scratch buffer serves both. Returns -@(median, variance, stddev)@ columns, or 'Nothing' on a non-numeric column -(the caller keeps the separate per-expression kernels). - -The Welford fold runs over each gathered slice in ascending original-row order -— exactly the recurrence, order and finalize of the var/std kernels -('DataFrame.Internal.AggKernelPar.varPar' and the sequential @varScatter@, -which agree bit-for-bit) — BEFORE the in-place median selection permutes the -slice, and the selection then proceeds exactly as 'groupedMedian'. Both -outputs are therefore bit-identical to the unfused paths; the second full -gather pass is what the fusion saves (measured ~35% off the median+sd pair at -1e4 groups / 1e8 rows on -N16). --} +-- | Fused grouped median and variance family over the same column. runMedianVarFused :: GroupedDataFrame -> Int -> Column -> Maybe (Column, Column, Column) runMedianVarFused gdf nGroups c = do @@ -277,7 +198,7 @@ welford (pos + 1) c' mu' mm' var <- welford s (0 :: Int) 0 0 VUM.unsafeWrite varOut g var - -- Median selection, as in 'medianByGroup' (permutes the slice). + -- Median selection, as in 'medianByGroup'. let slice = VUM.unsafeSlice s len buf !mid = len `div` 2 VA.select slice (mid + 1) @@ -298,20 +219,7 @@ pure (med, var) {-# NOINLINE medianVarByGroup #-} -------------------------------------------------------------------------------- --- Parallel holistic median -------------------------------------------------------------------------------- - -{- | Holistic per-group median over a single unboxed Int/Double column. The -@valueIndices@/@offsets@ layout already places each group's rows in a contiguous -run, so we copy each group's values into a scratch buffer at its own offset and -select the median-rank order statistics in that slice in place (O(len) per -group rather than the O(len log len) full sort) — each group's slice is -independent, so the per-group selections split across capabilities by group -range with no merge. Order statistics are value-determined, so the result is -identical to the sorting variant. Empty groups never occur, so the result is -total. --} +-- | 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" @@ -325,8 +233,7 @@ 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. + -- Each worker sorts and takes medians of its own disjoint group range. parallelBounds_ caps bounds $ \gs ge -> let grp !g | g >= ge = pure () @@ -342,9 +249,7 @@ fill s let slice = VUM.unsafeSlice s len buf !mid = len `div` 2 - {- Move the least mid+1 values to the front (in no - particular order); the two largest of those are the order - statistics at sorted positions mid and mid-1. -} + -- 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) @@ -361,13 +266,7 @@ VU.unsafeFreeze (VUM.unsafeSlice 0 nGroups out) {-# NOINLINE medianByGroup #-} -------------------------------------------------------------------------------- --- Group-range partitioning (shared with the median path) -------------------------------------------------------------------------------- - -{- | Split @[0, nGroups)@ into @caps@ contiguous group ranges balanced by row -count. Identical policy to 'DataFrame.Internal.Aggregation.Kernel.Scatter.groupRangeBounds'. --} +-- | 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) 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 @@ -86,7 +86,7 @@ columnNames, derivingExpressions, empty, - fromNamedColumns, + fromColumns, getColumn, insertColumn, null, @@ -98,94 +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 as a named column. Size mismatches are reconciled by making -the shorter side nullable (`Maybe a`) and padding 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) => @@ -199,34 +127,7 @@ insertVector name xs = insertColumn name (fromVector xs) {-# INLINE insertVector #-} -{- | Adds a foldable collection as a named column. Size mismatches are reconciled by -making the shorter side nullable (`Maybe a`) and padding with `Nothing`. -Do not pass infinite collections: they are fully forced. - -==== __Example__ -@ ->>> :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) => @@ -240,34 +141,7 @@ 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) => @@ -285,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) => @@ -330,9 +178,7 @@ values = xs' ++ replicate (rows - length xs') defaultValue in insertColumn name (fromList values) d -{- | /O(n)/ Like 'insertVector' but takes an already-unboxed vector, -skipping the boxed-to-unboxed conversion 'insertVector' would do for numbers. --} +-- | /O(n)/ Like 'insertVector' but takes an already-unboxed vector. insertUnboxedVector :: forall a. (Columnable a, VU.Unbox a) => @@ -345,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") @@ -412,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) @@ -511,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 @@ -590,81 +294,13 @@ 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' @@ -672,18 +308,7 @@ empty (zip [0 ..] names) -{- | 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)] - -@ --} +-- | 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 @@ -696,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 @@ -723,50 +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 -{- | The dataframe as a row-major matrix of floats, for handing data to ML systems. -'Left' 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' @@ -780,9 +355,7 @@ (fst (dataframeDimensions df)) (\i -> VU.generate (V.length m) (\j -> (m VG.! j) VG.! i)) -{- | The dataframe as a row-major matrix of doubles, for handing data to ML systems. -'Left' 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' @@ -796,9 +369,7 @@ (fst (dataframeDimensions df)) (\i -> VU.generate (V.length m) (\j -> (m VG.! j) VG.! i)) -{- | The dataframe as a row-major matrix of ints, for handing data to ML systems. -'Left' 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) @@ -821,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) @@ -848,9 +408,7 @@ Left e -> throw e Right (TColumn col) -> toVector col -{- | A column as an unboxed vector of 'Int' values. -'Left' if the column cannot be converted to ints (non-numeric or out of 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) @@ -863,9 +421,7 @@ Left e -> throw e Right (TColumn col) -> toIntVector col -{- | A column as an unboxed vector of 'Double' values. -'Left' if the column cannot be converted to doubles (e.g. 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) @@ -881,9 +437,7 @@ Left e -> throw e Right (TColumn col) -> toDoubleVector col -{- | A column as an unboxed vector of 'Float' values. -'Left' if the column cannot be converted to floats (e.g. 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) @@ -919,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 = @@ -950,19 +484,19 @@ 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) -> UExpr (Col @(Maybe a) name) - BoxedColumn Nothing (_ :: V.Vector a) -> UExpr (Col @a name) + 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: " 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 @@ -44,7 +44,7 @@ parseTimeOpt :: DateFormat -> T.Text -> Maybe Day parseTimeOpt dateFormat s = - parseTimeM {- Accept leading/trailing whitespace -} + parseTimeM True defaultTimeLocale dateFormat @@ -127,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) -> @@ -156,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 @@ -14,6 +14,7 @@ -- * Joins join, + joinOn, innerJoin, leftJoin, rightJoin, @@ -21,9 +22,7 @@ -- * Low-level join kernels - {- | Reused by the lazy executor and the parallel-join tests; not a - stable public API (candidates for a future @Join.Internal@ split). - -} + -- | Reused by the lazy executor and tests; not a stable public API. buildHashColumn, buildCompactIndex, hashProbeKernel, @@ -101,26 +100,30 @@ join RIGHT xs right = rightJoin xs right join FULL_OUTER xs right = fullOuterJoin xs right -{- | Build-side row count above which the single-threaded hash join gives way -to sort-merge (single-threaded) or the parallel chunked-probe hash join when -multiple capabilities are available (see 'shouldParallelizeJoin'). --} +-- | Join on @(leftKey, rightKey)@ pairs; the result keeps the left names. +joinOn :: + JoinType -> + [(T.Text, T.Text)] -> + DataFrame -> + DataFrame -> + DataFrame +joinOn jt keyPairs left right = + join + jt + (map fst keyPairs) + left + (renameMany [(r, l) | (l, r) <- keyPairs, l /= r] right) + +-- | Build-side row count above which the hash join switches strategy. joinStrategyThreshold :: Int joinStrategyThreshold = 500_000 -{- | Maps hash values to contiguous slices of original row indices. Indices -live in one hash-sorted vector; an open-addressing linear-probe table (keyed -by hash) records each run's offset and length. @-1@ starts mark empty slots. --} +-- | Maps hash values to contiguous slices of original row indices. data CompactIndex = CompactIndex { ciSortedIndices :: {-# UNPACK #-} !(VU.Vector Int) , ciKeys :: {-# UNPACK #-} !(VU.Vector Int) , ciRuns :: {-# UNPACK #-} !(VU.Vector Int) - {- ^ @(start, len)@ of each run packed 32/32 into one 'Int'; @-1@ = empty - slot. One vector instead of separate starts\/lens halves the table's - memory (a 1e8-row build side is a 2^28-slot table: 2.1GB instead of - 4.3GB) and saves a random cache-line read per lookup hit. - -} + -- ^ @(start, len)@ of each run packed 32/32 into one 'Int'; @-1@ = empty. , ciMask :: {-# UNPACK #-} !Int } @@ -139,12 +142,7 @@ ciRunLen !w = w .&. 0xFFFF_FFFF {-# INLINE ciRunLen #-} -{- | Home slot of a hash: the top @log2 cap@ bits of a Fibonacci multiply. -The row hash's final FxHash step is a multiply, which leaves its LOW bits -poorly diffused (text keys cluster catastrophically: contiguous-pileup chains -in the hundreds), so the table must never index by @h .&. mask@ directly. -@shift@ is @64 - log2 cap@. --} +-- | Home slot of a hash: top bits of a Fibonacci multiply, not low bits. ciSlot :: Int -> Int -> Int ciSlot !shift !h = fromIntegral @@ -156,9 +154,7 @@ ciShiftFor !mask = 64 - popCount mask {-# INLINE ciShiftFor #-} -{- | Look up a hash in the open-addressing table. -Returns @(start, len)@ of the matching run, or @(-1, 0)@ on a miss. --} +-- | 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 (ciSlot shift h) where @@ -176,9 +172,7 @@ 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 @@ -186,10 +180,7 @@ | p > n = p | otherwise = go (p * 2) -{- | Build a compact index from a vector of row hashes: sort by hash, scan for -contiguous runs, insert each run into an open-addressing table. Capacity is -sized for the worst case (every row distinct) so building never resizes. --} +-- | 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 = @@ -232,9 +223,7 @@ | otherwise = j {-# INLINE findGroupEnd #-} -{- | Sort a hash vector, returning sorted hashes and their original indices. -Sorts an index array keyed by hash, avoiding an intermediate pair vector. --} +-- | 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 @@ -287,38 +276,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; --- the probe side is partitioned across cores, so probing the larger maximises --- the parallel win. +-- 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 @@ -347,20 +311,14 @@ 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. Output is bit-for-bit identical to -'hashInnerKernel' (probe-row order preserved). --} +-- | 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 = @@ -370,9 +328,7 @@ in (pf, bf) {-# NOINLINE parInnerKernel #-} -{- | The raw table fields of a 'CompactIndex', in the closure-free shape the -parallel probe kernels consume. --} +-- | The raw table fields of a 'CompactIndex' for the parallel probe kernels. ciProbeTable :: CompactIndex -> ProbeTable ciProbeTable ci = ProbeTable @@ -389,10 +345,7 @@ keyIdxs = validatedKeyColIndices "buildHashColumn" csSet df in D.computeRowHashes keyIdxs df -{- | Probe one batch of rows against a pre-built 'CompactIndex', returning -@(probeExpandedIxs, buildExpandedIxs)@. Unlike 'hashInnerKernel' it neither -builds the index nor guards cross-product size — the caller sizes batches. --} +-- | Probe a batch of rows against a pre-built 'CompactIndex'. hashProbeKernel :: -- | Built once from the full right\/build side. CompactIndex -> @@ -456,14 +409,7 @@ <*> VU.unsafeFreeze (VUM.slice 0 total bv) in (VU.force pFrozen, VU.force bFrozen) -{- | Hash-based inner join kernel: build the index on @buildHashes@, probe with -@probeHashes@, return @(probeExpandedIndices, buildExpandedIndices)@. Grows its -output buffer dynamically rather than pre-allocating the full cross product. --} - -{- | Output-row ceiling for a join kernel; exceeding it signals 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 @@ -534,10 +480,7 @@ <*> VU.unsafeFreeze (VUM.slice 0 total bv) in (VU.force pFrozen, VU.force bFrozen) -{- | Sort-merge inner join kernel: sort both sides by hash, walk in lockstep, -return @(leftExpandedIndices, rightExpandedIndices)@. Grows its output buffer -dynamically so low-cardinality cross products don't OOM. --} +-- | 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 = @@ -623,7 +566,7 @@ <*> VU.unsafeFreeze (VUM.slice 0 total rv) 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 -> @@ -673,32 +616,7 @@ 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" @@ -709,8 +627,7 @@ | D.null left || D.nRows left == 0 = D.empty | otherwise = leftJoinNonEmpty callPoint cs left right --- The right side is always the build side; the left is probed (and drives --- parallelism). Unmatched right rows are marked with a @-1@ sentinel. +-- 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 @@ -736,10 +653,7 @@ in 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 = @@ -747,10 +661,7 @@ in unsafePerformIO (parLeftProbe (ciProbeTable ci) leftHashes) {-# NOINLINE parLeftKernel #-} -{- | Hash-based left join kernel, returning @(leftExpandedIndices, -rightExpandedIndices)@ with @-1@ marking unmatched right rows. Grows its output -buffer dynamically rather than pre-allocating the full cross product. --} +-- | 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 @@ -816,10 +727,7 @@ <$> VU.unsafeFreeze (VUM.slice 0 total lv) <*> VU.unsafeFreeze (VUM.slice 0 total rv) -{- | Sort-merge left join kernel, returning @(leftExpandedIndices, -rightExpandedIndices)@ with a @-1@ sentinel. Grows its output buffer -dynamically so low-cardinality cross products don't OOM. --} +-- | 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 @@ -919,9 +827,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 -> @@ -970,26 +876,7 @@ 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 @@ -1021,10 +908,7 @@ in 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 @@ -1134,9 +1018,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 @@ -1202,9 +1084,7 @@ fill 0 0 0 (,) <$> VU.unsafeFreeze lv <*> VU.unsafeFreeze rv -{- | Assemble the result DataFrame for a full outer join. Both index vectors -use a @-1@ sentinel; key columns are coalesced (first non-null wins). --} +-- | Assemble the result DataFrame for a full outer join. assembleFullOuter :: S.Set T.Text -> DataFrame -> diff --git a/src/DataFrame/Operations/Join/Parallel.hs b/src/DataFrame/Operations/Join/Parallel.hs --- a/src/DataFrame/Operations/Join/Parallel.hs +++ b/src/DataFrame/Operations/Join/Parallel.hs @@ -1,18 +1,7 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE ScopedTypeVariables #-} -{- | -Parallel chunked-probe join kernels. The build side is indexed once into a -shared, read-only 'CompactIndex' (open-addressing, from -"DataFrame.Operations.Join"); the probe side is split into @caps@ /contiguous/ -row ranges and probed in parallel by 'forkFinally' workers (no sparks). Each worker -makes two passes over its range — a count pass to size its slice, then a fill -pass — and writes into the single shared output buffers at a precomputed -prefix-sum offset. Because ranges are contiguous and laid out in range order, -the produced @(probeIxs, buildIxs)@ vectors are /bit-for-bit identical/ to the -sequential 'hashInnerKernel' \/ 'hashLeftKernel': probe rows appear in original -order and, within a probe row, build matches in @ciSortedIndices@ order. --} +-- | Parallel chunked-probe join kernels over a shared 'CompactIndex'. module DataFrame.Operations.Join.Parallel ( ProbeTable (..), parInnerProbe, @@ -31,26 +20,15 @@ import Data.Word (Word64) import DataFrame.Internal.Control.Concurrent (capabilities, forkJoin_) -{- | Below this many probe rows the fork/coordination overhead is not worth it; -the caller uses its sequential 'ST' kernel instead. --} +-- | Below this many probe rows the fork overhead outweighs parallel probing. 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. --} +-- | Below this many build rows the sequential hash probe beats forking. 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). --} +-- | Whether a join should take the parallel probe path. shouldParallelizeJoin :: Int -> Int -> Bool shouldParallelizeJoin probeRows buildRows = probeRows >= parJoinThreshold @@ -58,38 +36,18 @@ && capabilities > 1 {-# NOINLINE shouldParallelizeJoin #-} -{- | Above this many probe rows the probe-side row hashing and table lookups -dominate the join, so partitioning the probe across cores wins even when the -build side is small and cache-resident (the regime 'shouldParallelizeJoin' -deliberately leaves sequential). Sized at 1e6: below it the per-question gain is -swamped by 'forkFinally'/coordination overhead, so small/medium-inner joins stay -sequential (measured). This is the small-build large-probe lever closing the -medium-factor 1e7 join (1e7 probe x ~1e4 build). --} +-- | Probe-row count above which parallel probing wins even for small builds. 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. --} +-- | 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 needed by the probe: the raw -open-addressing table vectors of the @CompactIndex@ (which lives in -"DataFrame.Operations.Join"; passing the fields avoids an import cycle). -Passing concrete vectors instead of a lookup closure keeps the per-row probe -loop free of unknown calls and boxed-tuple allocation — the lookup is inlined -into the count and fill loops. --} +-- | 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) @@ -98,30 +56,21 @@ , ptMask :: {-# UNPACK #-} !Int } -{- | Parallel inner-join probe. @parInnerProbe table probeHashes@ returns -@(probeIxs, buildIxs)@ identical to a sequential probe of the same index. The -build index must already be constructed from the build side. --} +-- | 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. Like 'parInnerProbe' but every probe row emits at -least one output row; unmatched rows carry a @-1@ sentinel in the build column. --} +-- | 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. @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. --} +-- | Shared two-pass parallel probe covering inner and left join semantics. runProbe :: Bool -> ProbeTable -> @@ -136,9 +85,7 @@ !runs = ptRuns pt !mask = ptMask pt !shift = 64 - popCount mask - -- Packed (start,len) run for hash @h@, or -1 on a miss. Home slot is - -- the top log2(cap) bits of a Fibonacci multiply (the row hash's low - -- bits are poorly diffused); must match the build-side ciSlot exactly. + -- Packed (start,len) run for hash @h@, or -1 on a miss. findRun !h = go ( fromIntegral @@ -171,7 +118,7 @@ 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]. + -- Exclusive prefix sum gives per-chunk start offsets; total at [nChunks]. let scan !k !acc | k > nChunks = pure acc | otherwise = do @@ -214,9 +161,6 @@ 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. --} +-- | 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/Merge.hs b/src/DataFrame/Operations/Merge.hs --- a/src/DataFrame/Operations/Merge.hs +++ b/src/DataFrame/Operations/Merge.hs @@ -15,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 = 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 #-} @@ -8,7 +9,9 @@ module DataFrame.Operations.Permutation ( -- * Sorting SortOrder (..), + sortOrderColumns, sortBy, + topBy, -- * Shuffling shuffle, @@ -19,20 +22,35 @@ 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, materializeMerged, + sliceColumn, ) import DataFrame.Internal.Data.PackedText (packedSlice, sliceCmpBytes) import DataFrame.Internal.DataFrame ( @@ -43,6 +61,7 @@ import DataFrame.Internal.Expression (Expr (Col), getColumns) import DataFrame.Operations.Core (dimensions) import DataFrame.Operations.Transformations (derive) +import GHC.Float (castDoubleToWord64) import System.Random (Random (randomR), RandomGen) import Type.Reflection (typeRep) @@ -57,6 +76,12 @@ (==) (Desc _) (Desc _) = True (==) _ _ = False +instance Show SortOrder where + showsPrec d (Asc e) = showParen (d > 10) (showString "Asc " . showsPrec 11 e) + showsPrec d (Desc e) = + showParen (d > 10) (showString "Desc " . showsPrec 11 e) + +-- | Columns a sort order reads. sortOrderColumns :: SortOrder -> [T.Text] sortOrderColumns (Asc e) = getColumns e sortOrderColumns (Desc e) = getColumns e @@ -65,10 +90,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 @@ -91,10 +113,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 -> @@ -112,16 +131,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 @@ -164,6 +255,63 @@ 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) 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 @@ -20,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 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 @@ -53,7 +53,7 @@ DataFrame (..), columnNames, empty, - fromNamedColumns, + fromColumns, getColumn, ) import DataFrame.Internal.Expression @@ -67,24 +67,7 @@ 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 = @@ -144,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 = @@ -218,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 @@ -230,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 @@ -252,13 +235,6 @@ _ -> Nothing {-# INLINE _getColumnAsDouble #-} -optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double -optionalToDoubleVector = - VU.fromList - . V.foldl' - (\acc e -> if isJust e then realToFrac (fromMaybe 0 e) : acc else acc) - [] - -- | Calculates the sum of a given column as a standalone value. sum :: forall a. (Columnable a, Num a) => Expr a -> DataFrame -> a @@ -271,50 +247,14 @@ Just Refl -> VG.sum column Nothing -> 0 Just (PackedText _ _) -> 0 - Just (MergedColumn _ _) -> 0 -- matches the old eager These column (type never Num) + 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 @@ -362,7 +302,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 @@ -23,7 +23,6 @@ -- * Filtering filter, filterBy, - filterWhere, filterJust, filterNothing, filterAllJust, @@ -90,6 +89,7 @@ columnToTextVec, findIndices, hasMissing, + maskIndices, materializeMerged, materializePacked, mkRandom, @@ -182,36 +182,35 @@ where (r, c) = dataframeDimensions d start' = clip start 0 r - -- Clamp both endpoints before subtracting: end - start' on an unclamped - -- end wraps for very negative values and reopens the range. + -- 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 --} -filter :: +-- | O(n * k) Keep rows whose value of the expression satisfies the predicate. +filterBy :: forall a. (Columnable a) => - -- | Column to filter by - Expr a -> -- | Filter condition (a -> Bool) -> + -- | Column to filter by + Expr a -> -- | Dataframe to filter DataFrame -> DataFrame -filter e@(Col filterColumnName) condition df = case getColumn filterColumnName df of +filterBy condition e@(Col filterColumnName) df = case getColumn filterColumnName df of Nothing -> throw $ - ColumnsNotFoundException [filterColumnName] "filter" (M.keys $ columnIndices df) + ColumnsNotFoundException + [filterColumnName] + "filterBy" + (M.keys $ columnIndices df) Just c@(PackedText _ _) -> - filter e condition (insertColumn filterColumnName (materializePacked c) df) + filterBy condition e (insertColumn filterColumnName (materializePacked c) df) Just c@(MergedColumn _ _) -> - filter e condition (insertColumn filterColumnName (materializeMerged c) df) + filterBy condition e (insertColumn filterColumnName (materializeMerged c) df) Just _col@(BoxedColumn bm (column :: V.Vector b)) -> case testEquality (typeRep @a) (typeRep @b) of Just Refl -> filterByVector filterColumnName column condition df @@ -237,7 +236,7 @@ Nothing -> filterByVector filterColumnName column condition df Nothing -> filterByVector filterColumnName column condition df _ -> filterByVector filterColumnName column condition df -filter expr condition df = +filterBy condition expr df = let (TColumn col') = case interpret @a df (normalize expr) of Left e -> throw e @@ -264,7 +263,7 @@ { userType = Right $ typeRep @a , expectedType = Right $ typeRep @b , errorColumnName = Just (T.unpack filterColumnName) - , callingFunctionName = Just "filter" + , callingFunctionName = Just "filterBy" } ) Just Refl -> @@ -276,26 +275,19 @@ , dataframeDimensions = (VG.length ixs, snd (dataframeDimensions df)) } -{- | O(k) a version of filter where the predicate comes first. - -> filterBy even "x" df --} -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 --} -filterWhere :: Expr Bool -> DataFrame -> DataFrame -filterWhere expr df = +-- | O(k) Filter the dataframe with a boolean expression. +filter :: Expr Bool -> DataFrame -> DataFrame +filter 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 @@ -303,10 +295,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 -> @@ -314,57 +303,41 @@ ColumnsNotFoundException [colName] "filterJust" (M.keys $ columnIndices df) Just column | hasMissing column -> case column of BoxedColumn (Just _) (_col :: V.Vector a) -> - filter (Col @(Maybe a) colName) isJust df & apply @(Maybe a) fromJust colName + filterBy isJust (Col @(Maybe a) colName) df & apply @(Maybe a) fromJust colName UnboxedColumn (Just _) (_col :: VU.Vector a) -> - filter (Col @(Maybe a) colName) isJust df & apply @(Maybe a) fromJust colName + filterBy isJust (Col @(Maybe a) colName) df & apply @(Maybe a) fromJust colName c@(PackedText (Just _) _) -> filterJust colName (insertColumn colName (materializePacked c) df) _ -> 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 -> throw $ ColumnsNotFoundException [colName] "filterNothing" (M.keys $ columnIndices df) Just column | hasMissing column -> case column of - BoxedColumn (Just _) (_col :: V.Vector a) -> filter (Col @(Maybe a) colName) isNothing df - UnboxedColumn (Just _) (_col :: VU.Vector a) -> filter (Col @(Maybe a) colName) isNothing df + BoxedColumn (Just _) (_col :: V.Vector a) -> filterBy isNothing (Col @(Maybe a) colName) df + UnboxedColumn (Just _) (_col :: VU.Vector a) -> filterBy isNothing (Col @(Maybe a) colName) df c@(PackedText (Just _) _) -> filterNothing colName (insertColumn colName (materializePacked c) df) _ -> 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 -> @@ -393,42 +366,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 @@ -451,10 +405,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 @@ -464,10 +415,7 @@ where ixs' = VU.fromList ixs -{- | O(n) inverse of select - -> exclude ["Name"] df --} +-- | O(n) Inverse of select. exclude :: [T.Text] -> DataFrame -> @@ -476,15 +424,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 @@ -493,18 +433,10 @@ in df & insertColumn (name cRand) rand - & filterWhere (cRand .>=. Lit (1 - p)) + & filter (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 = @@ -514,23 +446,15 @@ withRand = df & insertColumn (name cRand) rand in ( withRand - & filterWhere (cRand .<=. Lit p) + & filter (cRand .<=. Lit p) & exclude [name cRand] , withRand - & filterWhere + & filter (cRand .>. Lit p) & 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 @@ -539,12 +463,12 @@ withRand = df & insertColumn (name cRand) rand partitionSize = 1 / fromIntegral folds singleFold n d = - d & filterWhere (cRand .>=. Lit (fromIntegral n * partitionSize)) + d & filter (cRand .>=. Lit (fromIntegral n * partitionSize)) go (-1) _ = [] go n d = let d' = singleFold n d - d'' = d & filterWhere (cRand .<. Lit (fromIntegral n * partitionSize)) + d'' = d & filter (cRand .<. Lit (fromIntegral n * partitionSize)) in d' : go (n - 1) d'' in @@ -569,14 +493,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) => @@ -595,14 +512,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) => 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 @@ -75,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) => @@ -93,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 @@ -119,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 @@ -166,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) => @@ -223,9 +210,7 @@ Left e -> throw e Right column' -> insertColumn columnName column' df -{- | Core impute implementation for nullable columns. A column with no null -bitmap cannot hold a missing value, so imputing it is a mistake, not a no-op. --} +-- | Core impute for nullable columns; a column with no nulls is an error. imputeCore :: forall b. (Columnable b) => @@ -246,9 +231,7 @@ (nonNullableColumnError columnName ("Maybe " ++ show (typeRep @b)) col) imputeCore expr _ _ = throw (NonColumnReferenceException (T.pack (show expr))) -{- | 'impute' was handed a column that carries no null bitmap, so there is -nothing it could replace. --} +-- | 'impute' was handed a column with no null bitmap, so nothing to replace. nonNullableColumnError :: T.Text -> String -> Column -> DataFrameException nonNullableColumnError columnName wanted col = TypeMismatchException @() @() @@ -284,9 +267,7 @@ runImpute = errorWithoutStackTrace "impute: unreachable" runImputeWith = errorWithoutStackTrace "impute: unreachable" -{- | Replace all instances of `Nothing` in a column with the given value. -Throws when the column carries no nulls to replace. --} +-- | 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, @@ -49,10 +47,7 @@ import Text.Read import Type.Reflection -{- | How parse failures are surfaced: 'NoSafeRead' throws, 'MaybeRead' yields -@Nothing@ (column wrapped @Maybe a@), 'EitherRead' yields @Left rawText@ -(column wrapped @Either Text a@, preserving the original input). --} +-- | How parse failures are surfaced: throw, 'Nothing', or 'Left' raw text. data SafeReadMode = NoSafeRead | MaybeRead @@ -64,20 +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' for columns without a 'parseSafeOverrides' entry. + -- ^ Default 'SafeReadMode' for columns with no 'parseSafeOverrides' entry. , parseSafeOverrides :: [(T.Text, SafeReadMode)] - {- ^ Per-column overrides taking precedence over 'parseSafe' — e.g. 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 @@ -88,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) @@ -150,10 +139,7 @@ where classify p t = if p t then Nothing else Just t -{- | For 'EitherRead' mode: parse under the chosen assumption into an -@Either Text a@ column. Successful parses become @Right@; failures (including -null/missing cells) become @Left@ carrying the raw input verbatim. --} +-- | 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 @@ -215,10 +201,7 @@ (parseUnboxedColumnWithPred False isNull readBool cols) (handleTextAssumption isNull cols) -{- | Int columns: one fused pass with in-place Int -> Double promotion; a cell -parsing as neither demotes the column to Text. 'readIntStrict' rejects overflow -so a huge integer promotes to 'Double' rather than wrapping. --} +-- | 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 @@ -231,9 +214,7 @@ (parseUnboxedColumnWithPred 0 isNull readDouble cols) (handleTextAssumption isNull cols) -{- | Text columns: no parse, just null-marking. An all-non-null column stays a -plain @V.Vector T.Text@; otherwise it becomes @V.Vector (Maybe T.Text)@. --} +-- | Text columns: no parse, just null-marking. handleTextAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column handleTextAssumption isNull cols | V.any isNull cols = @@ -241,10 +222,7 @@ (V.map (\t -> if isNull t then Nothing else Just t) cols) | otherwise = fromVector cols -{- | Date: single boxed parse pass ('Day' is not unboxable). Bails to -'handleTextAssumption' the moment a non-null cell fails to parse as a 'Day'. -A column with no nulls keeps type 'Day' rather than 'Maybe Day'. --} +-- | 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 = @@ -279,8 +257,7 @@ Nothing -> return Nothing loop 0 False --- Reached only when the sample was all-null: try each concrete type in turn, --- falling back to Text. A column with no nulls keeps type 'Day', not 'Maybe Day'. +-- 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 @@ -298,9 +275,7 @@ | otherwise -> fromVector (V.mapMaybe id vec) Nothing -> handleTextAssumption isNull cols -{- | True for nullish or explicitly-listed missing strings. ('convertNullish' -and 'convertOnlyEmpty' below are kept only for external callers.) --} +-- | True for nullish or explicitly-listed missing strings. isNullishOrMissing :: [T.Text] -> T.Text -> Bool isNullishOrMissing missing v = isNullish v || v `elem` missing @@ -329,9 +304,7 @@ hasSameConstructor Nothing Nothing = True hasSameConstructor _ _ = False -{- | Re-type columns of a 'DataFrame' according to a schema map. @resolveMode@ -maps a column name to its 'SafeReadMode' (typically via 'effectiveSafeRead'). --} +-- | Re-type columns of a 'DataFrame' according to a schema map. parseWithTypes :: (T.Text -> SafeReadMode) -> M.Map T.Text SchemaType -> @@ -398,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 @@ -41,9 +41,7 @@ ) 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 @@ -71,10 +69,7 @@ where colName = T.pack (symbolVal (Proxy @name)) -{- | Retrieve a column coerced to an unboxed 'Int' vector, named by type -application. The column must exist and be numeric — both are compile-time -checks via 'SafeLookup', so this is total (no 'Either', no runtime throw). --} +-- | Retrieve a column as an unboxed 'Int' vector, checked at compile time. columnAsIntVector :: forall name cols a. ( KnownSymbol name @@ -90,7 +85,7 @@ where colName = T.pack (symbolVal (Proxy @name)) --- | Retrieve a column coerced to an unboxed 'Double' vector. See 'columnAsIntVector'. +-- | Retrieve a column as an unboxed 'Double' vector. See 'columnAsIntVector'. columnAsDoubleVector :: forall name cols a. ( KnownSymbol name @@ -107,7 +102,7 @@ where colName = T.pack (symbolVal (Proxy @name)) --- | Retrieve a column coerced to an unboxed 'Float' vector. See 'columnAsIntVector'. +-- | Retrieve a column as an unboxed 'Float' vector. See 'columnAsIntVector'. columnAsFloatVector :: forall name cols a. ( KnownSymbol name @@ -124,9 +119,7 @@ where colName = T.pack (symbolVal (Proxy @name)) -{- | Retrieve a column as an unboxed vector of its own element type. The column -must exist and be unboxable — both compile-time checks, so this is total. --} +-- | Retrieve a column as an unboxed vector of its own element type. columnAsUnboxedVector :: forall name cols a. ( KnownSymbol name @@ -141,22 +134,19 @@ where colName = T.pack (symbolVal (Proxy @name)) -{- | Convert every column to 'Double' and transpose into a row-major matrix. -Total: 'AllColumnsReal' proves at compile time that every column is numeric and -unboxed, so the conversion cannot fail. --} +-- | 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) --- | Convert every column to 'Float' and transpose into a row-major matrix. See 'toDoubleMatrix'. +-- | 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) --- | Convert every column to 'Int' and transpose into a row-major matrix. See 'toDoubleMatrix'. +-- | Every column as 'Int', row-major. See 'toDoubleMatrix'. toIntMatrix :: (AllColumnsReal "toIntMatrix" cols) => TypedDataFrame cols -> V.Vector (VU.Vector Int) 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,50 +34,14 @@ 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) => @@ -86,25 +50,7 @@ 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 --- '[ '(\"region\", Text) --- , '(\"total\", Double) --- , '(\"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 --- a/src/DataFrame/Typed/Apply.hs +++ b/src/DataFrame/Typed/Apply.hs @@ -6,13 +6,11 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} -{- | Typed column transformations: the @apply@\/@derive@ family and -default-valued inserts, plus the horizontal merge @('|||')@. Schema changes are -tracked at the type level (e.g. 'applyColumn' rewrites a column's element type -via 'SetColumnType'). --} +-- | Typed column transformations that track schema changes in the types. module DataFrame.Typed.Apply ( - applyColumn, + apply, + applyInt, + applyDouble, applyMany, applyWhere, applyAtIndex, @@ -50,15 +48,8 @@ ) import DataFrame.Typed.Types (TExpr (..), TypedDataFrame (..)) -{- | Map a function over a column, rewriting its element type from @a@ to @b@. -The schema's entry for @name@ is updated via 'SetColumnType'. - -@ -df' = applyColumn \@\"age\" (show :: Int -> String) df --- the \"age\" column is now String-typed -@ --} -applyColumn :: +-- | Map a function over a column, rewriting its element type from @a@ to @b@. +apply :: forall name a b cols. ( KnownSymbol name , a ~ SafeLookup name cols @@ -69,11 +60,37 @@ (a -> b) -> TypedDataFrame cols -> TypedDataFrame (SetColumnType name b cols) -applyColumn f (TDF df) = unsafeFreeze (D.apply f colName df) +apply 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. +-- | 'apply' specialised to an @Int@ column. +applyInt :: + forall name b cols. + ( KnownSymbol name + , Int ~ SafeLookup name cols + , Columnable b + , AssertPresent name cols + ) => + (Int -> b) -> + TypedDataFrame cols -> + TypedDataFrame (SetColumnType name b cols) +applyInt = apply @name + +-- | 'apply' specialised to a @Double@ column. +applyDouble :: + forall name b cols. + ( KnownSymbol name + , Double ~ SafeLookup name cols + , Columnable b + , AssertPresent name cols + ) => + (Double -> b) -> + TypedDataFrame cols -> + TypedDataFrame (SetColumnType name b cols) +applyDouble = apply @name + +-- | Like 'apply' but returns the error instead of throwing. safeApply :: forall name a b cols. ( KnownSymbol name @@ -89,10 +106,7 @@ where colName = T.pack (symbolVal (Proxy @name)) -{- | Apply a type-preserving function to several columns at once. Every named -column must already share the element type @a@ (enforced by -'AssertAllColumnsHaveType'). --} +-- | Apply a type-preserving function to several columns at once. applyMany :: forall (names :: [Symbol]) a cols. (AllKnownSymbol names, Columnable a, AssertAllColumnsHaveType names a cols) => @@ -101,14 +115,7 @@ TypedDataFrame cols applyMany f (TDF df) = TDF (D.applyMany f (symbolVals @names) df) -{- | Apply a function to a target column only on rows where a condition holds on -a filter column. Both columns are named by type application; the target keeps -its type. - -@ -applyWhere \@\"flagged\" \@\"score\" id (* 2) df -@ --} +-- | Apply a function to a target column on rows passing a condition. applyWhere :: forall filterName targetName a b cols. ( KnownSymbol filterName @@ -145,9 +152,7 @@ where colName = T.pack (symbolVal (Proxy @name)) -{- | Derive a new column and also return a typed reference to it. The returned -expression lives in the extended schema, so it can feed later operations. --} +-- | Derive a new column and also return a typed reference to it. deriveWithExpr :: forall name a cols. ( KnownSymbol name @@ -206,10 +211,7 @@ where colName = T.pack (symbolVal (Proxy @name)) -{- | Horizontal merge: place two DataFrames side by side. The schemas must be -disjoint (no shared column names), enforced by 'AssertDisjoint'; the result -schema is their concatenation. --} +-- | Horizontal merge: place two DataFrames with disjoint schemas side by side. (|||) :: (AssertDisjoint left right) => TypedDataFrame left -> 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 = '[ '(\"age\", Int), '(\"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, @@ -128,7 +105,6 @@ module DataFrame.Typed.Expr.Extra, ) where -import Data.Either (fromRight) import Data.Proxy (Proxy (..)) import Data.String (IsString (..)) import qualified Data.Text as T @@ -139,15 +115,16 @@ import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.Column.Types (Promote, PromoteDiv) import DataFrame.Internal.Expression ( + BinOpId (..), BinUDF (..), Expr (..), - UnUDF (..), + OpRef (..), ) import DataFrame.Internal.Expression.Operators.Nullable ( BaseType, DivWidenOp, NullCmpResult, - NullLift1Op (applyNull1), + NullLift1Op, NullLift1Result, NullLift2Op (applyNull2), NullLift2Result, @@ -172,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 '[(\"salary\", Double)] Double -salary = col \@\"salary\" -@ --} +-- | Create a typed column reference, checked against the schema. col :: forall (name :: Symbol) cols a. ( KnownSymbol name @@ -192,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 @@ -209,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 @@ -221,6 +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) +-- | 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) @@ -259,36 +229,30 @@ -- | Lift a unary function into a typed expression. lift :: (Columnable a, Columnable b) => (a -> b) -> TExpr cols a -> TExpr cols b -lift f (TExpr e) = TExpr (Unary (MkUnaryOp f "unaryUdf" Nothing) e) +lift f (TExpr e) = TExpr (F.lift f e) -- | Lift a binary function into typed expressions. lift2 :: (Columnable a, Columnable b, Columnable c) => (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) +lift2 f (TExpr a) (TExpr b) = TExpr (F.lift2 f 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) -> TExpr cols a -> TExpr cols (NullLift1Result a r) -nullLift f (TExpr e) = TExpr (Unary (MkUnaryOp (applyNull1 f) "nullLift" Nothing) e) +nullLift f (TExpr e) = TExpr (F.nullLift f 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) -> TExpr cols a -> TExpr cols b -> TExpr cols (NullLift2Result a b r) -nullLift2 f (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (applyNull2 f) "nullLift2" Nothing False 0) a b) +nullLift2 f (TExpr a) (TExpr b) = TExpr (F.nullLift2 f a b) infixl 4 .==., ./=., .<., .<=., .>=., .>. infix 4 .==, ./=, .<, .<=, .>=, .> @@ -300,27 +264,27 @@ (.==.) :: (Columnable a, Eq a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool -(.==.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (==) "eq" (Just "==") True 4) a b) +(.==.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (==) (Builtin BEq) (Just "==") True 4) a b) (./=.) :: (Columnable a, Eq a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool -(./=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (/=) "neq" (Just "/=") True 4) a b) +(./=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (/=) (Builtin BNeq) (Just "/=") True 4) a b) (.<.) :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool -(.<.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (<) "lt" (Just "<") False 4) a b) +(.<.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (<) (Builtin BLt) (Just "<") False 4) a b) (.<=.) :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool -(.<=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (<=) "leq" (Just "<=") False 4) a b) +(.<=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (<=) (Builtin BLeq) (Just "<=") False 4) a b) (.>=.) :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool -(.>=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (>=) "geq" (Just ">=") False 4) a b) +(.>=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (>=) (Builtin BGeq) (Just ">=") False 4) a b) (.>.) :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool -(.>.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (>) "gt" (Just ">") False 4) a b) +(.>.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (>) (Builtin BGt) (Just ">") False 4) a b) -- Same-type arithmetic operators @@ -342,43 +306,43 @@ (.^^.) :: (Columnable a, Columnable b, Fractional a, Integral b) => TExpr cols a -> TExpr cols b -> TExpr cols a -(.^^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^^) "pow" (Just ".^^.") False 8) a b) +(.^^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^^) (Builtin BPowFrac) (Just ".^^.") False 8) a b) (.^.) :: (Columnable a, Columnable b, Num a, Integral b) => TExpr cols a -> TExpr cols b -> TExpr cols a -(.^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^) "pow" (Just ".^.") False 8) a b) +(.^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^) (Builtin BPow) (Just ".^.") False 8) a b) (.&&.) :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool -(.&&.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (&&) "and" (Just ".&&.") True 3) a b) +(.&&.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (&&) (Builtin BAnd) (Just ".&&.") True 3) a b) (.||.) :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool -(.||.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (||) "or" (Just ".||.") True 2) a b) +(.||.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (||) (Builtin BOr) (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 -> TExpr cols b -> TExpr cols (NullCmpResult a b) (.&&) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (nullCmpOp (&&)) "nulland" (Just ".&&") True 3) a b) + TExpr + (Binary (MkBinaryOp (nullCmpOp (&&)) (Builtin BNullAnd) (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 -> TExpr cols b -> TExpr cols (NullCmpResult a b) (.||) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (nullCmpOp (||)) "nullor" (Just ".||") True 2) a b) + TExpr + (Binary (MkBinaryOp (nullCmpOp (||)) (Builtin BNullOr) (Just ".||") True 2) a b) 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) @@ -390,7 +354,7 @@ (.+) (TExpr a) (TExpr b) = TExpr ( Binary - (MkBinaryOp (applyNull2 (widenArithOp (+))) "nulladd" (Just "+") True 6) + (MkBinaryOp (applyNull2 (widenArithOp (+))) (Builtin BNullAdd) (Just "+") True 6) a b ) @@ -407,7 +371,7 @@ (.-) (TExpr a) (TExpr b) = TExpr ( Binary - (MkBinaryOp (applyNull2 (widenArithOp (-))) "nullsub" (Just "-") False 6) + (MkBinaryOp (applyNull2 (widenArithOp (-))) (Builtin BNullSub) (Just "-") False 6) a b ) @@ -424,7 +388,7 @@ (.*) (TExpr a) (TExpr b) = TExpr ( Binary - (MkBinaryOp (applyNull2 (widenArithOp (*))) "nullmul" (Just "*") True 7) + (MkBinaryOp (applyNull2 (widenArithOp (*))) (Builtin BNullMul) (Just "*") True 7) a b ) @@ -441,7 +405,7 @@ (./) (TExpr a) (TExpr b) = TExpr ( Binary - (MkBinaryOp (applyNull2 (divArithOp (/))) "nulldiv" (Just "/") False 7) + (MkBinaryOp (applyNull2 (divArithOp (/))) (Builtin BNullDiv) (Just "/") False 7) a b ) @@ -458,7 +422,12 @@ ) => TExpr cols a -> TExpr cols b -> TExpr cols a (.^^) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (applyNull2 (^^)) "pow" (Just ".^^") False 8) a b) + TExpr + ( Binary + (MkBinaryOp (applyNull2 (^^)) (Builtin BPowFrac) (Just ".^^") False 8) + a + b + ) -- | Nullable-aware exponentiation (num base, integral exponent). (.^) :: @@ -472,12 +441,10 @@ ) => TExpr cols a -> TExpr cols b -> TExpr cols a (.^) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (applyNull2 (^)) "pow" (Just ".^") False 8) a b) + TExpr + (Binary (MkBinaryOp (applyNull2 (^)) (Builtin BPow) (Just ".^") False 8) a b) -{- | 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) @@ -488,7 +455,11 @@ TExpr cols (NullCmpResult a b) (.==) (TExpr a) (TExpr b) = TExpr - (Binary (MkBinaryOp (applyNull2 (widenCmpOp (==))) "eq" (Just "==") True 4) a b) + ( Binary + (MkBinaryOp (applyNull2 (widenCmpOp (==))) (Builtin BEq) (Just "==") True 4) + a + b + ) -- | Nullable-aware inequality. Widens numeric operands to their common type. (./=) :: @@ -501,7 +472,11 @@ TExpr cols (NullCmpResult a b) (./=) (TExpr a) (TExpr b) = TExpr - (Binary (MkBinaryOp (applyNull2 (widenCmpOp (/=))) "neq" (Just "/=") True 4) a b) + ( Binary + (MkBinaryOp (applyNull2 (widenCmpOp (/=))) (Builtin BNeq) (Just "/=") True 4) + a + b + ) -- | Nullable-aware less-than. Widens numeric operands to their common type. (.<) :: @@ -514,9 +489,13 @@ TExpr cols (NullCmpResult a b) (.<) (TExpr a) (TExpr b) = TExpr - (Binary (MkBinaryOp (applyNull2 (widenCmpOp (<))) "lt" (Just "<") False 4) a b) + ( Binary + (MkBinaryOp (applyNull2 (widenCmpOp (<))) (Builtin BLt) (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) @@ -527,9 +506,13 @@ TExpr cols (NullCmpResult a b) (.<=) (TExpr a) (TExpr b) = TExpr - (Binary (MkBinaryOp (applyNull2 (widenCmpOp (<=))) "leq" (Just "<=") False 4) a b) + ( Binary + (MkBinaryOp (applyNull2 (widenCmpOp (<=))) (Builtin BLeq) (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) @@ -540,7 +523,11 @@ TExpr cols (NullCmpResult a b) (.>=) (TExpr a) (TExpr b) = TExpr - (Binary (MkBinaryOp (applyNull2 (widenCmpOp (>=))) "geq" (Just ">=") False 4) a b) + ( Binary + (MkBinaryOp (applyNull2 (widenCmpOp (>=))) (Builtin BGeq) (Just ">=") False 4) + a + b + ) -- | Nullable-aware greater-than. Widens numeric operands to their common type. (.>) :: @@ -553,14 +540,14 @@ TExpr cols (NullCmpResult a b) (.>) (TExpr a) (TExpr b) = TExpr - (Binary (MkBinaryOp (applyNull2 (widenCmpOp (>))) "gt" (Just ">") False 4) a b) + ( Binary + (MkBinaryOp (applyNull2 (widenCmpOp (>))) (Builtin BGt) (Just ">") False 4) + a + b + ) not :: TExpr cols Bool -> TExpr cols Bool -not (TExpr e) = TExpr (Unary (MkUnaryOp Prelude.not "not" (Just "!")) e) - -------------------------------------------------------------------------------- --- Aggregation combinators -------------------------------------------------------------------------------- +not (TExpr e) = TExpr (F.not e) sum :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a sum (TExpr e) = TExpr (F.sum e) @@ -594,50 +581,26 @@ TExpr cols a -> TExpr cols a over (TExpr e) = TExpr{unTExpr = F.over (symbolVals @names) e} -------------------------------------------------------------------------------- --- Cast / coercion expressions -------------------------------------------------------------------------------- - castExpr :: forall b cols src. (Columnable b, Columnable src, Read b) => TExpr cols src -> TExpr cols (Maybe b) -castExpr (TExpr e) = - TExpr - (CastExprWith @b @(Maybe b) @src "castExpr" (either (const Nothing) Just) e) +castExpr (TExpr e) = TExpr (F.castExpr @b e) castExprWithDefault :: forall b cols src. (Columnable b, Columnable src, Read b) => b -> TExpr cols src -> TExpr cols b -castExprWithDefault def (TExpr e) = - TExpr - ( CastExprWith @b @b @src - ("castExprWithDefault:" <> T.pack (show def)) - (fromRight def) - e - ) +castExprWithDefault def (TExpr e) = TExpr (F.castExprWithDefault def e) castExprEither :: forall b cols src. (Columnable b, Columnable src, Read b) => TExpr cols src -> TExpr cols (Either T.Text b) -castExprEither (TExpr e) = - TExpr - ( CastExprWith @b @(Either T.Text b) @src - "castExprEither" - (either (Left . T.pack) Right) - e - ) +castExprEither (TExpr e) = TExpr (F.castExprEither @b e) unsafeCastExpr :: forall b cols src. (Columnable b, Columnable src, Read b) => TExpr cols src -> TExpr cols b -unsafeCastExpr (TExpr e) = - TExpr - ( CastExprWith @b @b @src - "unsafeCastExpr" - (fromRight (error "unsafeCastExpr: unexpected Nothing in column")) - e - ) +unsafeCastExpr (TExpr e) = TExpr (F.unsafeCastExpr @b e) toDouble :: (Columnable a, Real a) => TExpr cols a -> TExpr cols Double toDouble (TExpr e) = TExpr (F.toDouble e) diff --git a/src/DataFrame/Typed/Expr/Extra.hs b/src/DataFrame/Typed/Expr/Extra.hs --- a/src/DataFrame/Typed/Expr/Extra.hs +++ b/src/DataFrame/Typed/Expr/Extra.hs @@ -1,13 +1,31 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} -{- | Typed counterparts of the remaining "DataFrame.Functions" expression -combinators not already provided by "DataFrame.Typed.Expr". Each wraps the -untyped combinator 1:1, replacing @Expr@ with @'TExpr' cols@ so column -references stay schema-checked. Re-exported from "DataFrame.Typed.Expr". --} +-- | Typed counterparts of the remaining "DataFrame.Functions" combinators. module DataFrame.Typed.Expr.Extra ( + cast, + castWithDefault, + castEither, + unsafeCast, + eq, + lt, + gt, + leq, + geq, + and, + or, + countDistinct, + topK, + bottomK, + isNull, + isNotNull, + contains, + length, div, mod, mode, @@ -44,15 +62,116 @@ bind, ) where +import Data.Proxy (Proxy (..)) 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 GHC.TypeLits (KnownSymbol, Symbol, symbolVal) +import Prelude hiding (and, div, length, max, min, mod, or) import qualified DataFrame.Functions as F import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.Expression.Operators.Nullable (BaseType) +import DataFrame.Typed.Schema (AssertPresent) import DataFrame.Typed.Types (TExpr (..)) +colName :: forall (name :: Symbol). (KnownSymbol name) => T.Text +colName = T.pack (symbolVal (Proxy @name)) + +-- | Parse a schema-checked column as @a@; failures become 'Nothing'. +cast :: + forall (name :: Symbol) a cols. + (KnownSymbol name, AssertPresent name cols, Columnable a, Read a) => + TExpr cols (Maybe a) +cast = TExpr (F.cast @a (colName @name)) + +-- | Parse a schema-checked column as @a@, with a default for failures. +castWithDefault :: + forall (name :: Symbol) a cols. + (KnownSymbol name, AssertPresent name cols, Columnable a, Read a) => + a -> TExpr cols a +castWithDefault d = TExpr (F.castWithDefault d (colName @name)) + +-- | Parse a schema-checked column as @a@, keeping failures as 'Left'. +castEither :: + forall (name :: Symbol) a cols. + (KnownSymbol name, AssertPresent name cols, Columnable a, Read a) => + TExpr cols (Either T.Text a) +castEither = TExpr (F.castEither @a (colName @name)) + +-- | Parse a schema-checked column as @a@, erroring on failures. +unsafeCast :: + forall (name :: Symbol) a cols. + (KnownSymbol name, AssertPresent name cols, Columnable a, Read a) => + TExpr cols a +unsafeCast = TExpr (F.unsafeCast @a (colName @name)) + +-- | Equality. +eq :: + (Columnable a, Eq a, a ~ BaseType a) => + TExpr cols a -> TExpr cols a -> TExpr cols Bool +eq (TExpr a) (TExpr b) = TExpr (F.eq a b) + +-- | Less than. +lt :: + (Columnable a, Ord a, a ~ BaseType a) => + TExpr cols a -> TExpr cols a -> TExpr cols Bool +lt (TExpr a) (TExpr b) = TExpr (F.lt a b) + +-- | Greater than. +gt :: + (Columnable a, Ord a, a ~ BaseType a) => + TExpr cols a -> TExpr cols a -> TExpr cols Bool +gt (TExpr a) (TExpr b) = TExpr (F.gt a b) + +-- | Less than or equal. +leq :: + (Columnable a, Ord a, a ~ BaseType a) => + TExpr cols a -> TExpr cols a -> TExpr cols Bool +leq (TExpr a) (TExpr b) = TExpr (F.leq a b) + +-- | Greater than or equal. +geq :: + (Columnable a, Ord a, a ~ BaseType a) => + TExpr cols a -> TExpr cols a -> TExpr cols Bool +geq (TExpr a) (TExpr b) = TExpr (F.geq a b) + +-- | Logical and. +and :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool +and (TExpr a) (TExpr b) = TExpr (F.and a b) + +-- | Logical or. +or :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool +or (TExpr a) (TExpr b) = TExpr (F.or a b) + +-- | Number of distinct values (aggregation). +countDistinct :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols Int +countDistinct (TExpr e) = TExpr (F.countDistinct e) + +-- | The k largest values per group (aggregation). +topK :: (Columnable a, Ord a) => Int -> TExpr cols a -> TExpr cols [a] +topK k (TExpr e) = TExpr (F.topK k e) + +-- | The k smallest values per group (aggregation). +bottomK :: (Columnable a, Ord a) => Int -> TExpr cols a -> TExpr cols [a] +bottomK k (TExpr e) = TExpr (F.bottomK k e) + +-- | SQL spelling of 'isNothing'. +isNull :: (Columnable a) => TExpr cols (Maybe a) -> TExpr cols Bool +isNull (TExpr e) = TExpr (F.isNull e) + +-- | SQL spelling of 'isJust'. +isNotNull :: (Columnable a) => TExpr cols (Maybe a) -> TExpr cols Bool +isNotNull (TExpr e) = TExpr (F.isNotNull e) + +-- | Whether the text contains the literal substring. +contains :: T.Text -> TExpr cols T.Text -> TExpr cols Bool +contains needle (TExpr e) = TExpr (F.contains needle e) + +-- | Character count. +length :: TExpr cols T.Text -> TExpr cols Int +length (TExpr e) = TExpr (F.length e) + -- | Integer division. div :: (Integral a, Columnable a) => TExpr cols a -> TExpr cols a -> TExpr cols a @@ -119,7 +238,7 @@ max :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols a max (TExpr a) (TExpr b) = TExpr (F.max a b) --- | Fold a column into a single value with a seed and step function (aggregation). +-- | 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 diff --git a/src/DataFrame/Typed/Functions.hs b/src/DataFrame/Typed/Functions.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Functions.hs @@ -0,0 +1,6 @@ +-- | Typed expression combinators; the typed "DataFrame.Functions". +module DataFrame.Typed.Functions ( + module DataFrame.Typed.Expr, +) where + +import DataFrame.Typed.Expr diff --git a/src/DataFrame/Typed/Join.hs b/src/DataFrame/Typed/Join.hs --- a/src/DataFrame/Typed/Join.hs +++ b/src/DataFrame/Typed/Join.hs @@ -1,4 +1,5 @@ {-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} @@ -6,20 +7,32 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} module DataFrame.Typed.Join ( -- * Typed joins + join, + joinOn, + JoinSchema, + JoinOnSchema, + JoinKeys, + KnownJoinType (..), + Fsts, + SwapPairs, innerJoin, leftJoin, rightJoin, fullOuterJoin, ) where +import Data.Kind (Type) import GHC.TypeLits (Symbol) +import DataFrame.Operations.Join (JoinType (..)) import qualified DataFrame.Operations.Join as DJ import DataFrame.Typed.Freeze (unsafeFreeze) +import DataFrame.Typed.Operations (AllKnownPairs (..)) import DataFrame.Typed.Schema import DataFrame.Typed.Types (TypedDataFrame (..)) @@ -35,7 +48,7 @@ TypedDataFrame right -> TypedDataFrame (InnerJoinSchema keys left right) innerJoin (TDF l) (TDF r) = - unsafeFreeze (DJ.innerJoin keyNames r l) + unsafeFreeze (DJ.innerJoin keyNames l r) where keyNames = symbolVals @keys @@ -83,6 +96,74 @@ TypedDataFrame right -> TypedDataFrame (FullOuterJoinSchema keys left right) fullOuterJoin (TDF l) (TDF r) = - unsafeFreeze (DJ.fullOuterJoin keyNames r l) + unsafeFreeze (DJ.fullOuterJoin keyNames l r) where keyNames = symbolVals @keys + +-- | Keys present on both sides with matching types. +type JoinKeys keys left right = + ( AllKnownSymbol keys + , AssertAllPresent keys left + , AssertAllPresent keys right + , AssertKeyTypesMatch keys left right + ) + +-- | Result schema of a @jt@ join on @keys@. +type family + JoinSchema + (jt :: JoinType) + (keys :: [Symbol]) + (left :: [(Symbol, Type)]) + (right :: [(Symbol, Type)]) :: + [(Symbol, Type)] + where + JoinSchema 'INNER keys left right = InnerJoinSchema keys left right + JoinSchema 'LEFT keys left right = LeftJoinSchema keys left right + JoinSchema 'RIGHT keys left right = RightJoinSchema keys left right + JoinSchema 'FULL_OUTER keys left right = FullOuterJoinSchema keys left right + +-- | Reflect a promoted 'JoinType' to its value. +class KnownJoinType (jt :: JoinType) where + joinTypeVal :: JoinType + +instance KnownJoinType 'INNER where joinTypeVal = INNER +instance KnownJoinType 'LEFT where joinTypeVal = LEFT +instance KnownJoinType 'RIGHT where joinTypeVal = RIGHT +instance KnownJoinType 'FULL_OUTER where joinTypeVal = FULL_OUTER + +-- | Typed counterpart of untyped @join@: @join \@'INNER \@'["key"] l r@. +join :: + forall (jt :: JoinType) (keys :: [Symbol]) left right. + (KnownJoinType jt, JoinKeys keys left right) => + TypedDataFrame left -> + TypedDataFrame right -> + TypedDataFrame (JoinSchema jt keys left right) +join (TDF l) (TDF r) = + unsafeFreeze (DJ.join (joinTypeVal @jt) (symbolVals @keys) l r) + +-- | Left names of @(leftKey, rightKey)@ pairs. +type family Fsts (pairs :: [(Symbol, Symbol)]) :: [Symbol] where + Fsts '[] = '[] + Fsts ('(l, _) ': rest) = l ': Fsts rest + +-- | Flip each pair, turning @(left, right)@ into a right-to-left rename. +type family SwapPairs (pairs :: [(Symbol, Symbol)]) :: [(Symbol, Symbol)] where + SwapPairs '[] = '[] + SwapPairs ('(l, r) ': rest) = '(r, l) ': SwapPairs rest + +-- | Result schema of 'joinOn': right keys take their left names first. +type JoinOnSchema jt pairs left right = + JoinSchema jt (Fsts pairs) left (RenameManyInSchema (SwapPairs pairs) right) + +-- | Join on key pairs: @joinOn \@'INNER \@'[ '("id", "uid")] l r@. +joinOn :: + forall (jt :: JoinType) (pairs :: [(Symbol, Symbol)]) left right. + ( KnownJoinType jt + , AllKnownPairs pairs + , JoinKeys (Fsts pairs) left (RenameManyInSchema (SwapPairs pairs) right) + ) => + TypedDataFrame left -> + TypedDataFrame right -> + TypedDataFrame (JoinOnSchema jt pairs left right) +joinOn (TDF l) (TDF r) = + unsafeFreeze (DJ.joinOn (joinTypeVal @jt) (pairVals @pairs) l r) 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 @@ -11,7 +11,6 @@ module DataFrame.Typed.Operations ( -- * Schema-preserving operations - filterWhere, filter, filterBy, filterAllJust, @@ -19,6 +18,7 @@ filterNothing, filterAllNothing, sortBy, + topBy, take, takeLast, drop, @@ -31,6 +31,7 @@ -- * Schema-modifying operations derive, + deriveMany, impute, select, exclude, @@ -47,6 +48,11 @@ valueCounts, valueProportions, + -- * Construction and folds + empty, + null, + fold, + -- * Metadata dimensions, nRows, @@ -61,14 +67,19 @@ intersect, difference, symmetricDifference, + + -- * Helpers + AllKnownPairs (..), + toUntypedSort, ) where +import qualified Data.List as L import Data.Proxy (Proxy (..)) import qualified Data.Text as T import qualified Data.Vector as V import GHC.TypeLits (KnownSymbol, Symbol, symbolVal) import System.Random (RandomGen) -import Prelude hiding (drop, filter, take) +import Prelude hiding (drop, filter, null, take) import qualified DataFrame.Functions as DF import DataFrame.Internal.Column (Columnable) @@ -84,44 +95,29 @@ import DataFrame.Typed.Freeze (unsafeFreeze) import DataFrame.Typed.Schema -import DataFrame.Typed.Types (TExpr (..), TSortOrder (..), TypedDataFrame (..)) +import DataFrame.Typed.Types ( + TAgg (..), + TExpr (..), + TSortOrder (..), + TypedDataFrame (..), + taggToNamedExprs, + ) -{- | Filter rows where a boolean expression evaluates to True. -The expression is validated against the schema at compile time. --} -filterWhere :: TExpr cols Bool -> TypedDataFrame cols -> TypedDataFrame cols -filterWhere (TExpr expr) (TDF df) = TDF (D.filterWhere expr df) +-- | Filter rows where a boolean expression evaluates to True. +filter :: TExpr cols Bool -> TypedDataFrame cols -> TypedDataFrame cols +filter (TExpr expr) (TDF df) = TDF (D.filter expr df) -- | Filter rows by applying a predicate to a typed expression. -filter :: - (Columnable a) => - TExpr cols a -> (a -> Bool) -> TypedDataFrame cols -> TypedDataFrame cols -filter (TExpr expr) pred' (TDF df) = TDF (D.filter expr pred' df) - --- | Filter rows by a predicate on a column expression (flipped argument order). filterBy :: (Columnable a) => (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 '[ '(\"x\", Maybe Double), '(\"y\", Int)] -filterAllJust df :: TDF '[ '(\"x\", Double), '(\"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 +128,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,22 +139,24 @@ where colName = T.pack (symbolVal (Proxy @name)) -{- | Keep only rows where every nullable column has Nothing. -Schema is preserved. --} +-- | 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. -Sort orders reference columns that are validated against the schema. --} +-- | 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 - toUntypedSort :: TSortOrder cols -> D.SortOrder - toUntypedSort (Asc (TExpr e)) = D.Asc e - toUntypedSort (Desc (TExpr e)) = D.Desc e +-- | Rows @[offset, offset + count)@ of the frame sorted by the given orders. +topBy :: + Int -> Int -> [TSortOrder cols] -> TypedDataFrame cols -> TypedDataFrame cols +topBy offset count ords (TDF df) = + TDF (D.topBy offset count (map toUntypedSort ords) df) + +toUntypedSort :: TSortOrder cols -> D.SortOrder +toUntypedSort (Asc (TExpr e)) = D.Asc e +toUntypedSort (Desc (TExpr e)) = D.Desc e + -- | Take the first @n@ rows. take :: Int -> TypedDataFrame cols -> TypedDataFrame cols take n (TDF df) = TDF (D.take n df) @@ -198,15 +194,7 @@ shuffle :: (RandomGen g) => g -> TypedDataFrame cols -> TypedDataFrame cols shuffle g (TDF df) = TDF (D.shuffle g df) -{- | 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 ('(\"total\", Double ': originalCols)) -@ --} +-- | Derive a new column from a typed expression; its name must be absent. derive :: forall name a cols. ( KnownSymbol name @@ -220,6 +208,16 @@ where colName = T.pack (symbolVal (Proxy @name)) +-- | Add several columns at once; build them with 'DataFrame.Typed.as'. +deriveMany :: + forall new cols. + (AssertAllAbsent (Reverse new) cols) => + (TAgg '[] cols '[] -> TAgg '[] cols new) -> + TypedDataFrame cols -> + TypedDataFrame (Append cols (Reverse new)) +deriveMany build (TDF df) = + unsafeFreeze (D.deriveMany (taggToNamedExprs (build TAggNil)) df) + impute :: forall name a cols. ( KnownSymbol name @@ -331,9 +329,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 @@ -358,9 +354,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) @@ -368,6 +362,22 @@ symmetricDifference :: TypedDataFrame cols -> TypedDataFrame cols -> TypedDataFrame cols symmetricDifference (TDF a) (TDF b) = TDF (DS.symmetricDifference a b) + +-- | The frame with no rows and no columns. +empty :: TypedDataFrame '[] +empty = TDF D.empty + +-- | True when the frame has no rows or no columns. +null :: TypedDataFrame cols -> Bool +null (TDF df) = D.null df + +-- | Thread the frame through a schema-preserving step for each element. +fold :: + (a -> TypedDataFrame cols -> TypedDataFrame cols) -> + [a] -> + TypedDataFrame cols -> + TypedDataFrame cols +fold f xs df = L.foldl' (flip f) df xs dimensions :: TypedDataFrame cols -> (Int, Int) dimensions (TDF df) = D.dimensions df diff --git a/src/DataFrame/Typed/Sampling.hs b/src/DataFrame/Typed/Sampling.hs --- a/src/DataFrame/Typed/Sampling.hs +++ b/src/DataFrame/Typed/Sampling.hs @@ -1,9 +1,6 @@ {-# LANGUAGE FlexibleContexts #-} -{- | Typed sampling and splitting. All operations are schema-preserving: they -change which rows are present, never the columns, so every result reuses the -input schema @cols@. --} +-- | Typed sampling and splitting; all operations preserve the schema. module DataFrame.Typed.Sampling ( randomSplit, kFolds, @@ -30,10 +27,7 @@ (RandomGen g) => g -> Int -> TypedDataFrame cols -> [TypedDataFrame cols] kFolds g k (TDF df) = map TDF (D.kFolds g k df) -{- | Select rows by index. -| This may fail if the indices are out of bounds; -| use with caution or use 'filter' to select rows by a predicate instead. --} +-- | 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) diff --git a/src/DataFrame/Typed/Statistics.hs b/src/DataFrame/Typed/Statistics.hs --- a/src/DataFrame/Typed/Statistics.hs +++ b/src/DataFrame/Typed/Statistics.hs @@ -6,19 +6,7 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} -{- | Typed statistical reducers over a 'TypedDataFrame'. - -These mirror the untyped reducers in "DataFrame.Operations.Statistics", taking -a schema-checked 'TExpr' instead of a raw @Expr@. The names (@sum@, @mean@, -@median@, …) deliberately collide with the aggregation-expression combinators -in "DataFrame.Typed.Expr", so this module is meant to be imported qualified: - -@ -import qualified DataFrame.Typed.Statistics as TS - -avg = TS.mean (col \@\"salary\") employees -@ --} +-- | Typed statistical reducers over a 'TypedDataFrame'; import qualified. module DataFrame.Typed.Statistics ( mean, meanMaybe, @@ -117,14 +105,7 @@ sum :: (Columnable a, Num a) => TExpr cols a -> TypedDataFrame cols -> a sum (TExpr e) (TDF df) = Stats.sum e df -{- | Pearson's correlation coefficient between two columns, named by type -application. Both columns must exist in the schema and be numeric — these are -checked at compile time via 'SafeLookup' on each name. - -@ -TS.correlation \@\"height\" \@\"weight\" people -@ --} +-- | Pearson's correlation coefficient between two named numeric columns. correlation :: forall c1 c2 a b cols. ( KnownSymbol c1 @@ -147,17 +128,12 @@ (T.pack (symbolVal (Proxy @c2))) df -{- | Frequency table for a column. The result schema is data-dependent -(one column per distinct value), so an untyped 'D.DataFrame' is returned. --} +-- | 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 (e.g. the mean). -Schema-preserving: the imputed column keeps its type-level @Maybe@ even though -its runtime values are now fully populated. --} +-- | Impute missing values in a column using a derived scalar. imputeWith :: (ImputeOp a, Columnable (BaseType a)) => (TExpr cols (BaseType a) -> TExpr cols (BaseType a)) -> @@ -167,14 +143,10 @@ imputeWith f (TExpr e) (TDF df) = TDF (Stats.imputeWith (unTExpr . f . TExpr) e df) -{- | Descriptive statistics of the numeric columns. Returns an untyped -'D.DataFrame' (the result is a fixed set of statistic rows, not the input schema). --} +-- | 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 (non-null\/null counts, unique values, type). Returns an -untyped 'D.DataFrame'. --} +-- | Per-column summary of counts, unique values and type, untyped. describeColumns :: TypedDataFrame cols -> D.DataFrame describeColumns (TDF df) = Core.describeColumns df