dataframe-0.3.3.5: src/DataFrame/Internal/Expression.hs
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# 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 Data.String
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 :: VU.Vector d)) -> case testEquality (typeRep @(v b)) (typeOf col) of
Just Refl -> interpret @c df (Lit (f col))
Nothing -> case testEquality (typeRep @b) (typeRep @d) of
Just Refl -> interpret @c df (Lit (f (V.convert 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 (VU.zipWith f) 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 (V.zipWith f . V.convert) 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 (V.zipWith f . V.convert) 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.foldl' f s) col
Nothing -> case testEquality (typeRep @(VU.Vector b)) (typeRep @d) of
Just Refl -> case sUnbox @b of
STrue ->
Right $ Aggregated $ TColumn $ fromVector $ V.map (VU.foldl' f s) 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
(+) (Lit x) (Lit y) = Lit (x + y)
(+) (Lit x) expr = UnaryOp ("add " <> (T.pack . show) (Lit x)) (+ x) expr
(+) expr (Lit x) = UnaryOp ("add " <> (T.pack . show) (Lit x)) (+ x) expr
(+) e1 e2
| e1 == e2 = UnaryOp ("mult " <> (T.pack . show) (Lit @a 2)) (* 2) e1
| otherwise = BinaryOp "add" (+) e1 e2
(-) :: Expr a -> Expr a -> Expr a
(-) (Lit x) (Lit y) = Lit (x - y)
(-) e1 e2 = BinaryOp "sub" (-) e1 e2
(*) :: Expr a -> Expr a -> Expr a
(*) (Lit 0) _ = Lit 0
(*) _ (Lit 0) = Lit 0
(*) (Lit 1) e = e
(*) e (Lit 1) = e
(*) (Lit x) (Lit y) = Lit (x * y)
(*) (Lit x) expr = UnaryOp ("mult " <> (T.pack . show) (Lit x)) (* x) expr
(*) expr (Lit x) = UnaryOp ("mult " <> (T.pack . show) (Lit x)) (* x) expr
(*) e1 e2
| e1 == e2 = UnaryOp "pow 2" (^ 2) e1
| otherwise = BinaryOp "mult" (*) e1 e2
fromInteger :: Integer -> Expr a
fromInteger = Lit . fromInteger
negate :: Expr a -> Expr a
negate (Lit n) = Lit (negate n)
negate expr = case expr of
(UnaryOp "negate" negate e) -> case testEquality (typeOf e) (typeOf expr) of
Nothing -> expr -- This case is impossible.
Just Refl -> e
_ -> UnaryOp "negate" negate expr
abs :: (Num a) => Expr a -> Expr a
abs expr = case expr of
(UnaryOp "abs" abs e) -> case testEquality (typeOf e) (typeOf expr) of
Nothing -> expr -- This case is impossible.
Just Refl -> e
_ -> UnaryOp "abs" abs expr
signum :: (Num a) => Expr a -> Expr a
signum (Lit n) = Lit (signum n)
signum expr = UnaryOp "signum" signum expr
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 (IsString a, Columnable a) => IsString (Expr a) where
fromString :: String -> Expr a
fromString s = Lit (fromString s)
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 ++ ")"
normalize :: (Eq a, Ord a, Show a, Typeable a) => Expr a -> Expr a
normalize expr = case expr of
Col name -> Col name
Lit val -> Lit val
If cond th el -> If (normalize cond) (normalize th) (normalize el)
UnaryOp name f e -> UnaryOp name f (normalize e)
BinaryOp name f e1 e2
| isCommutative name ->
let n1 = normalize e1
n2 = normalize e2
in case testEquality (typeOf n1) (typeOf n2) of
Nothing -> expr
Just Refl ->
if compareExpr n1 n2 == GT
then BinaryOp name f n2 n1 -- Swap to canonical order
else BinaryOp name f n1 n2
| otherwise -> BinaryOp name f (normalize e1) (normalize e2)
AggVector e name f -> AggVector (normalize e) name f
AggReduce e name f -> AggReduce (normalize e) name f
AggNumericVector e name f -> AggNumericVector (normalize e) name f
AggFold e name init f -> AggFold (normalize e) name init f
isCommutative :: T.Text -> Bool
isCommutative name =
name
`elem` [ "add"
, "mult"
, "min"
, "max"
, "eq"
, "and"
, "or"
]
-- Compare expressions for ordering (used in normalization)
compareExpr :: Expr a -> Expr a -> Ordering
compareExpr e1 e2 = compare (exprKey e1) (exprKey e2)
where
exprKey :: Expr a -> String
exprKey (Col name) = "0:" ++ T.unpack name
exprKey (Lit val) = "1:" ++ show val
exprKey (If c t e) = "2:" ++ exprKey c ++ exprKey t ++ exprKey e
exprKey (UnaryOp name _ e) = "3:" ++ T.unpack name ++ exprKey e
exprKey (BinaryOp name _ e1 e2) = "4:" ++ T.unpack name ++ exprKey e1 ++ exprKey e2
exprKey (AggVector e name _) = "5:" ++ T.unpack name ++ exprKey e
exprKey (AggReduce e name _) = "6:" ++ T.unpack name ++ exprKey e
exprKey (AggNumericVector e name _) = "7:" ++ T.unpack name ++ exprKey e
exprKey (AggFold e name _ _) = "8:" ++ T.unpack name ++ exprKey e
instance (Eq a, Columnable a) => Eq (Expr a) where
(==) l r = eqNormalized (normalize l) (normalize r)
where
exprEq :: (Columnable b, Columnable c) => Expr b -> Expr c -> Bool
exprEq e1 e2 = case testEquality (typeOf e1) (typeOf e2) of
Just Refl -> e1 == e2
Nothing -> False
eqNormalized :: Expr a -> Expr a -> Bool
eqNormalized (Col n1) (Col n2) = n1 == n2
eqNormalized (Lit v1) (Lit v2) = v1 == v2
eqNormalized (If c1 t1 e1) (If c2 t2 e2) =
c1 == c2 && t1 `exprEq` t2 && e1 `exprEq` e2
eqNormalized (UnaryOp n1 _ e1) (UnaryOp n2 _ e2) =
n1 == n2 && e1 `exprEq` e2
eqNormalized (BinaryOp n1 _ e1a e1b) (BinaryOp n2 _ e2a e2b) =
n1 == n2 && e1a `exprEq` e2a && e1b `exprEq` e2b
eqNormalized (AggVector e1 n1 _) (AggVector e2 n2 _) =
n1 == n2 && e1 `exprEq` e2
eqNormalized (AggReduce e1 n1 _) (AggReduce e2 n2 _) =
n1 == n2 && e1 `exprEq` e2
eqNormalized (AggNumericVector e1 n1 _) (AggNumericVector e2 n2 _) =
n1 == n2 && e1 `exprEq` e2
eqNormalized (AggFold e1 n1 i1 _) (AggFold e2 n2 i2 _) =
n1 == n2 && e1 `exprEq` e2 && i1 == i2
eqNormalized _ _ = False
instance (Ord a, Columnable a) => Ord (Expr a) where
compare :: Expr a -> Expr a -> Ordering
compare e1 e2 = case (e1, e2) of
(Col n1, Col n2) -> compare n1 n2
(Lit v1, Lit v2) -> compare v1 v2
(If c1 t1 e1', If c2 t2 e2') ->
compare c1 c2 <> exprComp t1 t2 <> exprComp e1' e2'
(UnaryOp n1 _ e1', UnaryOp n2 _ e2') ->
compare n1 n2 <> exprComp e1' e2'
(BinaryOp n1 _ a1 b1, BinaryOp n2 _ a2 b2) ->
compare n1 n2 <> exprComp a1 a2 <> exprComp b1 b2
(AggVector e1' n1 _, AggVector e2' n2 _) ->
compare n1 n2 <> exprComp e1' e2'
(AggReduce e1' n1 _, AggReduce e2' n2 _) ->
compare n1 n2 <> exprComp e1' e2'
(AggNumericVector e1' n1 _, AggNumericVector e2' n2 _) ->
compare n1 n2 <> exprComp e1' e2'
(AggFold e1' n1 i1 _, AggFold e2' n2 i2 _) ->
compare n1 n2 <> exprComp e1' e2' <> compare i1 i2
-- Different constructors - compare by priority
(Col _, _) -> LT
(_, Col _) -> GT
(Lit _, _) -> LT
(_, Lit _) -> GT
(UnaryOp{}, _) -> LT
(_, UnaryOp{}) -> GT
(BinaryOp{}, _) -> LT
(_, BinaryOp{}) -> GT
(If{}, _) -> LT
(_, If{}) -> GT
(AggVector{}, _) -> LT
(_, AggVector{}) -> GT
(AggReduce{}, _) -> LT
(_, AggReduce{}) -> GT
(AggNumericVector{}, _) -> LT
(_, AggNumericVector{}) -> GT
exprComp :: (Columnable b, Columnable c) => Expr b -> Expr c -> Ordering
exprComp e1 e2 = case testEquality (typeOf e1) (typeOf e2) of
Just Refl -> e1 `compare` e2
Nothing -> LT
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