{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE IncoherentInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module DataFrame.Functions (
    module DataFrame.Functions,
    module DataFrame.Operators,
    add,
    sub,
    mult,
    divide,
    prettyPrint,
) where

import DataFrame.Internal.Column
import DataFrame.Internal.Expression
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.Text as T
import Data.Time
import Data.Type.Equality (testEquality, type (:~:) (Refl))
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU

import DataFrame.Internal.Nullable (
    BaseType,
    NullLift1Op (applyNull1),
    NullLift1Result,
    NullLift2Op (applyNull2),
    NullLift2Result,
 )
import DataFrame.Operators
import Text.Regex.TDFA
import Type.Reflection (typeRep)
import Prelude hiding (maximum, minimum)
import Prelude as P

lift :: (Columnable a, Columnable b) => (a -> b) -> Expr a -> Expr b
lift :: forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Expr a -> Expr b
lift a -> b
f =
    UnUDF a b -> Expr a -> Expr b
forall (op :: * -> * -> *) a b.
(UnaryOp op, Columnable a, Columnable b) =>
op b a -> Expr b -> Expr a
Unary (MkUnaryOp{unaryFn :: a -> b
unaryFn = a -> b
f, unaryName :: Text
unaryName = Text
"unaryUdf", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})

lift2 ::
    (Columnable c, Columnable b, Columnable a) =>
    (c -> b -> a) ->
    Expr c ->
    Expr b ->
    Expr a
lift2 :: forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a) -> Expr c -> Expr b -> Expr a
lift2 c -> b -> a
f =
    BinUDF c b a -> Expr c -> Expr b -> Expr a
forall (op :: * -> * -> * -> *) c b a.
(BinaryOp op, Columnable c, Columnable b, Columnable a) =>
op c b a -> Expr c -> Expr b -> Expr a
Binary
        ( MkBinaryOp
            { binaryFn :: c -> b -> a
binaryFn = c -> b -> a
f
            , binaryName :: Text
binaryName = Text
"binaryUdf"
            , binarySymbol :: Maybe Text
binarySymbol = Maybe Text
forall a. Maybe a
Nothing
            , binaryCommutative :: Bool
binaryCommutative = Bool
False
            , binaryPrecedence :: Int
binaryPrecedence = Int
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.
-}
nullLift ::
    (NullLift1Op a r (NullLift1Result a r), Columnable (NullLift1Result a r)) =>
    (BaseType a -> r) ->
    Expr a ->
    Expr (NullLift1Result a r)
nullLift :: forall a r.
(NullLift1Op a r (NullLift1Result a r),
 Columnable (NullLift1Result a r)) =>
(BaseType a -> r) -> Expr a -> Expr (NullLift1Result a r)
nullLift BaseType a -> r
f =
    UnUDF a (NullLift1Result a r)
-> Expr a -> Expr (NullLift1Result a r)
forall (op :: * -> * -> *) a b.
(UnaryOp op, Columnable a, Columnable b) =>
op b a -> Expr b -> Expr a
Unary
        (MkUnaryOp{unaryFn :: a -> NullLift1Result a r
unaryFn = (BaseType a -> r) -> a -> NullLift1Result a r
forall a r c. NullLift1Op a r c => (BaseType a -> r) -> a -> c
applyNull1 BaseType a -> r
f, unaryName :: Text
unaryName = Text
"nullLift", unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing})

{- | Lift a binary function over nullable or non-nullable column expressions.
Any 'Nothing' operand short-circuits to 'Nothing' in the result.

The return type is inferred via 'NullLift2Result': no annotation needed.
-}
nullLift2 ::
    (NullLift2Op a b r (NullLift2Result a b r), Columnable (NullLift2Result a b r)) =>
    (BaseType a -> BaseType b -> r) ->
    Expr a ->
    Expr b ->
    Expr (NullLift2Result a b r)
nullLift2 :: forall a b r.
(NullLift2Op a b r (NullLift2Result a b r),
 Columnable (NullLift2Result a b r)) =>
(BaseType a -> BaseType b -> r)
-> Expr a -> Expr b -> Expr (NullLift2Result a b r)
nullLift2 BaseType a -> BaseType b -> r
f =
    BinUDF a b (NullLift2Result a b r)
-> Expr a -> Expr b -> Expr (NullLift2Result a b r)
forall (op :: * -> * -> * -> *) c b a.
(BinaryOp op, Columnable c, Columnable b, Columnable a) =>
op c b a -> Expr c -> Expr b -> Expr a
Binary
        ( MkBinaryOp
            { binaryFn :: a -> b -> NullLift2Result a b r
binaryFn = (BaseType a -> BaseType b -> r) -> a -> b -> NullLift2Result a b r
forall a b r c.
NullLift2Op a b r c =>
(BaseType a -> BaseType b -> r) -> a -> b -> c
applyNull2 BaseType a -> BaseType b -> r
f
            , binaryName :: Text
binaryName = Text
"nullLift2"
            , binarySymbol :: Maybe Text
binarySymbol = Maybe Text
forall a. Maybe a
Nothing
            , binaryCommutative :: Bool
binaryCommutative = Bool
False
            , binaryPrecedence :: Int
binaryPrecedence = Int
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'.
-}
cast :: forall a. (Columnable a, Read a) => T.Text -> Expr (Maybe a)
cast :: forall a. (Columnable a, Read a) => Text -> Expr (Maybe a)
cast Text
colName = Text -> Text -> (Either [Char] a -> Maybe a) -> Expr (Maybe a)
forall a1 a.
(Columnable a1, Columnable a, Read a1) =>
Text -> Text -> (Either [Char] a1 -> a) -> Expr a
CastWith Text
colName Text
"cast" (([Char] -> Maybe a) -> (a -> Maybe a) -> Either [Char] a -> Maybe a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe a -> [Char] -> Maybe a
forall a b. a -> b -> a
const Maybe a
forall a. Maybe a
Nothing) a -> Maybe a
forall a. a -> Maybe a
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'.
-}
castWithDefault :: forall a. (Columnable a, Read a) => a -> T.Text -> Expr a
castWithDefault :: forall a. (Columnable a, Read a) => a -> Text -> Expr a
castWithDefault a
def Text
colName =
    Text -> Text -> (Either [Char] a -> a) -> Expr a
forall a1 a.
(Columnable a1, Columnable a, Read a1) =>
Text -> Text -> (Either [Char] a1 -> a) -> Expr a
CastWith Text
colName (Text
"castWithDefault:" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (a -> [Char]
forall a. Show a => a -> [Char]
show a
def)) (a -> Either [Char] a -> a
forall b a. b -> Either a b -> b
fromRight a
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\"@.
-}
castEither ::
    forall a. (Columnable a, Read a) => T.Text -> Expr (Either T.Text a)
castEither :: forall a. (Columnable a, Read a) => Text -> Expr (Either Text a)
castEither Text
colName = Text
-> Text
-> (Either [Char] a -> Either Text a)
-> Expr (Either Text a)
forall a1 a.
(Columnable a1, Columnable a, Read a1) =>
Text -> Text -> (Either [Char] a1 -> a) -> Expr a
CastWith Text
colName Text
"castEither" (([Char] -> Either Text a)
-> (a -> Either Text a) -> Either [Char] a -> Either Text a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Text -> Either Text a
forall a b. a -> Either a b
Left (Text -> Either Text a)
-> ([Char] -> Text) -> [Char] -> Either Text a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Text
T.pack) a -> Either Text a
forall a b. b -> Either a b
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.
-}
unsafeCast :: forall a. (Columnable a, Read a) => T.Text -> Expr a
unsafeCast :: forall a. (Columnable a, Read a) => Text -> Expr a
unsafeCast Text
colName =
    Text -> Text -> (Either [Char] a -> a) -> Expr a
