{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

{- | Type-safe expression construction for typed DataFrames.

Unlike the untyped @Expr a@ where column references are unchecked strings,
'TExpr' ensures at compile time that:

* Referenced columns exist in the schema
* Column types match the expression type

== Example

@
type Schema = '[Column \"age\" Int, Column \"salary\" Double]

-- This compiles:
goodExpr :: TExpr Schema Double
goodExpr = col \@\"salary\"

-- This gives a compile-time error (column not found):
badExpr :: TExpr Schema Double
badExpr = col \@\"nonexistent\"

-- This gives a compile-time error (type mismatch):
wrongType :: TExpr Schema Int
wrongType = col \@\"salary\"  -- salary is Double, not Int
@
-}
module DataFrame.Typed.Expr (
    -- * Core typed expression type (re-exported from Types)
    TExpr (..),

    -- * Column reference (schema-checked)
    col,

    -- * Literals
    lit,

    -- * Conditional
    ifThenElse,

    -- * Unary / binary lifting
    lift,
    lift2,

    -- * Comparison operators
    (.==.),
    (./=.),
    (.<.),
    (.<=.),
    (.>=.),
    (.>.),

    -- * Logical operators
    (.&&.),
    (.||.),
    DataFrame.Typed.Expr.not,

    -- * Aggregation combinators
    sum,
    mean,
    count,
    minimum,
    maximum,
    collect,

    -- * Named expression helper
    as,

    -- * Sort helpers
    asc,
    desc,
) where

import Data.Proxy (Proxy (..))
import Data.String (IsString (..))
import qualified Data.Text as T
import GHC.TypeLits (KnownSymbol, Symbol, symbolVal)

import DataFrame.Internal.Column (Columnable)
import DataFrame.Internal.Expression (
    AggStrategy (..),
    BinaryOp (..),
    Expr (..),
    MeanAcc (..),
    NamedExpr,
    UExpr (..),
    UnaryOp (..),
 )

import DataFrame.Typed.Schema (AssertPresent, Lookup)
import DataFrame.Typed.Types (TExpr (..), TSortOrder (..))
import Prelude hiding (maximum, minimum, sum)

{- | Create a typed column reference. This is the key type-safety entry point.

The column name must exist in @cols@ and its type must match @a@.
Both checks happen at compile time via type families.

@
salary :: TExpr '[Column \"salary\" Double] Double
salary = col \@\"salary\"
@
-}
col ::
    forall (name :: Symbol) cols a.
    ( KnownSymbol name
    , a ~ Lookup name cols
    , Columnable a
    , AssertPresent name cols
    ) =>
    TExpr cols a
col :: forall (name :: Symbol) (cols :: [*]) a.
(KnownSymbol name, a ~ Lookup name cols, Columnable a,
 AssertPresent name cols) =>
TExpr cols a
col = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Text -> Expr a
forall a. Columnable a => Text -> Expr a
Col (String -> Text
T.pack (Proxy name -> String
forall (n :: Symbol) (proxy :: Symbol -> *).
KnownSymbol n =>
proxy n -> String
symbolVal (forall {k} (t :: k). Proxy t
forall (t :: Symbol). Proxy t
Proxy @name))))

{- | Create a literal expression. Valid for any schema since it
references no columns.
-}
lit :: (Columnable a) => a -> TExpr cols a
lit :: forall a (cols :: [*]). Columnable a => a -> TExpr cols a
lit = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> TExpr cols a) -> (a -> Expr a) -> a -> TExpr cols a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Expr a
forall a. Columnable a => a -> Expr a
Lit

-- | Conditional expression.
ifThenElse ::
    (Columnable a) =>
    TExpr cols Bool -> TExpr cols a -> TExpr cols a -> TExpr cols a
ifThenElse :: forall a (cols :: [*]).
Columnable a =>
TExpr cols Bool -> TExpr cols a -> TExpr cols a -> TExpr cols a
ifThenElse (TExpr Expr Bool
c) (TExpr Expr a
t) (TExpr Expr a
e) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr Bool -> Expr a -> Expr a -> Expr a
forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
If Expr Bool
c Expr a
t Expr a
e)

-------------------------------------------------------------------------------
-- Numeric instances (mirror Expr's instances)
-------------------------------------------------------------------------------

instance (Num a, Columnable a) => Num (TExpr cols a) where
    (TExpr Expr a
a) + :: TExpr cols a -> TExpr cols a -> TExpr cols a
+ (TExpr Expr a
b) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a
a Expr a -> Expr a -> Expr a
forall a. Num a => a -> a -> a
+ Expr a
b)
    (TExpr Expr a
a) - :: TExpr cols a -> TExpr cols a -> TExpr cols a
- (TExpr Expr a
b) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a
a Expr a -> Expr a -> Expr a
forall a. Num a => a -> a -> a
- Expr a
b)
    (TExpr Expr a
a) * :: TExpr cols a -> TExpr cols a -> TExpr cols a
* (TExpr Expr a
b) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a
a Expr a -> Expr a -> Expr a
forall a. Num a => a -> a -> a
* Expr a
b)
    negate :: TExpr cols a -> TExpr cols a
