{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UndecidableInstances #-}

module DataFrame.Functions where

import DataFrame.Internal.Column
import DataFrame.Internal.DataFrame (
    DataFrame (..),
    unsafeGetColumn,
 )
import DataFrame.Internal.Expression (
    Expr (..),
    NamedExpr,
    UExpr (..),
 )
import DataFrame.Internal.Statistics

import Control.Monad
import Control.Monad.IO.Class
import qualified Data.Char as Char
import Data.Function
import Data.Functor
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 qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import qualified DataFrame.IO.CSV as CSV
import qualified DataFrame.IO.Parquet as Parquet
import Debug.Trace (trace)
import Language.Haskell.TH
import qualified Language.Haskell.TH.Syntax as TH
import Text.Regex.TDFA
import Prelude hiding (maximum, minimum)
import Prelude as P

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

name :: (Show a) => Expr a -> T.Text
name :: forall a. Show a => Expr a -> Text
name (Col Text
n) = Text
n
name Expr a
other =
    [Char] -> Text
forall a. HasCallStack => [Char] -> a
error ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$
        [Char]
"You must call `name` on a column reference. Not the expression: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Expr a -> [Char]
forall a. Show a => a -> [Char]
show Expr a
other

col :: (Columnable a) => T.Text -> Expr a
col :: forall a. Columnable a => Text -> Expr a
col = Text -> Expr a
forall a. Columnable a => Text -> Expr a
Col

as :: (Columnable a) => Expr a -> T.Text -> NamedExpr
as :: forall a. Columnable a => Expr a -> Text -> NamedExpr
as Expr a
expr Text
name = (Text
name, Expr a -> UExpr
forall a. Columnable a => Expr a -> UExpr
Wrap Expr a
expr)

infixr 0 .=
(.=) :: (Columnable a) => T.Text -> Expr a -> NamedExpr
.= :: forall a. Columnable a => Text -> Expr a -> NamedExpr
(.=) = (Expr a -> Text -> NamedExpr) -> Text -> Expr a -> NamedExpr
forall a b c. (a -> b -> c) -> b -> a -> c
flip Expr a -> Text -> NamedExpr
forall a. Columnable a => Expr a -> Text -> NamedExpr
as

ifThenElse :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a
ifThenElse :: forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
ifThenElse = Expr Bool -> Expr a -> Expr a -> Expr a
forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
If

lit :: (Columnable a) => a -> Expr a
lit :: forall a. Columnable a => a -> Expr a
lit = a -> Expr a
forall a. Columnable a => a -> Expr a
Lit

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 = Text -> (a -> b) -> Expr a -> Expr b
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
UnaryOp Text
"udf"

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 = Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"udf"

toDouble :: (Columnable a, Real a) => Expr a -> Expr Double
toDouble :: forall a. (Columnable a, Real a) => Expr a -> Expr Double
toDouble = Text -> (a -> Double) -> Expr a -> Expr Double
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
UnaryOp Text
"toDouble" a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac

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 = Text -> (a -> a -> a) -> Expr a -> Expr a -> Expr a
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"div" a -> a -> a
forall a. Integral a => a -> a -> a
Prelude.div

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 = Text -> (a -> a -> a) -> Expr a -> Expr a -> Expr a
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"mod" a -> a -> a
forall a. Integral a => a -> a -> a
Prelude.mod

(.==) :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
.== :: forall a. (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
(.==) = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"eq" a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==)

(./=) :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
./= :: forall a. (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
(./=) = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"eq" a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(/=)

eq :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
eq :: forall a. (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
eq = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"eq" a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==)

(.<) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
.< :: forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
(.<) = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"lt" a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<)

lt :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
lt :: forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
lt = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"lt" a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<)

-- TODO: Generalize this pattern for other equality functions.
(.>) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
.> :: forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
(.>) = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"gt" a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>)

gt :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
gt :: forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
gt = Expr a -> Expr a -> Expr Bool
forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
(.>)

(.<=) :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool
.<= :: forall a.
(Columnable a, Ord a, Eq a) =>
Expr a -> Expr a -> Expr Bool
(.<=) = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"leq" a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<=)