forall a1 a.
(Columnable a1, Columnable a, Read a1) =>
Text -> Text -> (Either [Char] a1 -> a) -> Expr a
CastWith
        Text
colName
        Text
"unsafeCast"
        (a -> Either [Char] a -> a
forall b a. b -> Either a b -> b
fromRight ([Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"unsafeCast: unexpected Nothing in column"))

castExpr ::
    forall b src.
    (Columnable b, Columnable src, Read b) =>
    Expr src ->
    Expr (Maybe b)
castExpr :: forall b src.
(Columnable b, Columnable src, Read b) =>
Expr src -> Expr (Maybe b)
castExpr = forall a1 a src.
(Columnable a1, Columnable a, Columnable src, Read a1) =>
Text -> (Either [Char] a1 -> a) -> Expr src -> Expr a
CastExprWith @b @(Maybe b) @src Text
"castExpr" (([Char] -> Maybe b) -> (b -> Maybe b) -> Either [Char] b -> Maybe b
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe b -> [Char] -> Maybe b
forall a b. a -> b -> a
const Maybe b
forall a. Maybe a
Nothing) b -> Maybe b
forall a. a -> Maybe a
Just)

castExprWithDefault ::
    forall b src. (Columnable b, Columnable src, Read b) => b -> Expr src -> Expr b
castExprWithDefault :: forall b src.
(Columnable b, Columnable src, Read b) =>
b -> Expr src -> Expr b
castExprWithDefault b
def =
    forall a1 a src.
(Columnable a1, Columnable a, Columnable src, Read a1) =>
Text -> (Either [Char] a1 -> a) -> Expr src -> Expr a
CastExprWith @b @b @src
        (Text
"castExprWithDefault:" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (b -> [Char]
forall a. Show a => a -> [Char]
show b
def))
        (b -> Either [Char] b -> b
forall b a. b -> Either a b -> b
fromRight b
def)

castExprEither ::
    forall b src.
    (Columnable b, Columnable src, Read b) =>
    Expr src ->
    Expr (Either T.Text b)
castExprEither :: forall b src.
(Columnable b, Columnable src, Read b) =>
Expr src -> Expr (Either Text b)
castExprEither =
    forall a1 a src.
(Columnable a1, Columnable a, Columnable src, Read a1) =>
Text -> (Either [Char] a1 -> a) -> Expr src -> Expr a
CastExprWith @b @(Either T.Text b) @src
        Text
"castExprEither"
        (([Char] -> Either Text b)
-> (b -> Either Text b) -> Either [Char] b -> Either Text b
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Text -> Either Text b
forall a b. a -> Either a b
Left (Text -> Either Text b)
-> ([Char] -> Text) -> [Char] -> Either Text b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Text
T.pack) b -> Either Text b
forall a b. b -> Either a b
Right)

unsafeCastExpr ::
    forall b src. (Columnable b, Columnable src, Read b) => Expr src -> Expr b
unsafeCastExpr :: forall b src.
(Columnable b, Columnable src, Read b) =>
Expr src -> Expr b
unsafeCastExpr =
    forall a1 a src.
(Columnable a1, Columnable a, Columnable src, Read a1) =>
Text -> (Either [Char] a1 -> a) -> Expr src -> Expr a
CastExprWith @b @b @src
        Text