negate (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Num a => a -> a
negate Expr a
a)
    abs :: TExpr cols a -> TExpr cols a
abs (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Num a => a -> a
abs Expr a
a)
    signum :: TExpr cols a -> TExpr cols a
signum (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Num a => a -> a
signum Expr a
a)
    fromInteger :: Integer -> TExpr cols a
fromInteger = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> TExpr cols a)
-> (Integer -> Expr a) -> Integer -> TExpr cols a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Expr a
forall a. Num a => Integer -> a
fromInteger

instance (Fractional a, Columnable a) => Fractional (TExpr cols a) where
    fromRational :: Rational -> TExpr cols a
fromRational = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> TExpr cols a)
-> (Rational -> Expr a) -> Rational -> TExpr cols a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rational -> Expr a
forall a. Fractional a => Rational -> a
fromRational
    (TExpr Expr a
a) / :: TExpr cols a -> TExpr cols a -> TExpr cols a
/ (TExpr Expr a
b) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a
a Expr a -> Expr a -> Expr a
forall a. Fractional a => a -> a -> a
/ Expr a
b)

instance (Floating a, Columnable a) => Floating (TExpr cols a) where
    pi :: TExpr cols a
pi = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr Expr a
forall a. Floating a => a
pi
    exp :: TExpr cols a -> TExpr cols a
exp (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
exp Expr a
a)
    sqrt :: TExpr cols a -> TExpr cols a
sqrt (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
sqrt Expr a
a)
    log :: TExpr cols a -> TExpr cols a
log (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
log Expr a
a)
    (TExpr Expr a
a) ** :: TExpr cols a -> TExpr cols a -> TExpr cols a
** (TExpr Expr a
b) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a
a Expr a -> Expr a -> Expr a
forall a. Floating a => a -> a -> a
** Expr a
b)
    logBase :: TExpr cols a -> TExpr cols a -> TExpr cols a
logBase (TExpr Expr a
a) (TExpr Expr a
b) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a -> Expr a
forall a. Floating a => a -> a -> a
logBase Expr a
a Expr a
b)
    sin :: TExpr cols a -> TExpr cols a
sin (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
sin Expr a
a)
    cos :: TExpr cols a -> TExpr cols a
cos (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
cos Expr a
a)
    tan :: TExpr cols a -> TExpr cols a
tan (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
tan Expr a
a)
    asin :: TExpr cols a -> TExpr cols a
asin (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
asin Expr a
a)
    acos :: TExpr cols a -> TExpr cols a
acos (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
acos Expr a
a)
    atan :: TExpr cols a -> TExpr cols a
atan (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
atan Expr a
a)
    sinh :: TExpr cols a -> TExpr cols a
sinh (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
sinh Expr a
a)
    cosh :: TExpr cols a -> TExpr cols a
cosh (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
cosh Expr a
a)
    asinh :: TExpr cols a -> TExpr cols a
asinh (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
asinh Expr a
a)
    acosh :: TExpr cols a -> TExpr cols a
acosh (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
acosh Expr a
a)
    atanh :: TExpr cols a -> TExpr cols a
atanh (TExpr Expr a
a) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> Expr a
forall a. Floating a => a -> a
atanh Expr a
a)

instance (IsString a, Columnable a) => IsString (TExpr cols a) where
    fromString :: String -> TExpr cols a
fromString = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (Expr a -> TExpr cols a)
-> (String -> Expr a) -> String -> TExpr cols a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Expr a
forall a. IsString a => String -> a
fromString

-------------------------------------------------------------------------------
-- Lifting arbitrary functions
-------------------------------------------------------------------------------

-- | Lift a unary function into a typed expression.
lift ::
    (Columnable a, Columnable b) => (a -> b) -> TExpr cols a -> TExpr cols b
