dataframe-0.3.3.0: src/DataFrame/Internal/Expression.hs
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UndecidableInstances #-}
module DataFrame.Internal.Expression where
import qualified Data.Map as M
import Data.Maybe (fromMaybe, isJust)
import qualified Data.Text as T
import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))
import qualified Data.Vector as V
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Unboxed as VU
import DataFrame.Errors
import DataFrame.Internal.Column
import DataFrame.Internal.DataFrame
import DataFrame.Internal.Types
import Type.Reflection (TypeRep, Typeable, typeOf, typeRep, pattern App)
data Expr a where
Col :: (Columnable a) => T.Text -> Expr a
Lit :: (Columnable a) => a -> Expr a
If ::
(Columnable a) =>
Expr Bool ->
Expr a ->
Expr a ->
Expr a
UnaryOp ::
( Columnable a
, Columnable b
) =>
T.Text -> -- Operation name
(b -> a) ->
Expr b ->
Expr a
BinaryOp ::
( Columnable c
, Columnable b
, Columnable a
) =>
T.Text -> -- operation name
(c -> b -> a) ->
Expr c ->
Expr b ->
Expr a
AggVector ::
( VG.Vector v b
, Typeable v
, Columnable a
, Columnable b
) =>
Expr b ->
T.Text -> -- Operation name
(v b -> a) ->
Expr a
AggReduce ::
(Columnable a) =>
Expr a ->
T.Text -> -- Operation name
(forall a. (Columnable a) => a -> a -> a) ->
Expr a
AggNumericVector ::
( Columnable a
, Columnable b
, VU.Unbox a
, VU.Unbox b
, Num a
, Num b
) =>
Expr b ->
T.Text -> -- Operation name
(VU.Vector b -> a) ->
Expr a
AggFold ::
forall a b.
(Columnable a, Columnable b) =>
Expr b ->
T.Text -> -- Operation name
a ->
(a -> b -> a) ->
Expr a
data UExpr where
Wrap :: (Columnable a) => Expr a -> UExpr
interpret ::
forall a.
(Columnable a) =>
DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)
interpret df (Lit value) =
pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value
interpret df (Col name) = case getColumn name df of
Nothing -> Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df)
Just col -> pure $ TColumn col
interpret df expr@(If cond l r) = case interpret @Bool df cond of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show cond)
}
)
Left e -> Left e
Right (TColumn conditions) -> case interpret @a df l of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
(context{callingFunctionName = Just "interpret", errorColumnName = Just (show l)})
Left e -> Left e
Right (TColumn left) -> case interpret @a df r of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
(context{callingFunctionName = Just "interpret", errorColumnName = Just (show r)})
Left e -> Left e
Right (TColumn right) -> case zipWithColumns
(\(c :: Bool) (l' :: a, r' :: a) -> if c then l' else r')
conditions
(zipColumns left right) of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show expr)
}
)
Left e -> Left e
Right res -> pure $ TColumn res
interpret df expr@(UnaryOp _ (f :: c -> d) value) = case interpret @c df value of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show value)
}
)
Left e -> Left e
Right (TColumn value') -> case mapColumn f value' of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show expr)
}
)
Left e -> Left e
Right res -> pure $ TColumn res
interpret df expr@(BinaryOp _ (f :: c -> d -> e) (Lit left) (Lit right)) =
pure $
TColumn $
fromVector $
V.replicate (fst $ dataframeDimensions df) (f left right)
interpret df expr@(BinaryOp _ (f :: c -> d -> e) (Lit left) right) = case interpret @d df right of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show right)
}
)
Left e -> Left e
Right (TColumn right') -> case mapColumn (f left) right' of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show expr)
}
)
Left e -> Left e
Right res -> pure $ TColumn res
interpret df expr@(BinaryOp _ (f :: c -> d -> e) left (Lit right)) = case interpret @c df left of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show left)
}
)
Left e -> Left e
Right (TColumn left') -> case mapColumn (`f` right) left' of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show expr)
}
)
Left e -> Left e
Right res -> pure $ TColumn res
interpret df expr@(BinaryOp _ (f :: c -> d -> e) left right) = case interpret @c df left of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show left)
}
)
Left e -> Left e
Right (TColumn left') -> case interpret @d df right of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show right)
}
)
Left e -> Left e
Right (TColumn right') -> case zipWithColumns f left' right' of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show expr)
}
)
Left e -> Left e
Right res -> pure $ TColumn res
interpret df expression@(AggVector expr op (f :: v b -> c)) = case interpret @b df expr of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show expr)
}
)
Left e -> Left e
Right (TColumn column) -> case column of
(BoxedColumn col) -> case testEquality (typeRep @(v b)) (typeOf col) of
Just Refl -> interpret @c df (Lit (f col))
Nothing ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @(v b))
, expectedType = Right (typeOf col)
, callingFunctionName = Just "interpret"
, errorColumnName = Nothing
}
)
(OptionalColumn col) -> case testEquality (typeRep @(v b)) (typeOf col) of
Just Refl -> interpret @c df (Lit (f col))
Nothing ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @(v b))
, expectedType = Right (typeOf col)
, callingFunctionName = Just "interpret"
, errorColumnName = Nothing
}
)
(UnboxedColumn col) -> case testEquality (typeRep @(v b)) (typeOf col) of
Just Refl -> interpret @c df (Lit (f col))
Nothing ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @(v b))
, expectedType = Right (typeOf col)
, callingFunctionName = Just "interpret"
, errorColumnName = Nothing
}
)
interpret df expression@(AggReduce expr op (f :: forall a. (Columnable a) => a -> a -> a)) = case interpret @a df expr of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show expr)
}
)
Left e -> Left e
Right (TColumn column) -> case headColumn @a column of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show expr)
}
)
Left (EmptyDataSetException loc) -> Left (EmptyDataSetException (T.pack $ show expr))
Left e -> Left e
Right h -> case ifoldlColumn (\acc _ v -> f acc v) h column of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show expression)
}
)
Left e -> Left e
Right value ->
pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value
interpret df expression@(AggNumericVector expr op (f :: VU.Vector b -> c)) = case interpret @b df expr of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show expr)
}
)
Left e -> Left e
Right (TColumn column) -> case column of
(UnboxedColumn (v :: VU.Vector d)) -> case testEquality (typeRep @d) (typeRep @b) of
Just Refl ->
pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) (f v)
Nothing ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @b)
, expectedType = Right (typeRep @d)
, callingFunctionName = Just "interpret"
, errorColumnName = Just (show expression)
}
)
(BoxedColumn (v :: V.Vector d)) -> case testEquality (typeRep @d) (typeRep @Integer) of
Just Refl ->
Right $
TColumn $
fromVector $
V.replicate
(fst $ dataframeDimensions df)
(f (VU.convert $ V.map fromInteger v))
Nothing ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @Integer)
, expectedType = Right (typeRep @d)
, callingFunctionName = Just "interpret"
, errorColumnName = Just (show expression)
}
)
(OptionalColumn (v :: V.Vector (Maybe d))) -> case sNumeric @d of
STrue -> case testEquality (typeRep @d) (typeRep @b) of
Nothing ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @b)
, expectedType = Right (typeRep @d)
, callingFunctionName = Just "interpret"
, errorColumnName = Just (show expression)
}
)
Just Refl ->
pure $
TColumn $
fromVector $
V.replicate
(fst $ dataframeDimensions df)
(f (VU.convert $ V.map (fromMaybe 0) $ V.filter isJust v))
SFalse ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @d)
, expectedType = Right (typeRep @b)
, callingFunctionName = Just "interpret"
, errorColumnName = Just (show expression)
}
)
interpret df expression@(AggFold expr op start (f :: (a -> b -> a))) = case interpret @b df expr of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show expr)
}
)
Left e -> Left e
Right (TColumn column) -> case ifoldlColumn (\acc _ v -> f acc v) start column of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpret"
, errorColumnName = Just (show expression)
}
)
Left e -> Left e
Right value ->
pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value
data AggregationResult a
= UnAggregated Column
| Aggregated (TypedColumn a)
mkUnaggregatedColumn ::
forall v a.
(VG.Vector v a, Columnable a) =>
v a -> VU.Vector Int -> VU.Vector Int -> V.Vector (v a)
mkUnaggregatedColumn col os indices =
V.generate
(VU.length os - 1)
( \i ->
( VG.generate
(os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i))
( \j ->
col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))
)
)
)
nestedTypeException ::
forall a b. (Typeable a, Typeable b) => String -> DataFrameException
nestedTypeException expression = case typeRep @a of
App t1 t2 ->
TypeMismatchException
( MkTypeErrorContext
{ userType = Left (show (typeRep @b)) :: Either String (TypeRep ())
, expectedType = Left (show (typeRep @a)) :: Either String (TypeRep ())
, callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just expression
}
)
t ->
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @(VU.Vector b))
, expectedType = Right (typeRep @b)
, callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just expression
}
)
interpretAggregation ::
forall a.
(Columnable a) =>
GroupedDataFrame -> Expr a -> Either DataFrameException (AggregationResult a)
interpretAggregation gdf (Lit value) =
Right $
Aggregated $
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 -> Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df)
Just (BoxedColumn col) -> Right $ UnAggregated $ fromVector $ mkUnaggregatedColumn col os indices
Just (OptionalColumn col) -> Right $ UnAggregated $ fromVector $ mkUnaggregatedColumn col os indices
Just (UnboxedColumn col) -> Right $ UnAggregated $ fromVector $ mkUnaggregatedColumn col os indices
interpretAggregation gdf expression@(UnaryOp _ (f :: c -> d) expr) =
case interpretAggregation @c gdf expr of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expr)
}
)
Left e -> Left e
Right (UnAggregated unaggregated) -> case unaggregated of
BoxedColumn (col :: V.Vector b) -> case testEquality (typeRep @b) (typeRep @(V.Vector c)) of
Just Refl -> Right $ UnAggregated $ fromVector $ V.map (V.map f) col
Nothing -> case testEquality (typeRep @b) (typeRep @(VU.Vector c)) of
Nothing -> Left $ nestedTypeException @b @c (show expression)
Just Refl -> case (sUnbox @c, sUnbox @a) of
(SFalse, _) -> Left $ InternalException "Boxed type inside an unboxed column"
(STrue, STrue) -> Right $ UnAggregated $ fromVector $ V.map (VU.map f) col
(STrue, _) -> Right $ UnAggregated $ fromVector $ V.map (V.map f . VU.convert) col
_ -> Left $ InternalException "Aggregated into a non-boxed column"
Right (Aggregated (TColumn aggregated)) -> case mapColumn f aggregated of
Left e -> Left e
Right col -> Right $ Aggregated $ TColumn col
interpretAggregation gdf expression@(BinaryOp _ (f :: c -> d -> e) left right) =
case (interpretAggregation @c gdf left, interpretAggregation @d gdf right) of
(Right (Aggregated (TColumn left')), Right (Aggregated (TColumn right'))) -> case zipWithColumns f left' right' of
Left e -> Left e
Right col -> Right $ Aggregated $ TColumn col
(Right (UnAggregated left'), Right (UnAggregated right')) -> case (left', right') of
(BoxedColumn (l :: V.Vector m), BoxedColumn (r :: V.Vector n)) -> case testEquality (typeRep @m) (typeRep @(VU.Vector c)) of
Just Refl -> case testEquality (typeRep @n) (typeRep @(VU.Vector d)) of
Just Refl -> case (sUnbox @c, sUnbox @d, sUnbox @e) of
(STrue, STrue, STrue) ->
Right $ UnAggregated $ fromVector $ V.zipWith (\l' r' -> VU.zipWith f l' r') l r
(STrue, STrue, SFalse) ->
Right $
UnAggregated $
fromVector $
V.zipWith (\l' r' -> V.zipWith f (V.convert l') (V.convert r')) l r
(_, _, _) -> Left $ InternalException "Boxed vectors contain unboxed types"
Nothing -> case testEquality (typeRep @n) (typeRep @(V.Vector d)) of
Just Refl -> case sUnbox @c of
STrue ->
Right $
UnAggregated $
fromVector $
V.zipWith (\l' r' -> V.zipWith f (V.convert l') r') l r
SFalse -> Left $ InternalException "Unboxed vectors contain boxed types"
Nothing -> Left $ nestedTypeException @n @d (show right)
Nothing -> case testEquality (typeRep @m) (typeRep @(V.Vector c)) of
Nothing -> Left $ nestedTypeException @m @c (show left)
Just Refl -> case testEquality (typeRep @n) (typeRep @(VU.Vector d)) of
Just Refl -> case (sUnbox @d, sUnbox @e) of
(STrue, STrue) ->
Right $
UnAggregated $
fromVector $
V.zipWith
(\l' r' -> V.convert @V.Vector @e @VU.Vector $ V.zipWith f l' (V.convert r'))
l
r
(STrue, SFalse) ->
Right $
UnAggregated $
fromVector $
V.zipWith (\l' r' -> V.zipWith f l' (V.convert r')) l r
(_, _) -> Left $ InternalException "Unboxed vectors contain boxed types"
Nothing -> case testEquality (typeRep @n) (typeRep @(V.Vector d)) of
Just Refl -> case sUnbox @c of
STrue -> case sUnbox @e of
SFalse ->
Right $
UnAggregated $
fromVector $
V.zipWith (\l' r' -> V.zipWith f (V.convert l') r') l r
STrue ->
Right $
UnAggregated $
fromVector $
V.zipWith (\l' r' -> V.convert @V.Vector @e @VU.Vector $ V.zipWith f l' r') l r
SFalse -> Left $ InternalException "Unboxed vectors contain boxed types"
Nothing -> Left $ nestedTypeException @n @d (show right)
_ -> Left $ InternalException "Aggregated into a non-boxed column"
(Right _, Right _) ->
Left $
AggregatedAndNonAggregatedException (T.pack $ show left) (T.pack $ show right)
(Left (TypeMismatchException context), _) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show left)
}
)
(Left e, _) -> Left e
(_, Left (TypeMismatchException context)) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show right)
}
)
(_, Left e) -> Left e
interpretAggregation gdf expression@(If cond l r) =
case ( interpretAggregation @Bool gdf cond
, interpretAggregation @a gdf l
, interpretAggregation @a gdf r
) of
( Right (Aggregated (TColumn conditions))
, Right (Aggregated (TColumn left))
, Right (Aggregated (TColumn right))
) -> case zipWithColumns
(\(c :: Bool) (l' :: a, r' :: a) -> if c then l' else r')
conditions
(zipColumns left right) of
Left e -> Left e
Right v -> Right $ Aggregated (TColumn v)
( Right (UnAggregated conditions)
, Right (UnAggregated left@(BoxedColumn (left' :: V.Vector b)))
, Right (UnAggregated right@(BoxedColumn (right' :: V.Vector c)))
) -> case testEquality (typeRep @b) (typeRep @c) of
Nothing ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @b)
, expectedType = Right (typeRep @c)
, callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expression)
}
)
Just Refl -> case testEquality (typeRep @(V.Vector a)) (typeRep @b) of
Just Refl -> case zipWithColumns
( \(c :: VU.Vector Bool) (l' :: V.Vector a, r' :: V.Vector a) ->
V.zipWith
(\c' (l'', r'') -> if c' then l'' else r'')
(V.convert c)
(V.zip l' r')
)
conditions
(zipColumns left right) of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expression)
}
)
Left e -> Left e
Right v -> Right $ UnAggregated v
Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @b) of
Nothing -> Left $ nestedTypeException @b @a (show expression)
Just Refl -> case sUnbox @a of
SFalse -> Left $ InternalException "Boxed type in unboxed column"
STrue -> case zipWithColumns
( \(c :: VU.Vector Bool) (l' :: VU.Vector a, r' :: VU.Vector a) ->
VU.zipWith
(\c' (l'', r'') -> if c' then l'' else r'')
c
(VU.zip l' r')
)
conditions
(zipColumns left right) of
Left (TypeMismatchException context) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expression)
}
)
Left e -> Left e
Right v -> Right $ UnAggregated v
(Right _, Right _, Right _) ->
Left $
AggregatedAndNonAggregatedException (T.pack $ show l) (T.pack $ show r)
(Left (TypeMismatchException context), _, _) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show cond)
}
)
(Left e, _, _) -> Left e
(_, Left (TypeMismatchException context), _) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show l)
}
)
(_, Left e, _) -> Left e
(_, _, Left (TypeMismatchException context)) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show r)
}
)
(_, _, Left e) -> Left e
interpretAggregation gdf@(Grouped df names indices os) expression@(AggVector expr op (f :: v b -> c)) =
case interpretAggregation @b gdf expr of
Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(v b)) (typeRep @d) of
Nothing -> Left $ nestedTypeException @d @b (show expr)
Just Refl -> case testEquality (typeRep @v) (typeRep @V.Vector) of
Nothing -> Right $ Aggregated $ TColumn $ fromVector $ V.map (f . V.convert) col
Just Refl -> Right $ Aggregated $ TColumn $ fromVector $ V.map f col
Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column"
Right (Aggregated (TColumn (BoxedColumn (col :: V.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of
Just Refl -> case testEquality (typeRep @v) (typeRep @V.Vector) of
Just Refl -> interpretAggregation @c gdf (Lit (f col))
Nothing -> interpretAggregation @c gdf (Lit ((f . V.convert) col))
Nothing ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @b)
, expectedType = Right (typeRep @d)
, callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expr)
}
)
Right (Aggregated (TColumn (UnboxedColumn (col :: VU.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of
Just Refl -> case testEquality (typeRep @v) (typeRep @VU.Vector) of
Just Refl -> interpretAggregation @c gdf (Lit (f col))
Nothing -> interpretAggregation @c gdf (Lit ((f . VU.convert) col))
Nothing ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @b)
, expectedType = Right (typeRep @d)
, callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expr)
}
)
Right (Aggregated (TColumn (OptionalColumn (col :: V.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of
Just Refl -> case testEquality (typeRep @v) (typeRep @V.Vector) of
Just Refl -> interpretAggregation @c gdf (Lit (f col))
Nothing -> interpretAggregation @c gdf (Lit ((f . V.convert) col))
Nothing ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @b)
, expectedType = Right (typeRep @d)
, callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expr)
}
)
(Left (TypeMismatchException context)) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expression)
}
)
(Left e) -> Left e
interpretAggregation gdf@(Grouped df names indices os) expression@(AggNumericVector expr op (f :: VU.Vector b -> c)) =
case interpretAggregation @b gdf expr of
(Left (TypeMismatchException context)) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expression)
}
)
(Left e) -> Left e
Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(VU.Vector b)) (typeRep @d) of
Nothing -> case testEquality (typeRep @(VU.Vector Int)) (typeRep @d) of
Nothing -> case testEquality (typeRep @(V.Vector Integer)) (typeRep @d) of
Nothing -> Left $ nestedTypeException @d @b (show expr)
Just Refl ->
Right $
Aggregated $
TColumn $
fromVector $
V.map (f . VU.convert . V.map fromIntegral) col
Just Refl ->
Right $
Aggregated $
TColumn $
fromVector $
V.map f (VG.map (VG.map fromIntegral) col)
Just Refl -> Right $ Aggregated $ TColumn $ fromVector $ V.map f col
Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column"
Right (Aggregated (TColumn (BoxedColumn (col :: V.Vector d)))) -> case testEquality (typeRep @Integer) (typeRep @d) of
Just Refl -> interpretAggregation @c gdf (Lit ((f . V.convert . V.map fromIntegral) col))
Nothing ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @b)
, expectedType = Right (typeRep @d)
, callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expr)
}
)
Right (Aggregated (TColumn (UnboxedColumn (col :: VU.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of
Just Refl -> interpretAggregation @c gdf (Lit (f col))
Nothing ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @b)
, expectedType = Right (typeRep @d)
, callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expr)
}
)
Right (Aggregated (TColumn (OptionalColumn (col :: V.Vector (Maybe d))))) -> case testEquality (typeRep @b) (typeRep @d) of
Just Refl ->
interpretAggregation @c
gdf
(Lit ((f . V.convert . V.map (fromMaybe 0) . V.filter isJust) col))
Nothing ->
Left $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right (typeRep @b)
, expectedType = Right (typeRep @d)
, callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expr)
}
)
interpretAggregation gdf@(Grouped df names indices os) expression@(AggReduce expr op (f :: forall a. (Columnable a) => a -> a -> a)) =
case interpretAggregation @a gdf expr of
(Left (TypeMismatchException context)) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expression)
}
)
(Left e) -> Left e
Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(V.Vector a)) (typeRep @d) of
Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @d) of
Nothing -> Left $ nestedTypeException @d @a (show expr)
Just Refl -> case sUnbox @a of
STrue ->
Right $
Aggregated $
TColumn $
fromVector $
V.map (\v -> VU.foldl' f (VG.head v) (VG.drop 1 v)) col
SFalse -> Left $ InternalException "Boxed type inside an unboxed column"
Just Refl ->
Right $
Aggregated $
TColumn $
fromVector $
V.map (\v -> VG.foldl' f (VG.head v) (VG.drop 1 v)) col
Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column"
Right (Aggregated (TColumn column)) -> case headColumn @a column of
Left e -> Left e
Right h -> case ifoldlColumn (\acc _ v -> f acc v) h column of
Left e -> Left e
Right value -> interpretAggregation @a gdf (Lit value)
interpretAggregation gdf@(Grouped df names indices os) expression@(AggFold expr op s (f :: (a -> b -> a))) =
case interpretAggregation @b gdf expr of
(Left (TypeMismatchException context)) ->
Left $
TypeMismatchException
( context
{ callingFunctionName = Just "interpretAggregation"
, errorColumnName = Just (show expression)
}
)
(Left e) -> Left e
Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(V.Vector b)) (typeRep @d) of
Just Refl -> Right $ Aggregated $ TColumn $ fromVector $ V.map (\v -> V.foldl' f s v) col
Nothing -> case testEquality (typeRep @(VU.Vector b)) (typeRep @d) of
Just Refl -> case sUnbox @b of
STrue ->
Right $ Aggregated $ TColumn $ fromVector $ V.map (\v -> VU.foldl' f s v) col
SFalse -> Left $ InternalException "Boxed type inside an unboxed column"
Nothing -> Left $ nestedTypeException @d @b (show expr)
Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column"
Right (Aggregated (TColumn column)) -> case ifoldlColumn (\acc _ v -> f acc v) s column of
Left e -> Left e
Right value -> interpretAggregation @a gdf (Lit value)
instance (Num a, Columnable a) => Num (Expr a) where
(+) :: Expr a -> Expr a -> Expr a
(+) = BinaryOp "add" (+)
(*) :: Expr a -> Expr a -> Expr a
(*) = BinaryOp "mult" (*)
fromInteger :: Integer -> Expr a
fromInteger = Lit . fromInteger
negate :: Expr a -> Expr a
negate = UnaryOp "negate" negate
abs :: (Num a) => Expr a -> Expr a
abs = UnaryOp "abs" abs
signum :: (Num a) => Expr a -> Expr a
signum = UnaryOp "signum" signum
add :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a
add = (+)
sub :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a
sub = (-)
mult :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a
mult = (*)
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
(/) = BinaryOp "divide" (/)
divide :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a
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 = UnaryOp "exp" exp
log :: (Floating a, Columnable a) => Expr a -> Expr a
log = UnaryOp "log" log
sin :: (Floating a, Columnable a) => Expr a -> Expr a
sin = UnaryOp "sin" sin
cos :: (Floating a, Columnable a) => Expr a -> Expr a
cos = UnaryOp "cos" cos
asin :: (Floating a, Columnable a) => Expr a -> Expr a
asin = UnaryOp "asin" asin
acos :: (Floating a, Columnable a) => Expr a -> Expr a
acos = UnaryOp "acos" acos
atan :: (Floating a, Columnable a) => Expr a -> Expr a
atan = UnaryOp "atan" atan
sinh :: (Floating a, Columnable a) => Expr a -> Expr a
sinh = UnaryOp "sinh" sinh
cosh :: (Floating a, Columnable a) => Expr a -> Expr a
cosh = UnaryOp "cosh" cosh
asinh :: (Floating a, Columnable a) => Expr a -> Expr a
asinh = UnaryOp "asinh" sinh
acosh :: (Floating a, Columnable a) => Expr a -> Expr a
acosh = UnaryOp "acosh" acosh
atanh :: (Floating a, Columnable a) => Expr a -> Expr a
atanh = UnaryOp "atanh" atanh
instance (Show a) => Show (Expr a) where
show :: forall a. (Show a) => Expr a -> String
show (Col name) = "(col @" ++ show (typeRep @a) ++ " " ++ show name ++ ")"
show (Lit value) = "(lit (" ++ show value ++ "))"
show (If cond l r) = "(ifThenElse " ++ show cond ++ " " ++ show l ++ " " ++ show r ++ ")"
show (UnaryOp name f value) = "(" ++ T.unpack name ++ " " ++ show value ++ ")"
show (BinaryOp name f a b) = "(" ++ T.unpack name ++ " " ++ show a ++ " " ++ show b ++ ")"
show (AggNumericVector expr op _) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"
show (AggVector expr op _) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"
show (AggReduce expr op _) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"
show (AggFold expr op _ _) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"
instance (Eq a, Show a) => Eq (Expr a) where
(==) :: (Eq a, Show a) => Expr a -> Expr a -> Bool
(==) l r = show l == show r
replaceExpr ::
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
Expr a -> Expr b -> Expr c -> Expr c
replaceExpr new old expr = case testEquality (typeRep @b) (typeRep @c) of
Just Refl -> case testEquality (typeRep @a) (typeRep @c) of
Just Refl -> if old == expr then new else replace'
Nothing -> expr
Nothing -> replace'
where
replace' = case expr of
(Col _) -> expr
(Lit _) -> expr
(If cond l r) ->
If (replaceExpr new old cond) (replaceExpr new old l) (replaceExpr new old r)
(UnaryOp name f value) -> UnaryOp name f (replaceExpr new old value)
(BinaryOp name f l r) -> BinaryOp name f (replaceExpr new old l) (replaceExpr new old r)
(AggNumericVector expr op f) -> AggNumericVector (replaceExpr new old expr) op f
(AggVector expr op f) -> AggVector (replaceExpr new old expr) op f
(AggReduce expr op f) -> AggReduce (replaceExpr new old expr) op f
(AggFold expr op acc f) -> AggFold (replaceExpr new old expr) op acc f
eSize :: Expr a -> Int
eSize (Col _) = 1
eSize (Lit _) = 1
eSize (If c l r) = 1 + eSize c + eSize l + eSize r
eSize (UnaryOp _ _ e) = 1 + eSize e
eSize (BinaryOp _ _ l r) = 1 + eSize l + eSize r
eSize (AggNumericVector expr op _) = eSize expr + 1
eSize (AggVector expr op _) = eSize expr + 1
eSize (AggReduce expr op _) = eSize expr + 1
eSize (AggFold expr op _ _) = eSize expr + 1