leq :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool
leq :: forall a.
(Columnable a, Ord a, Eq a) =>
Expr a -> Expr a -> Expr Bool
leq = Expr a -> Expr a -> Expr Bool
forall a.
(Columnable a, Ord a, Eq a) =>
Expr a -> Expr a -> Expr Bool
(.<=)

(.>=) :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool
.>= :: forall a.
(Columnable a, Ord a, Eq a) =>
Expr a -> Expr a -> Expr Bool
(.>=) = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"geq" a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>=)

geq :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool
geq :: forall a.
(Columnable a, Ord a, Eq a) =>
Expr a -> Expr a -> Expr Bool
geq = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"geq" a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>=)

and :: Expr Bool -> Expr Bool -> Expr Bool
and :: Expr Bool -> Expr Bool -> Expr Bool
and = Text
-> (Bool -> Bool -> Bool) -> Expr Bool -> Expr Bool -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"and" Bool -> Bool -> Bool
(&&)

(.&&) :: Expr Bool -> Expr Bool -> Expr Bool
.&& :: Expr Bool -> Expr Bool -> Expr Bool
(.&&) = Text
-> (Bool -> Bool -> Bool) -> Expr Bool -> Expr Bool -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"and" Bool -> Bool -> Bool
(&&)

or :: Expr Bool -> Expr Bool -> Expr Bool
or :: Expr Bool -> Expr Bool -> Expr Bool
or = Text
-> (Bool -> Bool -> Bool) -> Expr Bool -> Expr Bool -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"or" Bool -> Bool -> Bool
(||)

(.||) :: Expr Bool -> Expr Bool -> Expr Bool
.|| :: Expr Bool -> Expr Bool -> Expr Bool
(.||) = Text
-> (Bool -> Bool -> Bool) -> Expr Bool -> Expr Bool -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"or" Bool -> Bool -> Bool
(||)

not :: Expr Bool -> Expr Bool
not :: Expr Bool -> Expr Bool
not = Text -> (Bool -> Bool) -> Expr Bool -> Expr Bool
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
UnaryOp Text
"not" Bool -> Bool
Prelude.not

count :: (Columnable a) => Expr a -> Expr Int
count :: forall a. Columnable a => Expr a -> Expr Int
count Expr a
expr = Expr a -> Text -> Int -> (Int -> a -> Int) -> Expr Int
forall a b.
(Columnable a, Columnable b) =>
Expr b -> Text -> a -> (a -> b -> a) -> Expr a
AggFold Expr a
expr Text
"count" Int
0 (\Int
acc a
_ -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

collect :: (Columnable a) => Expr a -> Expr [a]
collect :: forall a. Columnable a => Expr a -> Expr [a]
collect Expr a
expr = Expr a -> Text -> [a] -> ([a] -> a -> [a]) -> Expr [a]
forall a b.
(Columnable a, Columnable b) =>
Expr b -> Text -> a -> (a -> b -> a) -> Expr a
AggFold Expr a
expr Text
"collect" [] ((a -> [a] -> [a]) -> [a] -> a -> [a]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:))

mode :: (Columnable a, Eq a) => Expr a -> Expr a
mode :: forall a. (Columnable a, Eq a) => Expr a -> Expr a
mode Expr a
expr =
    Expr a -> Text -> (Vector a -> a) -> Expr a
forall (v :: * -> *) b a.
(Vector v b, Typeable v, Columnable a, Columnable b) =>
Expr b -> Text -> (v b -> a) -> Expr a
AggVector
        Expr a
expr
        Text
"mode"
        ( (a, Integer) -> a
forall a b. (a, b) -> a
fst
            ((a, Integer) -> a) -> (Vector a -> (a, Integer)) -> Vector a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, Integer) -> (a, Integer) -> Ordering)
-> [(a, Integer)] -> (a, Integer)
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
L.maximumBy (Integer -> Integer -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Integer -> Integer -> Ordering)
-> ((a, Integer) -> Integer)
-> (a, Integer)
-> (a, Integer)
-> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (a, Integer) -> Integer
forall a b. (a, b) -> b
snd)
            ([(a, Integer)] -> (a, Integer))