lift :: forall a b (cols :: [*]).
(Columnable a, Columnable b) =>
(a -> b) -> TExpr cols a -> TExpr cols b
lift a -> b
f (TExpr Expr a
e) = Expr b -> TExpr cols b
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (UnaryOp a b -> Expr a -> Expr b
forall a b.
(Columnable a, Columnable b) =>
UnaryOp b a -> Expr b -> Expr a
Unary ((a -> b) -> Text -> Maybe Text -> UnaryOp a b
forall a b. (a -> b) -> Text -> Maybe Text -> UnaryOp a b
MkUnaryOp a -> b
f Text
"unaryUdf" Maybe Text
forall a. Maybe a
Nothing) Expr a
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 :: forall a b c (cols :: [*]).
(Columnable a, Columnable b, Columnable c) =>
(a -> b -> c) -> TExpr cols a -> TExpr cols b -> TExpr cols c
lift2 a -> b -> c
f (TExpr Expr a
a) (TExpr Expr b
b) = Expr c -> TExpr cols c
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (BinaryOp a b c -> Expr a -> Expr b -> Expr c
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary ((a -> b -> c)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a b c
forall a b c.
(a -> b -> c)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a b c
MkBinaryOp a -> b -> c
f Text
"binaryUdf" Maybe Text
forall a. Maybe a
Nothing Bool
False Int
0) Expr a
a Expr b
b)

infixl 4 .==., ./=., .<., .<=., .>=., .>.
infixr 3 .&&.
infixr 2 .||.

(.==.) ::
    (Columnable a, Eq a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool
.==. :: forall a (cols :: [*]).
(Columnable a, Eq a) =>
TExpr cols a -> TExpr cols a -> TExpr cols Bool
(.==.) (TExpr Expr a
a) (TExpr Expr a
b) = Expr Bool -> TExpr cols Bool
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (BinaryOp a a Bool -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary ((a -> a -> Bool)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a a Bool
forall a b c.
(a -> b -> c)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a b c
MkBinaryOp 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) Expr a
a Expr a
b)

(./=.) ::
    (Columnable a, Eq a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool
./=. :: forall a (cols :: [*]).
(Columnable a, Eq a) =>
TExpr cols a -> TExpr cols a -> TExpr cols Bool
(./=.) (TExpr Expr a
a) (TExpr Expr a
b) = Expr Bool -> TExpr cols Bool
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (BinaryOp a a Bool -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary ((a -> a -> Bool)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a a Bool
forall a b c.
(a -> b -> c)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a b c
MkBinaryOp a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(/=) Text
"neq" (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"/=") Bool
True Int
4) Expr a
a Expr a
b)

(.<.) ::
    (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool
.<. :: forall a (cols :: [*]).
(Columnable a, Ord a) =>
TExpr cols a -> TExpr cols a -> TExpr cols Bool
(.<.) (TExpr Expr a
a) (TExpr Expr a
b) = Expr Bool -> TExpr cols Bool
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (BinaryOp a a Bool -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary ((a -> a -> Bool)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a a Bool
forall a b c.
(a -> b -> c)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a b c
MkBinaryOp 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) Expr a
a Expr a
b)

(.<=.) ::
    (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool
.<=. :: forall a (cols :: [*]).
(Columnable a, Ord a) =>
TExpr cols a -> TExpr cols a -> TExpr cols Bool
(.<=.) (TExpr Expr a
a) (TExpr Expr a
b) = Expr Bool -> TExpr cols Bool
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (BinaryOp a a Bool -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary ((a -> a -> Bool)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a a Bool
forall a b c.
(a -> b -> c)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a b c
MkBinaryOp 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) Expr a
a Expr a
b)

(.>=.) ::
    (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool
.>=. :: forall a (cols :: [*]).
(Columnable a, Ord a) =>
TExpr cols a -> TExpr cols a -> TExpr cols Bool
(.>=.) (TExpr Expr a
a) (TExpr Expr a
b) = Expr Bool -> TExpr cols Bool
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (BinaryOp a a Bool -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary ((a -> a -> Bool)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a a Bool
forall a b c.
(a -> b -> c)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a b c
MkBinaryOp 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) Expr a
a Expr a
b)

(.>.) ::
    (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool
.>. :: forall a (cols :: [*]).
(Columnable a, Ord a) =>
TExpr cols a -> TExpr cols a -> TExpr cols Bool
(.>.) (TExpr Expr a
a) (TExpr Expr a
b) = Expr Bool -> TExpr cols Bool
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (BinaryOp a a Bool -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary ((a -> a -> Bool)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a a Bool
forall a b c.
(a -> b -> c)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a b c
MkBinaryOp 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) Expr a
a Expr a
b)

(.&&.) :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool
.&&. :: forall (cols :: [*]).
TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool
(.&&.) (TExpr Expr Bool
a) (TExpr Expr Bool
b) = Expr Bool -> TExpr cols Bool
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (BinaryOp Bool Bool Bool -> Expr Bool -> Expr Bool -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary ((Bool -> Bool -> Bool)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp Bool Bool Bool
forall a b c.
(a -> b -> c)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a b c
MkBinaryOp Bool -> Bool -> Bool
(&&) Text
"and" (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"&&") Bool
True Int
3) Expr Bool
a Expr Bool
b)

(.||.) :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool
.||. :: forall (cols :: [*]).
TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool
(.||.) (TExpr Expr Bool
a) (TExpr Expr Bool
b) = Expr Bool -> TExpr cols Bool
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (BinaryOp Bool Bool Bool -> Expr Bool -> Expr Bool -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary ((Bool -> Bool -> Bool)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp Bool Bool Bool
forall a b c.
(a -> b -> c)
-> Text -> Maybe Text -> Bool -> Int -> BinaryOp a b c
MkBinaryOp Bool -> Bool -> Bool
(||) Text
"or" (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"||") Bool
True Int
2) Expr Bool
a Expr Bool
b)

not :: TExpr cols Bool -> TExpr cols Bool
not :: forall (cols :: [*]). TExpr cols Bool -> TExpr cols Bool
not (TExpr Expr Bool
e) = Expr Bool -> TExpr cols Bool
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (UnaryOp Bool Bool -> Expr Bool -> Expr Bool
forall a b.
(Columnable a, Columnable b) =>
UnaryOp b a -> Expr b -> Expr a
Unary ((Bool -> Bool) -> Text -> Maybe Text -> UnaryOp Bool Bool
forall a b. (a -> b) -> Text -> Maybe Text -> UnaryOp a b
MkUnaryOp Bool -> Bool
Prelude.not Text
"not" (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"!")) Expr Bool
e)

-------------------------------------------------------------------------------
-- Aggregation combinators
-------------------------------------------------------------------------------

sum :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a
sum :: forall a (cols :: [*]).
(Columnable a, Num a) =>
TExpr cols a -> TExpr cols a
sum (TExpr Expr a
e) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (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
(+)) Expr a
e)

mean :: (Columnable a, Real a) => TExpr cols a -> TExpr cols Double
mean :: forall a (cols :: [*]).
(Columnable a, Real a) =>
TExpr cols a -> TExpr cols Double
mean (TExpr Expr a
e) =
    Expr Double -> TExpr cols Double
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr
        ( 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)
            )
            Expr a
e
        )

count :: (Columnable a) => TExpr cols a -> TExpr cols Int
count :: forall a (cols :: [*]).
Columnable a =>
TExpr cols a -> TExpr cols Int
count (TExpr Expr a
e) = Expr Int -> TExpr cols Int
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (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) Expr a
e)

minimum :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a
minimum :: forall a (cols :: [*]).
(Columnable a, Ord a) =>
TExpr cols a -> TExpr cols a
minimum (TExpr Expr a
e) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (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
min) Expr a
e)

maximum :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a
maximum :: forall a (cols :: [*]).
(Columnable a, Ord a) =>
TExpr cols a -> TExpr cols a
maximum (TExpr Expr a
e) = Expr a -> TExpr cols a
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (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
max) Expr a
e)

collect :: (Columnable a) => TExpr cols a -> TExpr cols [a]
collect :: forall a (cols :: [*]).
Columnable a =>
TExpr cols a -> TExpr cols [a]
collect (TExpr Expr a
e) = Expr [a] -> TExpr cols [a]
forall (cols :: [*]) a. Expr a -> TExpr cols a
TExpr (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 (:))) Expr a
e)