"unsafeCastExpr"
        (b -> Either [Char] b -> b
forall b a. b -> Either a b -> b
fromRight ([Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"unsafeCastExpr: unexpected Nothing in column"))

toDouble :: (Columnable a, Real a) => Expr a -> Expr Double
toDouble :: forall a. (Columnable a, Real a) => Expr a -> Expr Double
toDouble =
    UnUDF a Double -> Expr a -> Expr Double
forall (op :: * -> * -> *) a b.
(UnaryOp op, Columnable a, Columnable b) =>
op b a -> Expr b -> Expr a
Unary
        ( MkUnaryOp
            { unaryFn :: a -> Double
unaryFn = a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac
            , unaryName :: Text
unaryName = Text
"toDouble"
            , unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing
            }
        )

infix 8 `div`
div :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a
div :: forall a. (Integral a, Columnable a) => Expr a -> Expr a -> Expr a
div = (a -> a -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr a -> Expr a -> Expr a
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr c -> Expr b -> Expr a
lift2Decorated a -> a -> a
forall a. Integral a => a -> a -> a
Prelude.div Text
"div" (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"//") Bool
False Int
7

mod :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a
mod :: forall a. (Integral a, Columnable a) => Expr a -> Expr a -> Expr a
mod = (a -> a -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr a -> Expr a -> Expr a
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr c -> Expr b -> Expr a
lift2Decorated a -> a -> a
forall a. Integral a => a -> a -> a
Prelude.mod Text
"mod" Maybe Text
forall a. Maybe a
Nothing Bool
False Int
7

eq :: (Columnable a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool
eq :: forall a.
(Columnable a, Eq a, a ~ BaseType a) =>
Expr a -> Expr a -> Expr Bool
eq = (a -> a -> Bool)
-> Text
-> Maybe Text
-> Bool
-> Int
-> Expr a
-> Expr a
-> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr c -> Expr b -> Expr a
lift2Decorated a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==) Text
"eq" (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"==") Bool
True Int
4

lt :: (Columnable a, Ord a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool
lt :: forall a.
(Columnable a, Ord a, a ~ BaseType a) =>
Expr a -> Expr a -> Expr Bool
lt = (a -> a -> Bool)
-> Text
-> Maybe Text
-> Bool
-> Int
-> Expr a
-> Expr a
-> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr c -> Expr b -> Expr a
lift2Decorated a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<) Text
"lt" (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"<") Bool
False Int
4

gt :: (Columnable a, Ord a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool
gt :: forall a.
(Columnable a, Ord a, a ~ BaseType a) =>
Expr a -> Expr a -> Expr Bool
gt = (a -> a -> Bool)
-> Text
-> Maybe Text
-> Bool
-> Int
-> Expr a
-> Expr a
-> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr c -> Expr b -> Expr a
lift2Decorated a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>) Text
"gt" (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
">") Bool
False Int
4

leq ::
    (Columnable a, Ord a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool
leq :: forall a.
(Columnable a, Ord a, Eq a, a ~ BaseType a) =>
Expr a -> Expr a -> Expr Bool
leq = (a -> a -> Bool)
-> Text
-> Maybe Text
-> Bool
-> Int
-> Expr a
-> Expr a
-> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr c -> Expr b -> Expr a
lift2Decorated a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<=) Text
"leq" (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"<=") Bool
False Int
4

geq ::
    (Columnable a, Ord a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool
geq :: forall a.
(Columnable a, Ord a, Eq a, a ~ BaseType a) =>
Expr a -> Expr a -> Expr Bool
geq = (a -> a -> Bool)
-> Text
-> Maybe Text
-> Bool
-> Int
-> Expr a
-> Expr a
-> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr c -> Expr b -> Expr a
lift2Decorated a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>=) Text
"geq" (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
">=") Bool
False Int
4

and :: Expr Bool -> Expr Bool -> Expr Bool
and :: Expr Bool -> Expr Bool -> Expr Bool
and = Expr Bool -> Expr Bool -> Expr Bool
Expr Bool -> Expr Bool -> Expr (NullCmpResult Bool Bool)
forall a b.
(NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
(.&&)

or :: Expr Bool -> Expr Bool -> Expr Bool
or :: Expr Bool -> Expr Bool -> Expr Bool
or = Expr Bool -> Expr Bool -> Expr Bool
Expr Bool -> Expr Bool -> Expr (NullCmpResult Bool Bool)
forall a b.
(NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
(.||)

not :: Expr Bool -> Expr Bool
not :: Expr Bool -> Expr Bool
not =
    UnUDF Bool Bool -> Expr Bool -> Expr Bool
forall (op :: * -> * -> *) a b.
(UnaryOp op, Columnable a, Columnable b) =>
op b a -> Expr b -> Expr a
Unary
        (MkUnaryOp{unaryFn :: Bool -> Bool
unaryFn = Bool -> Bool
Prelude.not, unaryName :: Text
unaryName = Text
"not", unarySymbol :: Maybe Text
unarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"~"})

count :: (Columnable a) => Expr a -> Expr Int
count :: forall a. Columnable a => Expr a -> Expr Int
count = AggStrategy Int a -> Expr a -> Expr Int
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg (Text
-> Int
-> (Int -> a -> Int)
-> (Int -> Int -> Int)
-> (Int -> Int)
-> AggStrategy Int a
forall acc b a.
Columnable acc =>
Text
-> acc
-> (acc -> b -> acc)
-> (acc -> acc -> acc)
-> (acc -> a)
-> AggStrategy a b
MergeAgg Text
"count" (Int
0 :: Int) (\Int
c a
_ -> Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int -> Int
forall a. a -> a
id)
{-# SPECIALIZE count :: Expr Double -> Expr Int #-}
{-# SPECIALIZE count :: Expr Float -> Expr Int #-}
{-# SPECIALIZE count :: Expr Int -> Expr Int #-}
{-# SPECIALIZE count :: Expr Int8 -> Expr Int #-}
{-# SPECIALIZE count :: Expr Int16 -> Expr Int #-}
{-# SPECIALIZE count :: Expr Int32 -> Expr Int #-}
{-# SPECIALIZE count :: Expr Int64 -> Expr Int #-}
{-# INLINEABLE count #-}

-- | Row count, the equivalent of SQL's @COUNT(*)@.
countAll :: Expr Int
countAll :: Expr Int
countAll = Expr Int -> Expr Int
forall a. Columnable a => Expr a -> Expr Int
count (Int -> Expr Int
forall a. Columnable a => a -> Expr a
Lit (Int
0 :: Int))
{-# INLINE countAll #-}

collect :: (Columnable a) => Expr a -> Expr [a]
collect :: forall a. Columnable a => Expr a -> Expr [a]
collect = AggStrategy [a] a -> Expr a -> Expr [a]
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg (Text -> Maybe [a] -> ([a] -> a -> [a]) -> AggStrategy [a] a
forall a b. Text -> Maybe a -> (a -> b -> a) -> AggStrategy a b
FoldAgg Text
"collect" ([a] -> Maybe [a]
forall a. a -> Maybe a
Just []) ((a -> [a] -> [a]) -> [a] -> a -> [a]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)))
{-# SPECIALIZE collect :: Expr Double -> Expr [Double] #-}
{-# SPECIALIZE collect :: Expr Float -> Expr [Float] #-}
{-# SPECIALIZE collect :: Expr Int -> Expr [Int] #-}
{-# INLINEABLE collect #-}

mode :: (Ord a, Columnable a, Eq a) => Expr a -> Expr a
mode :: forall a. (Ord a, Columnable a, Eq a) => Expr a -> Expr a
mode =
    AggStrategy a a -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg
        ( Text -> (Vector a -> a) -> AggStrategy a a
forall (v :: * -> *) b a.
(Vector v b, Typeable v) =>
Text -> (v b -> a) -> AggStrategy a b
CollectAgg
            Text
"mode"
            ( (a, Int) -> a
forall a b. (a, b) -> a
fst
                ((a, Int) -> a) -> (Vector a -> (a, Int)) -> Vector a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, Int) -> (a, Int) -> Ordering) -> [(a, Int)] -> (a, Int)
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
L.maximumBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> ((a, Int) -> Int) -> (a, Int) -> (a, Int) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (a, Int) -> Int
forall a b. (a, b) -> b
snd)
                ([(a, Int)] -> (a, Int))
-> (Vector a -> [(a, Int)]) -> Vector a -> (a, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map a Int -> [(a, Int)]
forall k a. Map k a -> [(k, a)]
M.toList
                (Map a Int -> [(a, Int)])
-> (Vector a -> Map a Int) -> Vector a -> [(a, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map a Int -> a -> Map a Int) -> Map a Int -> Vector a -> Map a Int
forall a b. (a -> b -> a) -> a -> Vector b -> a
V.foldl' (\Map a Int
m a
e -> (Int -> Int -> Int) -> a -> Int -> Map a Int -> Map a Int
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) a
e (Int
1 :: Int) Map a Int
m) Map a Int
forall k a. Map k a
M.empty
            )
        )
{-# SPECIALIZE mode :: Expr Double -> Expr Double #-}
{-# SPECIALIZE mode :: Expr Float -> Expr Float #-}
{-# SPECIALIZE mode :: Expr Int -> Expr Int #-}
{-# SPECIALIZE mode :: Expr Int8 -> Expr Int8 #-}
{-# SPECIALIZE mode :: Expr Int16 -> Expr Int16 #-}
{-# SPECIALIZE mode :: Expr Int32 -> Expr Int32 #-}
{-# SPECIALIZE mode :: Expr Int64 -> Expr Int64 #-}
{-# INLINEABLE mode #-}

minimum :: (Columnable a, Ord a) => Expr a -> Expr a
minimum :: forall a. (Columnable a, Ord a) => Expr a -> Expr a
minimum = AggStrategy a a -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg (Text -> Maybe a -> (a -> a -> a) -> AggStrategy a a
forall a b. Text -> Maybe a -> (a -> b -> a) -> AggStrategy a b
FoldAgg Text
"minimum" Maybe a
forall a. Maybe a
Nothing a -> a -> a
forall a. Ord a => a -> a -> a
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 #-}
{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int8 -> Expr Int8 #-}
{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int16 -> Expr Int16 #-}
{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int32 -> Expr Int32 #-}
{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int64 -> Expr Int64 #-}
{-# INLINEABLE DataFrame.Functions.minimum #-}

maximum :: (Columnable a, Ord a) => Expr a -> Expr a
maximum :: forall a. (Columnable a, Ord a) => Expr a -> Expr a
maximum = AggStrategy a a -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg (Text -> Maybe a -> (a -> a -> a) -> AggStrategy a a
forall a b. Text -> Maybe a -> (a -> b -> a) -> AggStrategy a b
FoldAgg Text
"maximum" Maybe a
forall a. Maybe a
Nothing a -> a -> a
forall a. Ord a => a -> a -> a
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 #-}
{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int8 -> Expr Int8 #-}
{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int16 -> Expr Int16 #-}
{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int32 -> Expr Int32 #-}
{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int64 -> Expr Int64 #-}
{-# INLINEABLE DataFrame.Functions.maximum #-}

sum :: forall a. (Columnable a, Num a) => Expr a -> Expr a
sum :: forall a. (Columnable a, Num a) => Expr a -> Expr a
sum = AggStrategy a a -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg (Text -> Maybe a -> (a -> a -> a) -> AggStrategy a a
forall a b. Text -> Maybe a -> (a -> b -> a) -> AggStrategy a b
FoldAgg Text
"sum" Maybe a
forall a. Maybe a
Nothing a -> a -> a
forall a. Num a => a -> a -> a
(+))
{-# SPECIALIZE DataFrame.Functions.sum :: Expr Double -> Expr Double #-}
{-# SPECIALIZE DataFrame.Functions.sum :: Expr Float -> Expr Float #-}
{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int -> Expr Int #-}
{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int8 -> Expr Int8 #-}
{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int16 -> Expr Int16 #-}
{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int32 -> Expr Int32 #-}
{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int64 -> Expr Int64 #-}
{-# INLINEABLE DataFrame.Functions.sum #-}

sumMaybe :: forall a. (Columnable a, Num a) => Expr (Maybe a) -> Expr a
sumMaybe :: forall a. (Columnable a, Num a) => Expr (Maybe a) -> Expr a
sumMaybe = AggStrategy a (Maybe a) -> Expr (Maybe a) -> Expr a
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg (Text -> (Vector (Maybe a) -> a) -> AggStrategy a (Maybe a)
forall (v :: * -> *) b a.
(Vector v b, Typeable v) =>
Text -> (v b -> a) -> AggStrategy a b
CollectAgg Text
"sumMaybe" ([a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
P.sum ([a] -> a) -> (Vector (Maybe a) -> [a]) -> Vector (Maybe a) -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe a] -> [a]
forall a. [Maybe a] -> [a]
Maybe.catMaybes ([Maybe a] -> [a])
-> (Vector (Maybe a) -> [Maybe a]) -> Vector (Maybe a) -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Maybe a) -> [Maybe a]
forall a. Vector a -> [a]
V.toList))
{-# SPECIALIZE sumMaybe :: Expr (Maybe Double) -> Expr Double #-}
{-# SPECIALIZE sumMaybe :: Expr (Maybe Float) -> Expr Float #-}
{-# SPECIALIZE sumMaybe :: Expr (Maybe Int) -> Expr Int #-}
{-# SPECIALIZE sumMaybe :: Expr (Maybe Int8) -> Expr Int8 #-}
{-# SPECIALIZE sumMaybe :: Expr (Maybe Int16) -> Expr Int16 #-}
{-# SPECIALIZE sumMaybe :: Expr (Maybe Int32) -> Expr Int32 #-}
{-# SPECIALIZE sumMaybe :: Expr (Maybe Int64) -> Expr Int64 #-}
{-# INLINEABLE sumMaybe #-}

mean :: (Columnable a, Real a) => Expr a -> Expr Double
mean :: forall a. (Columnable a, Real a) => Expr a -> Expr Double
mean =
    AggStrategy Double a -> Expr a -> Expr Double
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg
        ( Text
-> MeanAcc
-> (MeanAcc -> a -> MeanAcc)
-> (MeanAcc -> MeanAcc -> MeanAcc)
-> (MeanAcc -> Double)
-> AggStrategy Double a
forall acc b a.
Columnable acc =>
Text
-> acc
-> (acc -> b -> acc)
-> (acc -> acc -> acc)
-> (acc -> a)
-> AggStrategy a b
MergeAgg
            Text
"mean"
            (Double -> Int -> MeanAcc
MeanAcc Double
0.0 Int
0)
            (\(MeanAcc Double
s Int
c) a
x -> Double -> Int -> MeanAcc
MeanAcc (Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
+ a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac a
x) (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
            (\(MeanAcc Double
s1 Int
c1) (MeanAcc Double
s2 Int
c2) -> Double -> Int -> MeanAcc
MeanAcc (Double
s1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
s2) (Int
c1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
c2))
            (\(MeanAcc Double
s Int
c) -> if Int
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Double
0 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else Double
s Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c)
        )
{-# SPECIALIZE mean :: Expr Double -> Expr Double #-}
{-# SPECIALIZE mean :: Expr Float -> Expr Double #-}
{-# SPECIALIZE mean :: Expr Int -> Expr Double #-}
{-# SPECIALIZE mean :: Expr Int8 -> Expr Double #-}
{-# SPECIALIZE mean :: Expr Int16 -> Expr Double #-}
{-# SPECIALIZE mean :: Expr Int32 -> Expr Double #-}
{-# SPECIALIZE mean :: Expr Int64 -> Expr Double #-}
{-# INLINEABLE mean #-}

-- | The k largest values per group.
topK :: (Columnable a, Ord a) => Int -> Expr a -> Expr [a]
topK :: forall a. (Columnable a, Ord a) => Int -> Expr a -> Expr [a]
topK = Text -> (a -> a -> Bool) -> Int -> Expr a -> Expr [a]
forall a.
(Columnable a, Ord a) =>
Text -> (a -> a -> Bool) -> Int -> Expr a -> Expr [a]
kExtremes Text
"topK" a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>)
{-# INLINEABLE topK #-}

-- | The k smallest values per group.
bottomK :: (Columnable a, Ord a) => Int -> Expr a -> Expr [a]
bottomK :: forall a. (Columnable a, Ord a) => Int -> Expr a -> Expr [a]
bottomK = Text -> (a -> a -> Bool) -> Int -> Expr a -> Expr [a]
forall a.
(Columnable a, Ord a) =>
Text -> (a -> a -> Bool) -> Int -> Expr a -> Expr [a]
kExtremes Text
"bottomK" a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<)
{-# INLINEABLE bottomK #-}

{- | Top k values after applying an ordering function. Floating-point NaNs
are dropped: they have no place in a total order, and admitting them would
make the merge non-associative (results would depend on chunk boundaries in
the lazy executor).
-}
kExtremes ::
    forall a.
    (Columnable a, Ord a) =>
    T.Text -> (a -> a -> Bool) -> Int -> Expr a -> Expr [a]
kExtremes :: forall a.
(Columnable a, Ord a) =>
Text -> (a -> a -> Bool) -> Int -> Expr a -> Expr [a]
kExtremes Text
opName a -> a -> Bool
wins Int
k =
    AggStrategy [a] a -> Expr a -> Expr [a]
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg
        ( Text
-> [a]
-> ([a] -> a -> [a])
-> ([a] -> [a] -> [a])
-> ([a] -> [a])
-> AggStrategy [a] a
forall acc b a.
Columnable acc =>
Text
-> acc
-> (acc -> b -> acc)
-> (acc -> acc -> acc)
-> (acc -> a)
-> AggStrategy a b
MergeAgg
            (Text
opName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (Int -> [Char]
forall a. Show a => a -> [Char]
show Int
k))
            []
            [a] -> a -> [a]
step
            (([a] -> a -> [a]) -> [a] -> [a] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' [a] -> a -> [a]
step)
            [a] -> [a]
forall a. a -> a
id
        )
  where
    step :: [a] -> a -> [a]
step [a]
acc a
x
        | a -> Bool
isNaNLike a
x = [a]
acc
        | Bool
otherwise = a -> [a] -> [a]
insert a
x [a]
acc
    insert :: a -> [a] -> [a]
insert a
x = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
P.take Int
k ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
go
      where
        go :: [a] -> [a]
go (a
z : [a]
zs) | a -> a -> Bool
wins a
z a
x = a
z a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a]
go [a]
zs
        go [a]
zs = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
zs

    isNaNLike :: a -> Bool
    isNaNLike :: a -> Bool
isNaNLike a
x
        | Just a :~: Double
Refl <- TypeRep a -> TypeRep Double -> Maybe (a :~: Double)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Double) = a -> Bool
forall a. RealFloat a => a -> Bool
isNaN a
x
        | Just a :~: Float
Refl <- TypeRep a -> TypeRep Float -> Maybe (a :~: Float)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Float) = a -> Bool
forall a. RealFloat a => a -> Bool
isNaN a
x
        | Bool
otherwise = Bool
False

meanMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double
meanMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double
meanMaybe = AggStrategy Double (Maybe a) -> Expr (Maybe a) -> Expr Double
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg (Text
-> (Vector (Maybe a) -> Double) -> AggStrategy Double (Maybe a)
forall (v :: * -> *) b a.
(Vector v b, Typeable v) =>
Text -> (v b -> a) -> AggStrategy a b
CollectAgg Text
"meanMaybe" (Vector Double -> Double
forall a. (Real a, Unbox a) => Vector a -> Double
mean' (Vector Double -> Double)
-> (Vector (Maybe a) -> Vector Double)
-> Vector (Maybe a)
-> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Maybe a) -> Vector Double
forall a. Real a => Vector (Maybe a) -> Vector Double
optionalToDoubleVector))
{-# SPECIALIZE meanMaybe :: Expr (Maybe Double) -> Expr Double #-}
{-# SPECIALIZE meanMaybe :: Expr (Maybe Float) -> Expr Double #-}
{-# SPECIALIZE meanMaybe :: Expr (Maybe Int) -> Expr Double #-}
{-# SPECIALIZE meanMaybe :: Expr (Maybe Int8) -> Expr Double #-}
{-# SPECIALIZE meanMaybe :: Expr (Maybe Int16) -> Expr Double #-}
{-# SPECIALIZE meanMaybe :: Expr (Maybe Int32) -> Expr Double #-}
{-# SPECIALIZE meanMaybe :: Expr (Maybe Int64) -> Expr Double #-}
{-# INLINEABLE meanMaybe #-}

variance :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double
variance :: forall a. (Columnable a, Real a, Unbox a) => Expr a -> Expr Double
variance = AggStrategy Double a -> Expr a -> Expr Double
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg (Text -> (Vector a -> Double) -> AggStrategy Double a
forall (v :: * -> *) b a.
(Vector v b, Typeable v) =>
Text -> (v b -> a) -> AggStrategy a b
CollectAgg Text
"variance" Vector a -> Double
forall a. (Real a, Unbox a) => Vector a -> Double
variance')
{-# SPECIALIZE variance :: Expr Double -> Expr Double #-}
{-# SPECIALIZE variance :: Expr Float -> Expr Double #-}
{-# SPECIALIZE variance :: Expr Int -> Expr Double #-}
{-# SPECIALIZE variance :: Expr Int8 -> Expr Double #-}
{-# SPECIALIZE variance :: Expr Int16 -> Expr Double #-}
{-# SPECIALIZE variance :: Expr Int32 -> Expr Double #-}
{-# SPECIALIZE variance :: Expr Int64 -> Expr Double #-}
{-# INLINEABLE variance #-}

median :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double
median :: forall a. (Columnable a, Real a, Unbox a) => Expr a -> Expr Double
median = AggStrategy Double a -> Expr a -> Expr Double
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg (Text -> (Vector a -> Double) -> AggStrategy Double a
forall (v :: * -> *) b a.
(Vector v b, Typeable v) =>
Text -> (v b -> a) -> AggStrategy a b
CollectAgg Text
"median" Vector a -> Double
forall a. (Real a, Unbox a) => Vector a -> Double
median')
{-# SPECIALIZE median :: Expr Double -> Expr Double #-}
{-# SPECIALIZE median :: Expr Float -> Expr Double #-}
{-# SPECIALIZE median :: Expr Int -> Expr Double #-}
{-# SPECIALIZE median :: Expr Int8 -> Expr Double #-}
{-# SPECIALIZE median :: Expr Int16 -> Expr Double #-}
{-# SPECIALIZE median :: Expr Int32 -> Expr Double #-}
{-# SPECIALIZE median :: Expr Int64 -> Expr Double #-}
{-# INLINEABLE median #-}

medianMaybe :: (Columnable a, Real a) => Expr (Maybe a) -> Expr Double
medianMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double
medianMaybe = AggStrategy Double (Maybe a) -> Expr (Maybe a) -> Expr Double
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg (Text
-> (Vector (Maybe a) -> Double) -> AggStrategy Double (Maybe a)
forall (v :: * -> *) b a.
(Vector v b, Typeable v) =>
Text -> (v b -> a) -> AggStrategy a b
CollectAgg Text
"meanMaybe" (Vector Double -> Double
forall a. (Real a, Unbox a) => Vector a -> Double
median' (Vector Double -> Double)
-> (Vector (Maybe a) -> Vector Double)
-> Vector (Maybe a)
-> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Maybe a) -> Vector Double
forall a. Real a => Vector (Maybe a) -> Vector Double
optionalToDoubleVector))
{-# SPECIALIZE medianMaybe :: Expr (Maybe Double) -> Expr Double #-}
{-# SPECIALIZE medianMaybe :: Expr (Maybe Float) -> Expr Double #-}
{-# SPECIALIZE medianMaybe :: Expr (Maybe Int) -> Expr Double #-}
{-# SPECIALIZE medianMaybe :: Expr (Maybe Int8) -> Expr Double #-}
{-# SPECIALIZE medianMaybe :: Expr (Maybe Int16) -> Expr Double #-}
{-# SPECIALIZE medianMaybe :: Expr (Maybe Int32) -> Expr Double #-}
{-# SPECIALIZE medianMaybe :: Expr (Maybe Int64) -> Expr Double #-}
{-# INLINEABLE medianMaybe #-}

optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double
optionalToDoubleVector :: forall a. Real a => Vector (Maybe a) -> Vector Double
optionalToDoubleVector =
    [Double] -> Vector Double
forall a. Unbox a => [a] -> Vector a
VU.fromList
        ([Double] -> Vector Double)
-> (Vector (Maybe a) -> [Double])
-> Vector (Maybe a)
-> Vector Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Double] -> Maybe a -> [Double])
-> [Double] -> Vector (Maybe a) -> [Double]
forall a b. (a -> b -> a) -> a -> Vector b -> a
V.foldl'
            (\[Double]
acc Maybe a
e -> if Maybe a -> Bool
forall a. Maybe a -> Bool
Maybe.isJust Maybe a
e then a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (a -> Maybe a -> a
forall a. a -> Maybe a -> a
Maybe.fromMaybe a
0 Maybe a
e) Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double]
acc else [Double]
acc)
            []

percentile :: Int -> Expr Double -> Expr Double
percentile :: Int -> Expr Double -> Expr Double
percentile Int
n =
    AggStrategy Double Double -> Expr Double -> Expr Double
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg
        ( Text -> (Vector Double -> Double) -> AggStrategy Double Double
forall (v :: * -> *) b a.
(Vector v b, Typeable v) =>
Text -> (v b -> a) -> AggStrategy a b
CollectAgg
            ([Char] -> Text
T.pack ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ [Char]
"percentile " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n)
            (Int -> Vector Double -> Double
forall a. (Unbox a, Num a, Real a) => Int -> Vector a -> Double
percentile' Int
n)
        )

stddev :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double
stddev :: forall a. (Columnable a, Real a, Unbox a) => Expr a -> Expr Double
stddev = AggStrategy Double a -> Expr a -> Expr Double
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg (Text -> (Vector a -> Double) -> AggStrategy Double a
forall (v :: * -> *) b a.
(Vector v b, Typeable v) =>
Text -> (v b -> a) -> AggStrategy a b
CollectAgg Text
"stddev" (Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> (Vector a -> Double) -> Vector a -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Double
forall a. (Real a, Unbox a) => Vector a -> Double
variance'))
{-# SPECIALIZE stddev :: Expr Double -> Expr Double #-}
{-# SPECIALIZE stddev :: Expr Float -> Expr Double #-}
{-# SPECIALIZE stddev :: Expr Int -> Expr Double #-}
{-# SPECIALIZE stddev :: Expr Int8 -> Expr Double #-}
{-# SPECIALIZE stddev :: Expr Int16 -> Expr Double #-}
{-# SPECIALIZE stddev :: Expr Int32 -> Expr Double #-}
{-# SPECIALIZE stddev :: Expr Int64 -> Expr Double #-}
{-# INLINEABLE stddev #-}

stddevMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double
stddevMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double
stddevMaybe = AggStrategy Double (Maybe a) -> Expr (Maybe a) -> Expr Double
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg (Text
-> (Vector (Maybe a) -> Double) -> AggStrategy Double (Maybe a)
forall (v :: * -> *) b a.
(Vector v b, Typeable v) =>
Text -> (v b -> a) -> AggStrategy a b
CollectAgg Text
"stddevMaybe" (Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double)
-> (Vector (Maybe a) -> Double) -> Vector (Maybe a) -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Double -> Double
forall a. (Real a, Unbox a) => Vector a -> Double
variance' (Vector Double -> Double)
-> (Vector (Maybe a) -> Vector Double)
-> Vector (Maybe a)
-> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Maybe a) -> Vector Double
forall a. Real a => Vector (Maybe a) -> Vector Double
optionalToDoubleVector))
{-# SPECIALIZE stddevMaybe :: Expr (Maybe Double) -> Expr Double #-}
{-# SPECIALIZE stddevMaybe :: Expr (Maybe Float) -> Expr Double #-}
{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int) -> Expr Double #-}
{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int8) -> Expr Double #-}
{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int16) -> Expr Double #-}
{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int32) -> Expr Double #-}
{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int64) -> Expr Double #-}
{-# INLINEABLE stddevMaybe #-}

zScore :: Expr Double -> Expr Double
zScore :: Expr Double -> Expr Double
zScore Expr Double
c = (Expr Double
c Expr Double -> Expr Double -> Expr Double
forall a. Num a => a -> a -> a
- Expr Double -> Expr Double
forall a. (Columnable a, Real a) => Expr a -> Expr Double
mean Expr Double
c) Expr Double -> Expr Double -> Expr Double
forall a. Fractional a => a -> a -> a
/ Expr Double -> Expr Double
forall a. (Columnable a, Real a, Unbox a) => Expr a -> Expr Double
stddev Expr Double
c

pow :: (Columnable a, Num a) => Expr a -> Int -> Expr a
pow :: forall a. (Columnable a, Num a) => Expr a -> Int -> Expr a
pow Expr a
expr Int
i = (a -> Int -> a)
-> Text
-> Maybe Text
-> Bool
-> Int
-> Expr a
-> Expr Int
-> Expr a
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr c -> Expr b -> Expr a
lift2Decorated a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
(^) Text
"pow" (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"^") Bool
True Int
8 Expr a
expr (Int -> Expr Int
forall a. Columnable a => a -> Expr a
Lit Int
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 :: forall a. (Columnable a, Num a, Ord a) => Expr a -> Expr a
relu = (a -> a) -> Text -> Maybe Text -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated (a -> a -> a
forall a. Ord a => a -> a -> a
Prelude.max a
0) Text
"relu" Maybe Text
forall a. Maybe a
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 :: forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr a
min = (a -> a -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr a -> Expr a -> Expr a
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr c -> Expr b -> Expr a
lift2Decorated a -> a -> a
forall a. Ord a => a -> a -> a
Prelude.min Text
"min" Maybe Text
forall a. Maybe a
Nothing Bool
True Int
1
{-# SPECIALIZE DataFrame.Functions.min ::
    Expr Double -> Expr Double -> Expr Double
    #-}
{-# SPECIALIZE DataFrame.Functions.min ::
    Expr Float -> Expr Float -> Expr Float
    #-}
{-# SPECIALIZE DataFrame.Functions.min :: Expr Int -> Expr Int -> Expr Int #-}
{-# INLINEABLE DataFrame.Functions.min #-}

max :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr a
max :: forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr a
max = (a -> a -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr a -> Expr a -> Expr a
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr c -> Expr b -> Expr a
lift2Decorated a -> a -> a
forall a. Ord a => a -> a -> a
Prelude.max Text
"max" Maybe Text
forall a. Maybe a
Nothing Bool
True Int
1
{-# SPECIALIZE DataFrame.Functions.max ::
    Expr Double -> Expr Double -> Expr Double
    #-}
{-# SPECIALIZE DataFrame.Functions.max ::
    Expr Float -> Expr Float -> Expr Float
    #-}
{-# SPECIALIZE DataFrame.Functions.max :: Expr Int -> Expr Int -> Expr Int #-}
{-# INLINEABLE DataFrame.Functions.max #-}

reduce ::
    forall a b.
    (Columnable a, Columnable b) =>
    Expr b ->
    a ->
    (a -> b -> a) ->
    Expr a
reduce :: forall a b.
(Columnable a, Columnable b) =>
Expr b -> a -> (a -> b -> a) -> Expr a
reduce Expr b
expr a
start a -> b -> a
f = AggStrategy a b -> Expr b -> Expr a
forall a b.
(Columnable a, Columnable b) =>
AggStrategy a b -> Expr b -> Expr a
Agg (Text -> Maybe a -> (a -> b -> a) -> AggStrategy a b
forall a b. Text -> Maybe a -> (a -> b -> a) -> AggStrategy a b
FoldAgg Text
"foldUdf" (a -> Maybe a
forall a. a -> Maybe a
Just a
start) a -> b -> a
f) Expr b
expr
{-# INLINEABLE reduce #-}

toMaybe :: (Columnable a) => Expr a -> Expr (Maybe a)
toMaybe :: forall a. Columnable a => Expr a -> Expr (Maybe a)
toMaybe = (a -> Maybe a) -> Text -> Maybe Text -> Expr a -> Expr (Maybe a)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated a -> Maybe a
forall a. a -> Maybe a
Just Text
"toMaybe" Maybe Text
forall a. Maybe a
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 :: forall a. Columnable a => a -> Expr (Maybe a) -> Expr a
fromMaybe a
d = (Maybe a -> a) -> Text -> Maybe Text -> Expr (Maybe a) -> Expr a
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated (a -> Maybe a -> a
forall a. a -> Maybe a -> a
Maybe.fromMaybe a
d) Text
"fromMaybe" Maybe Text
forall a. Maybe a
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 :: forall a. Columnable a => Expr (Maybe a) -> Expr Bool
isJust = (Maybe a -> Bool)
-> Text -> Maybe Text -> Expr (Maybe a) -> Expr Bool
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated Maybe a -> Bool
forall a. Maybe a -> Bool
Maybe.isJust Text
"isJust" Maybe Text
forall a. Maybe a
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 :: forall a. Columnable a => Expr (Maybe a) -> Expr Bool
isNothing = (Maybe a -> Bool)
-> Text -> Maybe Text -> Expr (Maybe a) -> Expr Bool
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated Maybe a -> Bool
forall a. Maybe a -> Bool
Maybe.isNothing Text
"isNothing" Maybe Text
forall a. Maybe a
Nothing
{-# SPECIALIZE isNothing :: Expr (Maybe Double) -> Expr Bool #-}
{-# SPECIALIZE isNothing :: Expr (Maybe Int) -> Expr Bool #-}
{-# INLINEABLE isNothing #-}

-- | SQL spelling of 'isNothing'.
isNull :: (Columnable a) => Expr (Maybe a) -> Expr Bool
isNull :: forall a. Columnable a => Expr (Maybe a) -> Expr Bool
isNull = Expr (Maybe a) -> Expr Bool
forall a. Columnable a => Expr (Maybe a) -> Expr Bool
isNothing
{-# INLINEABLE isNull #-}

-- | SQL spelling of 'isJust'.
isNotNull :: (Columnable a) => Expr (Maybe a) -> Expr Bool
isNotNull :: forall a. Columnable a => Expr (Maybe a) -> Expr Bool
isNotNull = Expr (Maybe a) -> Expr Bool
forall a. Columnable a => Expr (Maybe a) -> Expr Bool
isJust
{-# INLINEABLE isNotNull #-}

fromJust :: (Columnable a) => Expr (Maybe a) -> Expr a
fromJust :: forall a. Columnable a => Expr (Maybe a) -> Expr a
fromJust = (Maybe a -> a) -> Text -> Maybe Text -> Expr (Maybe a) -> Expr a
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated Maybe a -> a
forall a. HasCallStack => Maybe a -> a
Maybe.fromJust Text
"fromJust" Maybe Text
forall a. Maybe a
Nothing
{-# SPECIALIZE fromJust :: Expr (Maybe Double) -> Expr Double #-}
{-# SPECIALIZE fromJust :: Expr (Maybe Int) -> Expr Int #-}
{-# INLINEABLE fromJust #-}

whenPresent ::
    forall a b.
    (Columnable a, Columnable b) =>
    (a -> b) ->
    Expr (Maybe a) ->
    Expr (Maybe b)
whenPresent :: forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Expr (Maybe a) -> Expr (Maybe b)
whenPresent a -> b
f = (Maybe a -> Maybe b)
-> Text -> Maybe Text -> Expr (Maybe a) -> Expr (Maybe b)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated ((a -> b) -> Maybe a -> Maybe b
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) Text
"whenPresent" Maybe Text
forall a. Maybe a
Nothing
{-# INLINEABLE whenPresent #-}

whenBothPresent ::
    forall a b c.
    (Columnable a, Columnable b, Columnable c) =>
    (a -> b -> c) ->
    Expr (Maybe a) ->
    Expr (Maybe b) ->
    Expr (Maybe c)
whenBothPresent :: forall a b c.
(Columnable a, Columnable b, Columnable c) =>
(a -> b -> c) -> Expr (Maybe a) -> Expr (Maybe b) -> Expr (Maybe c)
whenBothPresent a -> b -> c
f = (Maybe a -> Maybe b -> Maybe c)
-> Text
-> Maybe Text
-> Bool
-> Int
-> Expr (Maybe a)
-> Expr (Maybe b)
-> Expr (Maybe c)
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr c -> Expr b -> Expr a
lift2Decorated (\Maybe a
l Maybe b
r -> a -> b -> c
f (a -> b -> c) -> Maybe a -> Maybe (b -> c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe a
l Maybe (b -> c) -> Maybe b -> Maybe c
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe b
r) Text
"whenBothPresent" Maybe Text
forall a. Maybe a
Nothing Bool
False Int
0
{-# INLINEABLE whenBothPresent #-}

recode ::
    forall a b.
    (Columnable a, Columnable b, Show (a, b)) =>
    [(a, b)] ->
    Expr a ->
    Expr (Maybe b)
recode :: forall a b.
(Columnable a, Columnable b, Show (a, b)) =>
[(a, b)] -> Expr a -> Expr (Maybe b)
recode [(a, b)]
mapping =
    UnUDF a (Maybe b) -> Expr a -> Expr (Maybe b)
forall (op :: * -> * -> *) a b.
(UnaryOp op, Columnable a, Columnable b) =>
op b a -> Expr b -> Expr a
Unary
        ( MkUnaryOp
            { unaryFn :: a -> Maybe b
unaryFn = (a -> [(a, b)] -> Maybe b
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(a, b)]
mapping)
            , unaryName :: Text
unaryName = Text
"recode " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack ([(a, b)] -> [Char]
forall a. Show a => a -> [Char]
show [(a, b)]
mapping)
            , unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing
            }
        )

recodeWithCondition ::
    forall a b.
    (Columnable a, Columnable b) =>
    Expr b ->
    [(Expr a -> Expr Bool, b)] ->
    Expr a ->
    Expr b
recodeWithCondition :: forall a b.
(Columnable a, Columnable b) =>
Expr b -> [(Expr a -> Expr Bool, b)] -> Expr a -> Expr b
recodeWithCondition Expr b
fallback [] Expr a
_val = Expr b
fallback
recodeWithCondition Expr b
fallback ((Expr a -> Expr Bool
cond, b
val) : [(Expr a -> Expr Bool, b)]
rest) Expr a
expr = Expr Bool -> Expr b -> Expr b -> Expr b
forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
ifThenElse (Expr a -> Expr Bool
cond Expr a
expr) (b -> Expr b
forall a. Columnable a => a -> Expr a
lit b
val) (Expr b -> [(Expr a -> Expr Bool, b)] -> Expr a -> Expr b
forall a b.
(Columnable a, Columnable b) =>
Expr b -> [(Expr a -> Expr Bool, b)] -> Expr a -> Expr b
recodeWithCondition Expr b
fallback [(Expr a -> Expr Bool, b)]
rest Expr a
expr)

recodeWithDefault ::
    forall a b.
    (Columnable a, Columnable b, Show (a, b)) =>
    b ->
    [(a, b)] ->
    Expr a ->
    Expr b
recodeWithDefault :: forall a b.
(Columnable a, Columnable b, Show (a, b)) =>
b -> [(a, b)] -> Expr a -> Expr b
recodeWithDefault b
d [(a, b)]
mapping =
    UnUDF a b -> Expr a -> Expr b
forall (op :: * -> * -> *) a b.
(UnaryOp op, Columnable a, Columnable b) =>
op b a -> Expr b -> Expr a
Unary
        ( MkUnaryOp
            { unaryFn :: a -> b
unaryFn = b -> Maybe b -> b
forall a. a -> Maybe a -> a
Maybe.fromMaybe b
d (Maybe b -> b) -> (a -> Maybe b) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> [(a, b)] -> Maybe b
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(a, b)]
mapping)
            , unaryName :: Text
unaryName =
                Text
"recodeWithDefault " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (b -> [Char]
forall a. Show a => a -> [Char]
show b
d) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack ([(a, b)] -> [Char]
forall a. Show a => a -> [Char]
show [(a, b)]
mapping)
            , unarySymbol :: Maybe Text
unarySymbol = Maybe Text
forall a. Maybe a
Nothing
            }
        )

firstOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a)
firstOrNothing :: forall a. Columnable a => Expr [a] -> Expr (Maybe a)
firstOrNothing = ([a] -> Maybe a)
-> Text -> Maybe Text -> Expr [a] -> Expr (Maybe a)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated [a] -> Maybe a
forall a. [a] -> Maybe a
Maybe.listToMaybe Text
"firstOrNothing" Maybe Text
forall a. Maybe a
Nothing

lastOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a)
lastOrNothing :: forall a. Columnable a => Expr [a] -> Expr (Maybe a)
lastOrNothing = ([a] -> Maybe a)
-> Text -> Maybe Text -> Expr [a] -> Expr (Maybe a)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated ([a] -> Maybe a
forall a. [a] -> Maybe a
Maybe.listToMaybe ([a] -> Maybe a) -> ([a] -> [a]) -> [a] -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. [a] -> [a]
reverse) Text
"lastOrNothing" Maybe Text
forall a. Maybe a
Nothing

splitOn :: T.Text -> Expr T.Text -> Expr [T.Text]
splitOn :: Text -> Expr Text -> Expr [Text]
splitOn Text
delim = (Text -> [Text]) -> Text -> Maybe Text -> Expr Text -> Expr [Text]
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated (HasCallStack => Text -> Text -> [Text]
Text -> Text -> [Text]
T.splitOn Text
delim) Text
"splitOn" Maybe Text
forall a. Maybe a
Nothing

match :: T.Text -> Expr T.Text -> Expr (Maybe T.Text)
match :: Text -> Expr Text -> Expr (Maybe Text)
match Text
regex =
    (Text -> Maybe Text)
-> Text -> Maybe Text -> Expr Text -> Expr (Maybe Text)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated
        ((\Text
r -> if Text -> Bool
T.null Text
r then Maybe Text
forall a. Maybe a
Nothing else Text -> Maybe Text
forall a. a -> Maybe a
Just Text
r) (Text -> Maybe Text) -> (Text -> Text) -> Text -> Maybe Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> Text -> Text
forall source source1 target.
(RegexMaker Regex CompOption ExecOption source,
 RegexContext Regex source1 target) =>
source1 -> source -> target
=~ Text
regex))
        (Text
"match " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (Text -> [Char]
forall a. Show a => a -> [Char]
show Text
regex))
        Maybe Text
forall a. Maybe a
Nothing

matchAll :: T.Text -> Expr T.Text -> Expr [T.Text]
matchAll :: Text -> Expr Text -> Expr [Text]
matchAll Text
regex =
    (Text -> [Text]) -> Text -> Maybe Text -> Expr Text -> Expr [Text]
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated
        (AllTextMatches [] Text -> [Text]
forall (f :: * -> *) b. AllTextMatches f b -> f b
getAllTextMatches (AllTextMatches [] Text -> [Text])
-> (Text -> AllTextMatches [] Text) -> Text -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> Text -> AllTextMatches [] Text
forall source source1 target.
(RegexMaker Regex CompOption ExecOption source,
 RegexContext Regex source1 target) =>
source1 -> source -> target
=~ Text
regex))
        (Text
"matchAll " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (Text -> [Char]
forall a. Show a => a -> [Char]
show Text
regex))
        Maybe Text
forall a. Maybe a
Nothing

parseDate ::
    (ParseTime t, Columnable t) => T.Text -> Expr T.Text -> Expr (Maybe t)
parseDate :: forall t.
(ParseTime t, Columnable t) =>
Text -> Expr Text -> Expr (Maybe t)
parseDate Text
format =
    (Text -> Maybe t)
-> Text -> Maybe Text -> Expr Text -> Expr (Maybe t)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated
        (Bool -> TimeLocale -> [Char] -> [Char] -> Maybe t
forall (m :: * -> *) t.
(MonadFail m, ParseTime t) =>
Bool -> TimeLocale -> [Char] -> [Char] -> m t
parseTimeM Bool
True TimeLocale
defaultTimeLocale (Text -> [Char]
T.unpack Text
format) ([Char] -> Maybe t) -> (Text -> [Char]) -> Text -> Maybe t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Char]
T.unpack)
        (Text
"parseDate " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
format)
        Maybe Text
forall a. Maybe a
Nothing

daysBetween :: Expr Day -> Expr Day -> Expr Int
daysBetween :: Expr Day -> Expr Day -> Expr Int
daysBetween =
    (Day -> Day -> Int)
-> Text
-> Maybe Text
-> Bool
-> Int
-> Expr Day
-> Expr Day
-> Expr Int
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Text -> Maybe Text -> Bool -> Int -> Expr c -> Expr b -> Expr a
lift2Decorated
        (\Day
d1 Day
d2 -> Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Day -> Day -> Integer
diffDays Day
d1 Day
d2))
        Text
"daysBetween"
        Maybe Text
forall a. Maybe a
Nothing
        Bool
True
        Int
2

bind ::
    forall a b m.
    (Columnable a, Columnable (m a), Monad m, Columnable b, Columnable (m b)) =>
    (a -> m b) ->
    Expr (m a) ->
    Expr (m b)
bind :: forall a b (m :: * -> *).
(Columnable a, Columnable (m a), Monad m, Columnable b,
 Columnable (m b)) =>
(a -> m b) -> Expr (m a) -> Expr (m b)
bind a -> m b
f = (m a -> m b) -> Text -> Maybe Text -> Expr (m a) -> Expr (m b)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Text -> Maybe Text -> Expr a -> Expr b
liftDecorated (m a -> (a -> m b) -> m b
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> m b
f) Text
"bind" Maybe Text
forall a. Maybe a
Nothing

{- | Window function: evaluate an expression partitioned by the given columns.

Each partition computes the inner expression independently, and the result
is broadcast back to every row in that partition. This is analogous to
Polars' @.over()@ or SQL @OVER (PARTITION BY ...)@.

@
-- Per-country median, broadcast to every row:
F.over [\"country\"] (F.median (F.col \@Double \"amount\"))

-- Deviation from group mean:
F.col \@Double \"amount\" - F.over [\"group\"] (F.mean (F.col \@Double \"amount\"))
@
-}
over :: (Columnable a) => [T.Text] -> Expr a -> Expr a
over :: forall a. Columnable a => [Text] -> Expr a -> Expr a
over = [Text] -> Expr a -> Expr a
forall a. Columnable a => [Text] -> Expr a -> Expr a
Over

-- See Section 2.4 of the Haskell Report https://www.haskell.org/definition/haskell2010.pdf
isReservedId :: T.Text -> Bool
isReservedId :: Text -> Bool
isReservedId Text
t = case Text
t of
    Text
"case" -> Bool
True
    Text
"class" -> Bool
True
    Text
"data" -> Bool
True
    Text
"default" -> Bool
True
    Text
"deriving" -> Bool
True
    Text
"do" -> Bool
True
    Text
"else" -> Bool
True
    Text
"foreign" -> Bool
True
    Text
"if" -> Bool
True
    Text
"import" -> Bool
True
    Text
"in" -> Bool
True
    Text
"infix" -> Bool
True
    Text
"infixl" -> Bool
True
    Text
"infixr" -> Bool
True
    Text
"instance" -> Bool
True
    Text
"let" -> Bool
True
    Text
"module" -> Bool
True
    Text
"newtype" -> Bool
True
    Text
"of" -> Bool
True
    Text
"then" -> Bool
True
    Text
"type" -> Bool
True
    Text
"where" -> Bool
True
    Text
_ -> Bool
False

isVarId :: T.Text -> Bool
isVarId :: Text -> Bool
isVarId Text
t = case Text -> Maybe (Char, Text)
T.uncons Text
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 (Char
c, Text
_) -> Char -> Bool
Char.isLower Char
c Bool -> Bool -> Bool
&& Char -> Bool
Char.isAlpha Char
c
    Maybe (Char, Text)
Nothing -> Bool
False

isHaskellIdentifier :: T.Text -> Bool
isHaskellIdentifier :: Text -> Bool
isHaskellIdentifier Text
t = Bool -> Bool
Prelude.not (Text -> Bool
isVarId Text
t) Bool -> Bool -> Bool
|| Text -> Bool
isReservedId Text
t

sanitize :: T.Text -> T.Text
sanitize :: Text -> Text
sanitize Text
t
    | Bool
isValid = Text
t
    | Text -> Bool
isHaskellIdentifier Text
t' = Text
"_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
t' Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_"
    | Bool
otherwise = Text
t'
  where
    isValid :: Bool
isValid =
        Bool -> Bool
Prelude.not (Text -> Bool
isHaskellIdentifier Text
t)
            Bool -> Bool -> Bool
&& Text -> Bool
isVarId Text
t
            Bool -> Bool -> Bool
&& (Char -> Bool) -> Text -> Bool
T.all Char -> Bool
Char.isAlphaNum Text
t
    t' :: Text
t' = (Char -> Char) -> Text -> Text
T.map Char -> Char
replaceInvalidCharacters (Text -> Text) -> (Text -> Text) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> Text -> Text
T.filter (Bool -> Bool
Prelude.not (Bool -> Bool) -> (Char -> Bool) -> Char -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Bool
parentheses) (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Text
t
    replaceInvalidCharacters :: Char -> Char
replaceInvalidCharacters Char
c
        | Char -> Bool
Char.isUpper Char
c = Char -> Char
Char.toLower Char
c
        | Char -> Bool
Char.isSpace Char
c = Char
'_'
        | Char -> Bool
Char.isPunctuation Char
c = Char
'_' -- '-' will also become a '_'
        | Char -> Bool
Char.isSymbol Char
c = Char
'_'
        | Char -> Bool
Char.isAlphaNum Char
c = Char
c -- Blanket condition
        | Bool
otherwise = Char
'_' -- If we're unsure we'll default to an underscore
    parentheses :: Char -> Bool
parentheses Char
c = case Char
c of
        Char
'(' -> Bool
True
        Char
')' -> Bool
True
        Char
'{' -> Bool
True
        Char
'}' -> Bool
True
        Char
'[' -> Bool
True
        Char
']' -> Bool
True
        Char
_ -> Bool
False