-> (Vector a -> [(a, Integer)]) -> Vector a -> (a, Integer)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map a Integer -> [(a, Integer)]
forall k a. Map k a -> [(k, a)]
M.toList
            (Map a Integer -> [(a, Integer)])
-> (Vector a -> Map a Integer) -> Vector a -> [(a, Integer)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map a Integer -> a -> Map a Integer)
-> Map a Integer -> Vector a -> Map a Integer
forall a b. (a -> b -> a) -> a -> Vector b -> a
V.foldl' (\Map a Integer
m a
e -> (Integer -> Integer -> Integer)
-> a -> Integer -> Map a Integer -> Map a Integer
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+) a
e Integer
1 Map a Integer
m) Map a Integer
forall k a. Map k a
M.empty
        )

minimum :: (Columnable a, Ord a) => Expr a -> Expr a
minimum :: forall a. (Columnable a, Ord a) => Expr a -> Expr a
minimum Expr a
expr = Expr a -> Text -> (a -> a -> a) -> Expr a
forall a. Columnable a => Expr a -> Text -> (a -> a -> a) -> Expr a
AggReduce Expr a
expr Text
"minimum" a -> a -> a
forall a. Ord a => a -> a -> a
Prelude.min

maximum :: (Columnable a, Ord a) => Expr a -> Expr a
maximum :: forall a. (Columnable a, Ord a) => Expr a -> Expr a
maximum Expr a
expr = Expr a -> Text -> (a -> a -> a) -> Expr a
forall a. Columnable a => Expr a -> Text -> (a -> a -> a) -> Expr a
AggReduce Expr a
expr Text
"maximum" a -> a -> a
forall a. Ord a => a -> a -> a
Prelude.max

