dataframe-0.3.0.2: src/DataFrame/Internal/Expression.hs
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE BangPatterns #-}
module DataFrame.Internal.Expression where
import qualified Data.Map as M
import Data.Type.Equality (type (:~:)(Refl), TestEquality (testEquality))
import Data.Data (Typeable)
import DataFrame.Internal.Column
import DataFrame.Internal.DataFrame
import DataFrame.Internal.Types
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Generic as VG
import Type.Reflection (typeRep)
import DataFrame.Errors (DataFrameException(ColumnNotFoundException))
import Control.Exception (throw)
import Data.Maybe (fromMaybe)
data Expr a where
Col :: Columnable a => T.Text -> Expr a
Lit :: Columnable a => a -> Expr a
Apply :: (Columnable a,
Columnable b)
=> T.Text -- Operation name
-> (b -> a)
-> Expr b
-> Expr a
BinOp :: (Columnable c,
Columnable b,
Columnable a)
=> T.Text -- operation name
-> (c -> b -> a)
-> Expr c
-> Expr b
-> Expr a
GeneralAggregate :: (Columnable a)
=> T.Text -- Column name
-> T.Text -- Operation name
-> (forall v b. (VG.Vector v b, Columnable b) => v b -> a)
-> Expr a
ReductionAggregate :: (Columnable a)
=> T.Text -- Column name
-> T.Text -- Operation name
-> (forall a . Columnable a => a -> a -> a)
-> Expr a
NumericAggregate :: (Columnable a,
Columnable b,
VU.Unbox a,
VU.Unbox b,
Num a,
Num b)
=> T.Text -- Column name
-> T.Text -- Operation name
-> (VU.Vector b -> a)
-> Expr a
data UExpr where
Wrap :: Columnable a => Expr a -> UExpr
interpret :: forall a . (Columnable a) => DataFrame -> Expr a -> TypedColumn a
interpret df (Lit value) = TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value
interpret df (Col name) = case getColumn name df of
Nothing -> throw $ ColumnNotFoundException name "" (map fst $ M.toList $ columnIndices df)
Just col -> TColumn col
interpret df (Apply _ (f :: c -> d) value) = let
(TColumn value') = interpret @c df value
-- TODO: Handle this gracefully.
in TColumn $ fromMaybe (error "mapColumn returned nothing") (mapColumn f value')
interpret df (BinOp _ (f :: c -> d -> e) left right) = let
(TColumn left') = interpret @c df left
(TColumn right') = interpret @d df right
in TColumn $ fromMaybe (error "mapColumn returned nothing") (zipWithColumns f left' right')
interpret df (ReductionAggregate name op (f :: forall a . Columnable a => a -> a -> a)) = let
(TColumn column) = interpret @a df (Col name)
in case headColumn @a column of
Nothing -> error "Invalid operation"
Just h -> case ifoldlColumn (\acc _ v -> f acc v) h column of
Nothing -> error "Invalid operation"
Just value -> TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value
interpret _ expr = error ("Invalid operation for dataframe: " ++ show expr)
interpretAggregation :: forall a . (Columnable a) => GroupedDataFrame -> Expr a -> TypedColumn a
interpretAggregation gdf (Lit value) = TColumn $ fromVector $ V.replicate (VG.length (offsets gdf) - 1) value
interpretAggregation gdf@(Grouped df names indices os) (Col name) = case getColumn name df of
Nothing -> throw $ ColumnNotFoundException name "" (map fst $ M.toList $ columnIndices df)
Just col -> TColumn $ atIndicesStable (VG.map (indices `VG.unsafeIndex`) (VG.init os)) col
interpretAggregation gdf (Apply _ (f :: c -> d) expr) = let
(TColumn value) = interpretAggregation @c gdf expr
in case mapColumn f value of
Nothing -> error "Type error in interpretation"
Just col -> TColumn col
interpretAggregation gdf (BinOp _ (f :: c -> d -> e) left right) = let
(TColumn left') = interpretAggregation @c gdf left
(TColumn right') = interpretAggregation @d gdf right
in case zipWithColumns f left' right' of
Nothing -> error "Type error in binary operation"
Just col -> TColumn col
interpretAggregation gdf@(Grouped df names indices os) (GeneralAggregate name op (f :: forall v b. (VG.Vector v b, Columnable b) => v b -> c)) = case getColumn name df of
Nothing -> throw $ ColumnNotFoundException name "" (map fst $ M.toList $ columnIndices df)
Just (BoxedColumn col) -> TColumn $ fromVector $
V.generate (VG.length os - 1)
(\i -> f (V.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i))
(\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i))))
)
)
Just (UnboxedColumn col) -> case sUnbox @c of
SFalse -> TColumn $ fromVector $
V.generate (VG.length os - 1)
(\i -> f (VU.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i))
(\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i))))
)
)
STrue -> TColumn $ fromUnboxedVector $
VU.generate (VG.length os - 1)
(\i -> f (VU.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i))
(\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i))))
)
)
Just (OptionalColumn col) -> TColumn $ fromVector $
V.generate (VG.length os - 1)
(\i -> f (V.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i))
(\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i))))
)
)
interpretAggregation gdf@(Grouped df names indices os) (ReductionAggregate name op (f :: forall a . Columnable a => a -> a -> a)) = case getColumn name df of
Nothing -> throw $ ColumnNotFoundException name "" (map fst $ M.toList $ columnIndices df)
Just (BoxedColumn col) -> TColumn $ fromVector $
VG.generate (VG.length os - 1) $ \g ->
let !start = os `VG.unsafeIndex` g
!end = os `VG.unsafeIndex` (g+1)
in go (col `VG.unsafeIndex` (indices `VG.unsafeIndex` start)) (start + 1) end
where
{-# INLINE go #-}
go !acc j e
| j == e = acc
| otherwise =
let !x = col `VG.unsafeIndex` (indices `VG.unsafeIndex` j)
in go (f acc x) (j + 1) e
Just (UnboxedColumn col) -> case sUnbox @a of
SFalse -> TColumn $ fromVector $
VG.generate (VG.length os - 1) $ \g ->
let !start = os `VG.unsafeIndex` g
!end = os `VG.unsafeIndex` (g+1)
in go (col `VG.unsafeIndex` (indices `VG.unsafeIndex` start)) (start + 1) end
where
{-# INLINE go #-}
go !acc j e
| j == e = acc
| otherwise =
let !x = col `VG.unsafeIndex` (indices `VG.unsafeIndex` j)
in go (f acc x) (j + 1) e
STrue -> TColumn $ fromVector $
VG.generate (VG.length os - 1) $ \g ->
let !start = os `VG.unsafeIndex` g
!end = os `VG.unsafeIndex` (g+1)
in go (col `VG.unsafeIndex` (indices `VG.unsafeIndex` start)) (start + 1) end
where
{-# INLINE go #-}
go !acc j e
| j == e = acc
| otherwise =
let !x = col `VG.unsafeIndex` (indices `VG.unsafeIndex` j)
in go (f acc x) (j + 1) e
Just (OptionalColumn col) -> TColumn $ fromVector $
VG.generate (VG.length os - 1) $ \g ->
let !start = os `VG.unsafeIndex` g
!end = os `VG.unsafeIndex` (g+1)
in go (col `VG.unsafeIndex` (indices `VG.unsafeIndex` start)) (start + 1) end
where
{-# INLINE go #-}
go !acc j e
| j == e = acc
| otherwise =
let !x = col `VG.unsafeIndex` (indices `VG.unsafeIndex` j)
in go (f acc x) (j + 1) e
interpretAggregation gdf@(Grouped df names indices os) (NumericAggregate name op (f :: VU.Vector b -> c)) = case getColumn name df of
Nothing -> throw $ ColumnNotFoundException name "" (map fst $ M.toList $ columnIndices df)
Just (UnboxedColumn (col :: VU.Vector d)) -> case testEquality (typeRep @b) (typeRep @d) of
Nothing -> case testEquality (typeRep @d) (typeRep @Int) of
Just Refl -> case sUnbox @c of
SFalse -> TColumn $ fromVector $
V.generate (VG.length os - 1)
(\i -> f (VU.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i))
(\j -> fromIntegral (col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))))
)
)
STrue -> TColumn $ fromUnboxedVector $
VU.generate (VG.length os - 1)
(\i -> f (VU.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i))
(\j -> fromIntegral (col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))))
)
)
Just Refl -> case sNumeric @d of
SFalse -> error $ "Cannot apply numeric aggregation to non-numeric column: " ++ (T.unpack name)
STrue -> case sUnbox @c of
SFalse -> TColumn $ fromVector $
V.generate (VG.length os - 1)
(\i -> f (VU.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i))
(\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i))))
)
)
STrue -> TColumn $ fromUnboxedVector $
VU.generate (VG.length os - 1)
(\i -> f (VU.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i))
(\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i))))
)
)
_ -> error $ "Cannot apply numeric aggregation to non-numeric column: " ++ (T.unpack name)
instance (Num a, Columnable a) => Num (Expr a) where
(+) :: Expr a -> Expr a -> Expr a
(+) = BinOp "add" (+)
(*) :: Expr a -> Expr a -> Expr a
(*) = BinOp "mult" (*)
fromInteger :: Integer -> Expr a
fromInteger = Lit . fromInteger
negate :: Expr a -> Expr a
negate = Apply "negate" negate
abs :: Num a => Expr a -> Expr a
abs = Apply "abs" abs
signum :: Num a => Expr a -> Expr a
signum = Apply "signum" signum
instance (Fractional a, Columnable a) => Fractional (Expr a) where
fromRational :: (Fractional a, Columnable a) => Rational -> Expr a
fromRational = Lit . fromRational
(/) :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a
(/) = BinOp "divide" (/)
instance (Floating a, Columnable a) => Floating (Expr a) where
pi :: (Floating a, Columnable a) => Expr a
pi = Lit pi
exp :: (Floating a, Columnable a) => Expr a -> Expr a
exp = Apply "exp" exp
log :: (Floating a, Columnable a) => Expr a -> Expr a
log = Apply "log" log
sin :: (Floating a, Columnable a) => Expr a -> Expr a
sin = Apply "sin" sin
cos :: (Floating a, Columnable a) => Expr a -> Expr a
cos = Apply "cos" cos
asin :: (Floating a, Columnable a) => Expr a -> Expr a
asin = Apply "asin" asin
acos :: (Floating a, Columnable a) => Expr a -> Expr a
acos = Apply "acos" acos
atan :: (Floating a, Columnable a) => Expr a -> Expr a
atan = Apply "atan" atan
sinh :: (Floating a, Columnable a) => Expr a -> Expr a
sinh = Apply "sinh" sinh
cosh :: (Floating a, Columnable a) => Expr a -> Expr a
cosh = Apply "cosh" cosh
asinh :: (Floating a, Columnable a) => Expr a -> Expr a
asinh = Apply "asinh" sinh
acosh :: (Floating a, Columnable a) => Expr a -> Expr a
acosh = Apply "acosh" acosh
atanh :: (Floating a, Columnable a) => Expr a -> Expr a
atanh = Apply "atanh" atanh
instance (Show a) => Show (Expr a) where
show :: forall a . Show a => Expr a -> String
show (Col name) = "col@" ++ show (typeRep @a) ++ "(" ++ T.unpack name ++ ")"
show (Lit value) = show value
show (Apply name f value) = T.unpack name ++ "(" ++ show value ++ ")"
show (BinOp name f a b) = T.unpack name ++ "(" ++ show a ++ ", " ++ show b ++ ")"