-------------------------------------------------------------------------------
-- Named expression helper
-------------------------------------------------------------------------------

-- | Create a 'NamedExpr' for use with 'aggregateUntyped'.
as :: (Columnable a) => TExpr cols a -> T.Text -> NamedExpr
as :: forall a (cols :: [*]).
Columnable a =>
TExpr cols a -> Text -> NamedExpr
as (TExpr Expr a
e) Text
name = (Text
name, Expr a -> UExpr
forall a. Columnable a => Expr a -> UExpr
UExpr Expr a
e)

-- | Create an ascending sort order from a typed expression.
asc :: (Columnable a) => TExpr cols a -> TSortOrder cols
asc :: forall a (cols :: [*]).
Columnable a =>
TExpr cols a -> TSortOrder cols
asc = TExpr cols a -> TSortOrder cols
forall a (cols :: [*]).
Columnable a =>
TExpr cols a -> TSortOrder cols
Asc

-- | Create a descending sort order from a typed expression.
desc :: (Columnable a) => TExpr cols a -> TSortOrder cols
desc :: forall a (cols :: [*]).
Columnable a =>
TExpr cols a -> TSortOrder cols
desc = TExpr cols a -> TSortOrder cols
forall a (cols :: [*]).
Columnable a =>
TExpr cols a -> TSortOrder cols
Desc