sum :: forall a. (Columnable a, Num a) => Expr a -> Expr a
sum :: forall a. (Columnable a, Num a) => Expr a -> Expr a
sum Expr a
expr = Expr a -> Text -> (a -> a -> a) -> Expr a
forall a. Columnable a => Expr a -> Text -> (a -> a -> a) -> Expr a
AggReduce Expr a
expr Text
"sum" a -> a -> a
forall a. Num a => a -> a -> a
(+)
{-# SPECIALIZE DataFrame.Functions.sum :: Expr Double -> Expr Double #-}
{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int -> Expr Int #-}
{-# 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 Expr (Maybe a)
expr = Expr (Maybe a) -> Text -> (Vector (Maybe a) -> a) -> Expr a
forall (v :: * -> *) b a.
(Vector v b, Typeable v, Columnable a, Columnable b) =>
Expr b -> Text -> (v b -> a) -> Expr a
AggVector Expr (Maybe a)
expr 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)

mean :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double
mean :: forall a. (Columnable a, Real a, Unbox a) => Expr a -> Expr Double
mean Expr a
expr = Expr a -> Text -> (Vector a -> Double) -> Expr Double
forall a b.
(Columnable a, Columnable b, Unbox a, Unbox b, Num a, Num b) =>
Expr b -> Text -> (Vector b -> a) -> Expr a
AggNumericVector Expr a
expr Text
"mean" Vector a -> Double
forall a. (Real a, Unbox a) => Vector a -> Double
mean'
{-# SPECIALIZE DataFrame.Functions.mean :: Expr Double -> Expr Double #-}
{-# SPECIALIZE DataFrame.Functions.mean :: Expr Int -> Expr Double #-}
{-# INLINEABLE DataFrame.Functions.mean #-}

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 Expr (Maybe a)
expr = Expr (Maybe a)
-> Text -> (Vector (Maybe a) -> Double) -> Expr Double
forall (v :: * -> *) b a.
(Vector v b, Typeable v, Columnable a, Columnable b) =>
Expr b -> Text -> (v b -> a) -> Expr a
AggVector Expr (Maybe a)
expr 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)

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 Expr a
expr = Expr a -> Text -> (Vector a -> Double) -> Expr Double
forall a b.
(Columnable a, Columnable b, Unbox a, Unbox b, Num a, Num b) =>
Expr b -> Text -> (Vector b -> a) -> Expr a
AggNumericVector Expr a
expr Text
"variance" Vector a -> Double
forall a. (Real a, Unbox a) => Vector a -> Double
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 Expr a
expr = Expr a -> Text -> (Vector a -> Double) -> Expr Double
forall a b.
(Columnable a, Columnable b, Unbox a, Unbox b, Num a, Num b) =>
Expr b -> Text -> (Vector b -> a) -> Expr a
AggNumericVector Expr a
expr Text
"median" Vector a -> Double
forall a. (Real a, Unbox a) => Vector a -> Double
median'

medianMaybe :: (Columnable a, Real a) => Expr (Maybe a) -> Expr Double
medianMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double
medianMaybe Expr (Maybe a)
expr = Expr (Maybe a)
-> Text -> (Vector (Maybe a) -> Double) -> Expr Double
forall (v :: * -> *) b a.
(Vector v b, Typeable v, Columnable a, Columnable b) =>
Expr b -> Text -> (v b -> a) -> Expr a
AggVector Expr (Maybe a)
expr 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)

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 Expr Double
expr =
    Expr Double -> Text -> (Vector Double -> Double) -> Expr Double
forall a b.
(Columnable a, Columnable b, Unbox a, Unbox b, Num a, Num b) =>
Expr b -> Text -> (Vector b -> a) -> Expr a
AggNumericVector
        Expr Double
expr
        ([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 Expr a
expr = Expr a -> Text -> (Vector a -> Double) -> Expr Double
forall a b.
(Columnable a, Columnable b, Unbox a, Unbox b, Num a, Num b) =>
Expr b -> Text -> (Vector b -> a) -> Expr a
AggNumericVector Expr a
expr 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')

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 Expr (Maybe a)
expr = Expr (Maybe a)
-> Text -> (Vector (Maybe a) -> Double) -> Expr Double
forall (v :: * -> *) b a.
(Vector v b, Typeable v, Columnable a, Columnable b) =>
Expr b -> Text -> (v b -> a) -> Expr a
AggVector Expr (Maybe a)
expr 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)

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, Unbox 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
_ Int
0 = a -> Expr a
forall a. Columnable a => a -> Expr a
Lit a
1
pow (Lit a
n) Int
i = a -> Expr a
forall a. Columnable a => a -> Expr a
Lit (a
n a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
i)
pow Expr a
expr Int
1 = Expr a
expr
pow Expr a
expr Int
i = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
UnaryOp (Text
"pow " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i)) (a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
i) Expr a
expr

relu :: (Columnable a, Num a) => Expr a -> Expr a
relu :: forall a. (Columnable a, Num a) => Expr a -> Expr a
relu = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
UnaryOp Text
"relu" (a -> a -> a
forall a. Ord a => a -> a -> a
Prelude.max a
0)

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 = Text -> (a -> a -> a) -> Expr a -> Expr a -> Expr a
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"min" a -> a -> a
forall a. Ord a => a -> a -> a
Prelude.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 = Text -> (a -> a -> a) -> Expr a -> Expr a -> Expr a
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinaryOp Text
"max" a -> a -> a
forall a. Ord a => a -> a -> a
Prelude.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 = Expr b -> Text -> a -> (a -> b -> a) -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Expr b -> Text -> a -> (a -> b -> a) -> Expr a
AggFold Expr b
expr Text
"foldUdf"

toMaybe :: (Columnable a) => Expr a -> Expr (Maybe a)
toMaybe :: forall a. Columnable a => Expr a -> Expr (Maybe a)
toMaybe = Text -> (a -> Maybe a) -> Expr a -> Expr (Maybe a)
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
UnaryOp Text
"toMaybe" a -> Maybe a
forall a. a -> Maybe a
Just

fromMaybe :: (Columnable a) => a -> Expr (Maybe a) -> Expr a
fromMaybe :: forall a. Columnable a => a -> Expr (Maybe a) -> Expr a
fromMaybe a
d = Text -> (Maybe a -> a) -> Expr (Maybe a) -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
UnaryOp (Text
"fromMaybe " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack (a -> [Char]
forall a. Show a => a -> [Char]
show a
d)) (a -> Maybe a -> a
forall a. a -> Maybe a -> a
Maybe.fromMaybe a
d)

isJust :: (Columnable a) => Expr (Maybe a) -> Expr Bool
isJust :: forall a. Columnable a => Expr (Maybe a) -> Expr Bool
isJust = Text -> (Maybe a -> Bool) -> Expr (Maybe a) -> Expr Bool
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
UnaryOp Text
"isJust" Maybe a -> Bool
forall a. Maybe a -> Bool
Maybe.isJust

isNothing :: (Columnable a) => Expr (Maybe a) -> Expr Bool
isNothing :: forall a. Columnable a => Expr (Maybe a) -> Expr Bool
isNothing = Text -> (Maybe a -> Bool) -> Expr (Maybe a) -> Expr Bool
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
UnaryOp Text
"isNothing" Maybe a -> Bool
forall a. Maybe a -> Bool
Maybe.isNothing

fromJust :: (Columnable a) => Expr (Maybe a) -> Expr a
fromJust :: forall a. Columnable a => Expr (Maybe a) -> Expr a
fromJust = Text -> (Maybe a -> a) -> Expr (Maybe a) -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
UnaryOp Text
"fromJust" Maybe a -> a
forall a. HasCallStack => Maybe a -> a
Maybe.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) -> Expr (Maybe a) -> Expr (Maybe b)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Expr a -> Expr b
lift ((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)

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)
-> Expr (Maybe a) -> Expr (Maybe b) -> Expr (Maybe c)
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a) -> Expr c -> Expr b -> Expr a
lift2 (\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)

recode ::
    forall a b.
    (Columnable a, Columnable b) => [(a, b)] -> Expr a -> Expr (Maybe b)
recode :: forall a b.
(Columnable a, Columnable b) =>
[(a, b)] -> Expr a -> Expr (Maybe b)
recode [(a, b)]
mapping = Text -> (a -> Maybe b) -> Expr a -> Expr (Maybe b)
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
UnaryOp ([Char] -> Text
T.pack ([Char]
"recode " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [(a, b)] -> [Char]
forall a. Show a => a -> [Char]
show [(a, b)]
mapping)) (a -> [(a, b)] -> Maybe b
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(a, b)]
mapping)

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
value = Expr b
fallback
recodeWithCondition Expr b
fallback ((Expr a -> Expr Bool
cond, b
value) : [(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
value) (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) => b -> [(a, b)] -> Expr a -> Expr b
recodeWithDefault :: forall a b.
(Columnable a, Columnable b) =>
b -> [(a, b)] -> Expr a -> Expr b
recodeWithDefault b
d [(a, b)]
mapping =
    Text -> (a -> b) -> Expr a -> Expr b
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
UnaryOp
        ([Char] -> Text
T.pack ([Char]
"recode " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [(a, b)] -> [Char]
forall a. Show a => a -> [Char]
show [(a, b)]
mapping))
        (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))

firstOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a)
firstOrNothing :: forall a. Columnable a => Expr [a] -> Expr (Maybe a)
firstOrNothing = ([a] -> Maybe a) -> Expr [a] -> Expr (Maybe a)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Expr a -> Expr b
lift [a] -> Maybe a
forall a. [a] -> Maybe a
Maybe.listToMaybe

lastOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a)
lastOrNothing :: forall a. Columnable a => Expr [a] -> Expr (Maybe a)
lastOrNothing = ([a] -> Maybe a) -> Expr [a] -> Expr (Maybe a)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Expr a -> Expr b
lift ([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)

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

match :: T.Text -> Expr T.Text -> Expr (Maybe T.Text)
match :: Text -> Expr Text -> Expr (Maybe Text)
match Text
regex = (Text -> Maybe Text) -> Expr Text -> Expr (Maybe Text)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Expr a -> Expr b
lift ((\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))

matchAll :: T.Text -> Expr T.Text -> Expr [T.Text]
matchAll :: Text -> Expr Text -> Expr [Text]
matchAll Text
regex = (Text -> [Text]) -> Expr Text -> Expr [Text]
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Expr a -> Expr b
lift (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))

parseDate :: T.Text -> Expr T.Text -> Expr (Maybe Day)
parseDate :: Text -> Expr Text -> Expr (Maybe Day)
parseDate Text
format = (Text -> Maybe Day) -> Expr Text -> Expr (Maybe Day)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Expr a -> Expr b
lift (Bool -> TimeLocale -> [Char] -> [Char] -> Maybe Day
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 Day) -> (Text -> [Char]) -> Text -> Maybe Day
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Char]
T.unpack)

daysBetween :: Expr Day -> Expr Day -> Expr Int
daysBetween :: Expr Day -> Expr Day -> Expr Int
daysBetween Expr Day
d1 Expr Day
d2 = (Integer -> Int) -> Expr Integer -> Expr Int
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Expr a -> Expr b
lift Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Day -> Day -> Integer) -> Expr Day -> Expr Day -> Expr Integer
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a) -> Expr c -> Expr b -> Expr a
lift2 Day -> Day -> Integer
diffDays Expr Day
d1 Expr Day
d2)

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) -> Expr (m a) -> Expr (m b)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Expr a -> Expr b
lift (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)

-- 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

typeFromString :: [String] -> Q Type
typeFromString :: [[Char]] -> Q Type
typeFromString [] = [Char] -> Q Type
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"No type specified"
typeFromString [[Char]
t] = do
    Maybe Name
maybeType <- [Char] -> Q (Maybe Name)
lookupTypeName [Char]
t
    case Maybe Name
maybeType of
        Just Name
name -> Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> Type
ConT Name
name)
        Maybe Name
Nothing ->
            if Int -> [Char] -> [Char]
forall a. Int -> [a] -> [a]
take Int
1 [Char]
t [Char] -> [Char] -> Bool
forall a. Eq a => a -> a -> Bool
== [Char]
"["
                then [[Char]] -> Q Type
typeFromString [[Char] -> [Char]
forall a. [a] -> [a]
dropFirstAndLast [Char]
t] Q Type -> (Type -> Type) -> Q Type
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> Type -> Type -> Type
AppT Type
ListT
                else [Char] -> Q Type
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> Q Type) -> [Char] -> Q Type
forall a b. (a -> b) -> a -> b
$ [Char]
"Unsupported type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
t
typeFromString [[Char]
tycon, [Char]
t1] = do
    Type
outer <- [[Char]] -> Q Type
typeFromString [[Char]
tycon]
    Type
inner <- [[Char]] -> Q Type
typeFromString [[Char]
t1]
    Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Type -> Type
AppT Type
outer Type
inner)
typeFromString [[Char]
tycon, [Char]
t1, [Char]
t2] = do
    Type
outer <- [[Char]] -> Q Type
typeFromString [[Char]
tycon]
    Type
lhs <- [[Char]] -> Q Type
typeFromString [[Char]
t1]
    Type
rhs <- [[Char]] -> Q Type
typeFromString [[Char]
t2]
    Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Type -> Type
AppT (Type -> Type -> Type
AppT Type
outer Type
lhs) Type
rhs)
typeFromString [[Char]]
s = [Char] -> Q Type
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> Q Type) -> [Char] -> Q Type
forall a b. (a -> b) -> a -> b
$ [Char]
"Unsupported types: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [[Char]] -> [Char]
unwords [[Char]]
s

dropFirstAndLast :: [a] -> [a]
dropFirstAndLast :: forall a. [a] -> [a]
dropFirstAndLast = [a] -> [a]
forall a. [a] -> [a]
reverse ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
1 ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. [a] -> [a]
reverse ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
1

declareColumnsFromCsvFile :: String -> DecsQ
declareColumnsFromCsvFile :: [Char] -> DecsQ
declareColumnsFromCsvFile [Char]
path = do
    DataFrame
df <-
        IO DataFrame -> Q DataFrame
forall a. IO a -> Q a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
            (ReadOptions -> [Char] -> IO DataFrame
CSV.readSeparated (ReadOptions
CSV.defaultReadOptions{CSV.numColumns = Just 100}) [Char]
path)
    DataFrame -> DecsQ
declareColumns DataFrame
df

-- TODO: We don't have to read the whole file, we can just read the schema.
declareColumnsFromParquetFile :: String -> DecsQ
declareColumnsFromParquetFile :: [Char] -> DecsQ
declareColumnsFromParquetFile [Char]
path = do
    DataFrame
df <- IO DataFrame -> Q DataFrame
forall a. IO a -> Q a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO ([Char] -> IO DataFrame
Parquet.readParquet [Char]
path)
    DataFrame -> DecsQ
declareColumns DataFrame
df

declareColumnsFromCsvWithOpts :: CSV.ReadOptions -> String -> DecsQ
declareColumnsFromCsvWithOpts :: ReadOptions -> [Char] -> DecsQ
declareColumnsFromCsvWithOpts ReadOptions
opts [Char]
path = do
    DataFrame
df <- IO DataFrame -> Q DataFrame
forall a. IO a -> Q a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ReadOptions -> [Char] -> IO DataFrame
CSV.readSeparated ReadOptions
opts [Char]
path)
    DataFrame -> DecsQ
declareColumns DataFrame
df

declareColumns :: DataFrame -> DecsQ
declareColumns :: DataFrame -> DecsQ
declareColumns DataFrame
df =
    let
        names :: [Text]
names = (((Text, Int) -> Text) -> [(Text, Int)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Int) -> Text
forall a b. (a, b) -> a
fst ([(Text, Int)] -> [Text])
-> (DataFrame -> [(Text, Int)]) -> DataFrame -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Text, Int) -> (Text, Int) -> Ordering)
-> [(Text, Int)] -> [(Text, Int)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
L.sortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> ((Text, Int) -> Int) -> (Text, Int) -> (Text, Int) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (Text, Int) -> Int
forall a b. (a, b) -> b
snd) ([(Text, Int)] -> [(Text, Int)])
-> (DataFrame -> [(Text, Int)]) -> DataFrame -> [(Text, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Text Int -> [(Text, Int)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Text Int -> [(Text, Int)])
-> (DataFrame -> Map Text Int) -> DataFrame -> [(Text, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame -> Map Text Int
columnIndices) DataFrame
df
        types :: [[Char]]
types = (Text -> [Char]) -> [Text] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map (Column -> [Char]
columnTypeString (Column -> [Char]) -> (Text -> Column) -> Text -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> DataFrame -> Column
`unsafeGetColumn` DataFrame
df)) [Text]
names
        specs :: [(Text, Text, [Char])]
specs = (Text -> [Char] -> (Text, Text, [Char]))
-> [Text] -> [[Char]] -> [(Text, Text, [Char])]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Text
name [Char]
type_ -> (Text
name, Text -> Text
sanitize Text
name, [Char]
type_)) [Text]
names [[Char]]
types
     in
        ([[Dec]] -> [Dec]) -> Q [[Dec]] -> DecsQ
forall a b. (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Q [[Dec]] -> DecsQ) -> Q [[Dec]] -> DecsQ
forall a b. (a -> b) -> a -> b
$ [(Text, Text, [Char])]
-> ((Text, Text, [Char]) -> DecsQ) -> Q [[Dec]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Text, Text, [Char])]
specs (((Text, Text, [Char]) -> DecsQ) -> Q [[Dec]])
-> ((Text, Text, [Char]) -> DecsQ) -> Q [[Dec]]
forall a b. (a -> b) -> a -> b
$ \(Text
raw, Text
nm, [Char]
tyStr) -> do
            Type
ty <- [[Char]] -> Q Type
typeFromString ([Char] -> [[Char]]
words [Char]
tyStr)
            [Char] -> (() -> Q ()) -> () -> Q ()
forall a. [Char] -> a -> a
trace (Text -> [Char]
T.unpack (Text
nm Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" :: Expr " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
T.pack [Char]
tyStr)) () -> Q ()
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            let n :: Name
n = [Char] -> Name
mkName (Text -> [Char]
T.unpack Text
nm)
            Dec
sig <- Name -> Q Type -> Q Dec
forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD Name
n [t|Expr $(Type -> Q Type
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
ty)|]
            Dec
val <- Q Pat -> Q Body -> [Q Dec] -> Q Dec
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
n) (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [|col $(Text -> Q Exp
forall t (m :: * -> *). (Lift t, Quote m) => t -> m Exp
forall (m :: * -> *). Quote m => Text -> m Exp
TH.lift Text
raw)|]) []
            [Dec] -> DecsQ
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Dec
sig, Dec
val]