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

module DataFrame.Internal.Interpreter (
    -- * New core API
    Value (..),
    Ctx (..),
    eval,
    materialize,

    -- * Backward-compatible API
    interpret,
    interpretAggregation,
    AggregationResult (..),
) where

import Data.Bifunctor (first)
import qualified Data.Map as M
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.Expression
import DataFrame.Internal.Types
import Type.Reflection (Typeable, typeRep)

-------------------------------------------------------------------------------
-- Value: the unified result type
-------------------------------------------------------------------------------

{- | The result of interpreting an expression.  Keeps literals as scalars
until the point where a concrete column is needed, avoiding premature
broadcast allocations.
-}
data Value a where
    -- | A single value, not yet broadcast to any length.
    Scalar :: (Columnable a) => a -> Value a
    {- | A flat column (one element per row in the flat case, or one
    element per group after aggregation).
    -}
    Flat :: (Columnable a) => Column -> Value a
    {- | A grouped column: one 'Column' slice per group.  Only produced
    when interpreting inside a 'GroupCtx'.
    -}
    Group :: (Columnable a) => V.Vector Column -> Value a

-- | The interpretation context.
data Ctx
    = FlatCtx DataFrame
    | GroupCtx GroupedDataFrame

-------------------------------------------------------------------------------
-- Materialisation
-------------------------------------------------------------------------------

{- | Force a 'Value' into a flat 'Column' of the given length.  Scalars
are broadcast; flat columns are returned as-is.
-}
materialize :: forall a. (Columnable a) => Int -> Value a -> Column
materialize :: forall a. Columnable a => Int -> Value a -> Column
materialize Int
n (Scalar a
v) = forall a. Columnable a => Int -> a -> Column
broadcastScalar @a Int
n a
v
materialize Int
_ (Flat Column
c) = Column
c
materialize Int
_ (Group Vector Column
_) =
    [Char] -> Column
forall a. HasCallStack => [Char] -> a
error [Char]
"materialize: cannot flatten a grouped value to a single column"

{- | Replicate a scalar to a column of length @n@, choosing the most
efficient representation.
-}
broadcastScalar :: forall a. (Columnable a) => Int -> a -> Column
broadcastScalar :: forall a. Columnable a => Int -> a -> Column
broadcastScalar Int
n a
v = case forall a. SBoolI (Unboxable a) => SBool (Unboxable a)
sUnbox @a of
    SBool (Unboxable a)
STrue -> Vector a -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (Int -> a -> Vector a
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
n a
v)
    SBool (Unboxable a)
SFalse -> Vector a -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector (Int -> a -> Vector a
forall a. Int -> a -> Vector a
V.replicate Int
n a
v)

-------------------------------------------------------------------------------
-- Lifting: the core combinators
-------------------------------------------------------------------------------

-- | Apply a pure function to a 'Value'.
liftValue ::
    (Columnable b, Columnable a) =>
    (b -> a) -> Value b -> Either DataFrameException (Value a)
liftValue :: forall b a.
(Columnable b, Columnable a) =>
(b -> a) -> Value b -> Either DataFrameException (Value a)
liftValue b -> a
f (Scalar b
v) = Value a -> Either DataFrameException (Value a)
forall a b. b -> Either a b
Right (a -> Value a
forall a. Columnable a => a -> Value a
Scalar (b -> a
f b
v))
liftValue b -> a
f (Flat Column
col) = Column -> Value a
forall a. Columnable a => Column -> Value a
Flat (Column -> Value a)
-> Either DataFrameException Column
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (b -> a) -> Column -> Either DataFrameException Column
forall b c.
(Columnable b, Columnable c) =>
(b -> c) -> Column -> Either DataFrameException Column
mapColumn b -> a
f Column
col
liftValue b -> a
f (Group Vector Column
gs) = Vector Column -> Value a
forall a. Columnable a => Vector Column -> Value a
Group (Vector Column -> Value a)
-> Either DataFrameException (Vector Column)
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Column -> Either DataFrameException Column)
-> Vector Column -> Either DataFrameException (Vector Column)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM ((b -> a) -> Column -> Either DataFrameException Column
forall b c.
(Columnable b, Columnable c) =>
(b -> c) -> Column -> Either DataFrameException Column
mapColumn b -> a
f) Vector Column
gs

{- | Apply a binary function to two 'Value's.  When one side is a
'Scalar' the operation degenerates to a 'liftValue' — this is how the
old @Binary op (Lit l) right@ special cases are recovered without
explicit pattern matches in the evaluator.
-}
liftValue2 ::
    (Columnable c, Columnable b, Columnable a) =>
    (c -> b -> a) ->
    Value c ->
    Value b ->
    Either DataFrameException (Value a)
liftValue2 :: forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Value c -> Value b -> Either DataFrameException (Value a)
liftValue2 c -> b -> a
f (Scalar c
l) (Scalar b
r) = Value a -> Either DataFrameException (Value a)
forall a b. b -> Either a b
Right (a -> Value a
forall a. Columnable a => a -> Value a
Scalar (c -> b -> a
f c
l b
r))
liftValue2 c -> b -> a
f (Scalar c
l) Value b
v = (b -> a) -> Value b -> Either DataFrameException (Value a)
forall b a.
(Columnable b, Columnable a) =>
(b -> a) -> Value b -> Either DataFrameException (Value a)
liftValue (c -> b -> a
f c
l) Value b
v
liftValue2 c -> b -> a
f Value c
v (Scalar b
r) = (c -> a) -> Value c -> Either DataFrameException (Value a)
forall b a.
(Columnable b, Columnable a) =>
(b -> a) -> Value b -> Either DataFrameException (Value a)
liftValue (c -> b -> a
`f` b
r) Value c
v
liftValue2 c -> b -> a
f (Flat Column
l) (Flat Column
r) = Column -> Value a
forall a. Columnable a => Column -> Value a
Flat (Column -> Value a)
-> Either DataFrameException Column
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (c -> b -> a)
-> Column -> Column -> Either DataFrameException Column
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
(a -> b -> c)
-> Column -> Column -> Either DataFrameException Column
zipWithColumns c -> b -> a
f Column
l Column
r
liftValue2 c -> b -> a
f (Group Vector Column
ls) (Group Vector Column
rs)
    | Vector Column -> Int
forall a. Vector a -> Int
V.length Vector Column
ls Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector Column -> Int
forall a. Vector a -> Int
V.length Vector Column
rs =
        Vector Column -> Value a
forall a. Columnable a => Vector Column -> Value a
Group (Vector Column -> Value a)
-> Either DataFrameException (Vector Column)
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Column -> Column -> Either DataFrameException Column)
-> Vector Column
-> Vector Column
-> Either DataFrameException (Vector Column)
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> m c) -> Vector a -> Vector b -> m (Vector c)
V.zipWithM ((c -> b -> a)
-> Column -> Column -> Either DataFrameException Column
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
(a -> b -> c)
-> Column -> Column -> Either DataFrameException Column
zipWithColumns c -> b -> a
f) Vector Column
ls Vector Column
rs
-- Shape mismatches: aggregated vs. non-aggregated.
liftValue2 c -> b -> a
_ (Flat Column
_) (Group Vector Column
_) =
    DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$ Text -> Text -> DataFrameException
AggregatedAndNonAggregatedException Text
"aggregated" Text
"non-aggregated"
liftValue2 c -> b -> a
_ (Group Vector Column
_) (Flat Column
_) =
    DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$ Text -> Text -> DataFrameException
AggregatedAndNonAggregatedException Text
"non-aggregated" Text
"aggregated"
liftValue2 c -> b -> a
_ (Group Vector Column
_) (Group Vector Column
_) =
    DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$ Text -> DataFrameException
InternalException Text
"Group count mismatch in binary operation"

-- | Branch on a boolean 'Value', selecting from two same-typed 'Value's.
branchValue ::
    forall a.
    (Columnable a) =>
    Value Bool ->
    Value a ->
    Value a ->
    Either DataFrameException (Value a)
branchValue :: forall a.
Columnable a =>
Value Bool
-> Value a -> Value a -> Either DataFrameException (Value a)
branchValue (Scalar Bool
True) Value a
l Value a
_ = Value a -> Either DataFrameException (Value a)
forall a b. b -> Either a b
Right Value a
l
branchValue (Scalar Bool
False) Value a
_ Value a
r = Value a -> Either DataFrameException (Value a)
forall a b. b -> Either a b
Right Value a
r
branchValue Value Bool
cond (Scalar a
l) (Scalar a
r) =
    (Bool -> a) -> Value Bool -> Either DataFrameException (Value a)
forall b a.
(Columnable b, Columnable a) =>
(b -> a) -> Value b -> Either DataFrameException (Value a)
liftValue (\Bool
c -> if Bool
c then a
l else a
r) Value Bool
cond
branchValue Value Bool
cond (Scalar a
l) Value a
r =
    (Bool -> a -> a)
-> Value Bool -> Value a -> Either DataFrameException (Value a)
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Value c -> Value b -> Either DataFrameException (Value a)
liftValue2 (\Bool
c a
rv -> if Bool
c then a
l else a
rv) Value Bool
cond Value a
r
branchValue Value Bool
cond Value a
l (Scalar a
r) =
    (Bool -> a -> a)
-> Value Bool -> Value a -> Either DataFrameException (Value a)
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Value c -> Value b -> Either DataFrameException (Value a)
liftValue2 (\Bool
c a
lv -> if Bool
c then a
lv else a
r) Value Bool
cond Value a
l
branchValue (Flat Column
cc) (Flat Column
lc) (Flat Column
rc) =
    Column -> Value a
forall a. Columnable a => Column -> Value a
Flat (Column -> Value a)
-> Either DataFrameException Column
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a.
Columnable a =>
Column -> Column -> Column -> Either DataFrameException Column
branchColumn @a Column
cc Column
lc Column
rc
branchValue (Group Vector Column
cgs) (Group Vector Column
lgs) (Group Vector Column
rgs)
    | Vector Column -> Int
forall a. Vector a -> Int
V.length Vector Column
cgs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector Column -> Int
forall a. Vector a -> Int
V.length Vector Column
lgs
        Bool -> Bool -> Bool
&& Vector Column -> Int
forall a. Vector a -> Int
V.length Vector Column
lgs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector Column -> Int
forall a. Vector a -> Int
V.length Vector Column
rgs =
        Vector Column -> Value a
forall a. Columnable a => Vector Column -> Value a
Group
            (Vector Column -> Value a)
-> Either DataFrameException (Vector Column)
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int
-> (Int -> Either DataFrameException Column)
-> Either DataFrameException (Vector Column)
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> m a) -> m (Vector a)
V.generateM
                (Vector Column -> Int
forall a. Vector a -> Int
V.length Vector Column
cgs)
                ( \Int
i ->
                    forall a.
Columnable a =>
Column -> Column -> Column -> Either DataFrameException Column
branchColumn @a (Vector Column
cgs Vector Column -> Int -> Column
forall a. Vector a -> Int -> a
V.! Int
i) (Vector Column
lgs Vector Column -> Int -> Column
forall a. Vector a -> Int -> a
V.! Int
i) (Vector Column
rgs Vector Column -> Int -> Column
forall a. Vector a -> Int -> a
V.! Int
i)
                )
branchValue Value Bool
_ Value a
_ Value a
_ =
    DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
        Text -> Text -> DataFrameException
AggregatedAndNonAggregatedException
            Text
"if-then-else branches"
            Text
"mismatched shapes"

{- | Low-level column branch: given a boolean column and two same-typed
columns, produce the element-wise selection.
-}
branchColumn ::
    forall a.
    (Columnable a) =>
    Column ->
    Column ->
    Column ->
    Either DataFrameException Column
branchColumn :: forall a.
Columnable a =>
Column -> Column -> Column -> Either DataFrameException Column
branchColumn Column
cc Column
lc Column
rc = do
    Vector Bool
cs <- forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector @Bool @V.Vector Column
cc
    Vector a
ls <- forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector @a @V.Vector Column
lc
    Vector a
rs <- forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector @a @V.Vector Column
rc
    Column -> Either DataFrameException Column
forall a. a -> Either DataFrameException a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Column -> Either DataFrameException Column)
-> Column -> Either DataFrameException Column
forall a b. (a -> b) -> a -> b
$
        forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector @a (Vector a -> Column) -> Vector a -> Column
forall a b. (a -> b) -> a -> b
$
            (Bool -> a -> a -> a)
-> Vector Bool -> Vector a -> Vector a -> Vector a
forall a b c d.
(a -> b -> c -> d) -> Vector a -> Vector b -> Vector c -> Vector d
V.zipWith3 (\Bool
c a
l a
r -> if Bool
c then a
l else a
r) Vector Bool
cs Vector a
ls Vector a
rs

-------------------------------------------------------------------------------
-- Error enrichment
-------------------------------------------------------------------------------

{- | Wrap an interpretation step so that any 'TypeMismatchException' gets
annotated with the expression that was being evaluated.
-}
addContext ::
    (Show a) => Expr a -> Either DataFrameException b -> Either DataFrameException b
addContext :: forall a b.
Show a =>
Expr a
-> Either DataFrameException b -> Either DataFrameException b
addContext Expr a
expr = (DataFrameException -> DataFrameException)
-> Either DataFrameException b -> Either DataFrameException b
forall a b c. (a -> b) -> Either a c -> Either b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ([Char] -> DataFrameException -> DataFrameException
enrichError (Expr a -> [Char]
forall a. Show a => a -> [Char]
show Expr a
expr))

enrichError :: String -> DataFrameException -> DataFrameException
enrichError :: [Char] -> DataFrameException -> DataFrameException
enrichError [Char]
loc (TypeMismatchException TypeErrorContext a b
ctx) =
    TypeErrorContext a b -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
        TypeErrorContext a b
ctx
            { callingFunctionName =
                callingFunctionName ctx <|+> Just "eval"
            , errorColumnName =
                errorColumnName ctx <|+> Just loc
            }
  where
    -- Prefer the existing value; fall back to the new one.
    Maybe a
Nothing <|+> :: Maybe a -> Maybe a -> Maybe a
<|+> Maybe a
b = Maybe a
b
    Maybe a
a <|+> Maybe a
_ = Maybe a
a
enrichError [Char]
_ DataFrameException
e = DataFrameException
e

-------------------------------------------------------------------------------
-- Group slicing
-------------------------------------------------------------------------------

{- | Given a flat column and grouping metadata, produce one 'Column' per
group.  Each result column is an O(1) slice into a sorted copy of the
input — the sort happens once, not per-group.
-}
sliceGroups :: Column -> VU.Vector Int -> VU.Vector Int -> V.Vector Column
sliceGroups :: Column -> Vector Int -> Vector Int -> Vector Column
sliceGroups Column
col Vector Int
os Vector Int
indices = case Column
col of
    BoxedColumn Vector a
vec ->
        let !sorted :: Vector a
sorted = Vector a -> Vector Int -> Vector a
forall a. Vector a -> Vector Int -> Vector a
V.unsafeBackpermute Vector a
vec (Vector Int -> Vector Int
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
V.convert Vector Int
indices)
         in Int -> (Int -> Column) -> Vector Column
forall a. Int -> (Int -> a) -> Vector a
V.generate Int
nGroups ((Int -> Column) -> Vector Column)
-> (Int -> Column) -> Vector Column
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                Vector a -> Column
forall a. Columnable a => Vector a -> Column
BoxedColumn (Int -> Int -> Vector a -> Vector a
forall a. Int -> Int -> Vector a -> Vector a
V.unsafeSlice (Int -> Int
start Int
i) (Int -> Int
len Int
i) Vector a
sorted)
    UnboxedColumn Vector a
vec ->
        let !sorted :: Vector a
sorted = Vector a -> Vector Int -> Vector a
forall a. Unbox a => Vector a -> Vector Int -> Vector a
VU.unsafeBackpermute Vector a
vec Vector Int
indices
         in Int -> (Int -> Column) -> Vector Column
forall a. Int -> (Int -> a) -> Vector a
V.generate Int
nGroups ((Int -> Column) -> Vector Column)
-> (Int -> Column) -> Vector Column
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                Vector a -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
UnboxedColumn (Int -> Int -> Vector a -> Vector a
forall a. Unbox a => Int -> Int -> Vector a -> Vector a
VU.unsafeSlice (Int -> Int
start Int
i) (Int -> Int
len Int
i) Vector a
sorted)
    OptionalColumn Vector (Maybe a)
vec ->
        let !sorted :: Vector (Maybe a)
sorted = Vector (Maybe a) -> Vector Int -> Vector (Maybe a)
forall a. Vector a -> Vector Int -> Vector a
V.unsafeBackpermute Vector (Maybe a)
vec (Vector Int -> Vector Int
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
V.convert Vector Int
indices)
         in Int -> (Int -> Column) -> Vector Column
forall a. Int -> (Int -> a) -> Vector a
V.generate Int
nGroups ((Int -> Column) -> Vector Column)
-> (Int -> Column) -> Vector Column
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                Vector (Maybe a) -> Column
forall a. Columnable a => Vector (Maybe a) -> Column
OptionalColumn (Int -> Int -> Vector (Maybe a) -> Vector (Maybe a)
forall a. Int -> Int -> Vector a -> Vector a
V.unsafeSlice (Int -> Int
start Int
i) (Int -> Int
len Int
i) Vector (Maybe a)
sorted)
  where
    !nGroups :: Int
nGroups = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
os Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    start :: Int -> Int
start Int
i = Vector Int
os Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
`VU.unsafeIndex` Int
i
    len :: Int -> Int
len Int
i = Vector Int
os Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
`VU.unsafeIndex` (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int
start Int
i
{-# INLINE sliceGroups #-}

numGroups :: GroupedDataFrame -> Int
numGroups :: GroupedDataFrame -> Int
numGroups GroupedDataFrame
gdf = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length (GroupedDataFrame -> Vector Int
offsets GroupedDataFrame
gdf) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

-------------------------------------------------------------------------------
-- eval: the unified interpreter
-------------------------------------------------------------------------------

{- | Evaluate an expression in a given context, producing a 'Value'.
This single function replaces both the old @interpret@ (flat) and
@interpretAggregation@ (grouped) code paths.
-}
eval ::
    forall a.
    (Columnable a) =>
    Ctx -> Expr a -> Either DataFrameException (Value a)
-- Leaves -----------------------------------------------------------------

eval :: forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval Ctx
_ (Lit a
v) = Value a -> Either DataFrameException (Value a)
forall a b. b -> Either a b
Right (a -> Value a
forall a. Columnable a => a -> Value a
Scalar a
v)
eval (FlatCtx DataFrame
df) (Col Text
name) =
    case Text -> DataFrame -> Maybe Column
getColumn Text
name DataFrame
df of
        Maybe Column
Nothing ->
            DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
                Text -> Text -> [Text] -> DataFrameException
ColumnNotFoundException Text
name Text
"" (Map Text Int -> [Text]
forall k a. Map k a -> [k]
M.keys (Map Text Int -> [Text]) -> Map Text Int -> [Text]
forall a b. (a -> b) -> a -> b
$ DataFrame -> Map Text Int
columnIndices DataFrame
df)
        Just Column
c -> Value a -> Either DataFrameException (Value a)
forall a b. b -> Either a b
Right (Column -> Value a
forall a. Columnable a => Column -> Value a
Flat Column
c)
eval (GroupCtx GroupedDataFrame
gdf) (Col Text
name) =
    case Text -> DataFrame -> Maybe Column
getColumn Text
name (GroupedDataFrame -> DataFrame
fullDataframe GroupedDataFrame
gdf) of
        Maybe Column
Nothing ->
            DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
                Text -> Text -> [Text] -> DataFrameException
ColumnNotFoundException
                    Text
name
                    Text
""
                    (Map Text Int -> [Text]
forall k a. Map k a -> [k]
M.keys (Map Text Int -> [Text]) -> Map Text Int -> [Text]
forall a b. (a -> b) -> a -> b
$ DataFrame -> Map Text Int
columnIndices (DataFrame -> Map Text Int) -> DataFrame -> Map Text Int
forall a b. (a -> b) -> a -> b
$ GroupedDataFrame -> DataFrame
fullDataframe GroupedDataFrame
gdf)
        Just Column
c ->
            Value a -> Either DataFrameException (Value a)
forall a b. b -> Either a b
Right
                ( Vector Column -> Value a
forall a. Columnable a => Vector Column -> Value a
Group
                    (Column -> Vector Int -> Vector Int -> Vector Column
sliceGroups Column
c (GroupedDataFrame -> Vector Int
offsets GroupedDataFrame
gdf) (GroupedDataFrame -> Vector Int
valueIndices GroupedDataFrame
gdf))
                )
-- Unary ------------------------------------------------------------------

eval Ctx
ctx expr :: Expr a
expr@(Unary (UnaryOp b a
op :: UnaryOp b a) Expr b
inner) = Expr a
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b.
Show a =>
Expr a
-> Either DataFrameException b -> Either DataFrameException b
addContext Expr a
expr (Either DataFrameException (Value a)
 -> Either DataFrameException (Value a))
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$ do
    Value b
v <- forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval @b Ctx
ctx Expr b
inner
    (b -> a) -> Value b -> Either DataFrameException (Value a)
forall b a.
(Columnable b, Columnable a) =>
(b -> a) -> Value b -> Either DataFrameException (Value a)
liftValue (UnaryOp b a -> b -> a
forall a b. UnaryOp a b -> a -> b
unaryFn UnaryOp b a
op) Value b
v

-- Binary -----------------------------------------------------------------

eval Ctx
ctx expr :: Expr a
expr@(Binary (BinaryOp c b a
op :: BinaryOp c b a) Expr c
left Expr b
right) =
    Expr a
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b.
Show a =>
Expr a
-> Either DataFrameException b -> Either DataFrameException b
addContext Expr a
expr (Either DataFrameException (Value a)
 -> Either DataFrameException (Value a))
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$ do
        Value c
l <- forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval @c Ctx
ctx Expr c
left
        Value b
r <- forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval @b Ctx
ctx Expr b
right
        (c -> b -> a)
-> Value c -> Value b -> Either DataFrameException (Value a)
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a)
-> Value c -> Value b -> Either DataFrameException (Value a)
liftValue2 (BinaryOp c b a -> c -> b -> a
forall a b c. BinaryOp a b c -> a -> b -> c
binaryFn BinaryOp c b a
op) Value c
l Value b
r

-- If ---------------------------------------------------------------------

eval Ctx
ctx expr :: Expr a
expr@(If Expr Bool
cond Expr a
l Expr a
r) = Expr a
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b.
Show a =>
Expr a
-> Either DataFrameException b -> Either DataFrameException b
addContext Expr a
expr (Either DataFrameException (Value a)
 -> Either DataFrameException (Value a))
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$ do
    Value Bool
c <- forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval @Bool Ctx
ctx Expr Bool
cond
    Value a
lv <- forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval @a Ctx
ctx Expr a
l
    Value a
rv <- forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval @a Ctx
ctx Expr a
r
    Value Bool
-> Value a -> Value a -> Either DataFrameException (Value a)
forall a.
Columnable a =>
Value Bool
-> Value a -> Value a -> Either DataFrameException (Value a)
branchValue Value Bool
c Value a
lv Value a
rv

-- Fast path: FoldAgg (seeded) on a bare Col in GroupCtx.
-- Avoids the O(n) backpermute in sliceGroups by folding directly over
-- permuted indices.  Only matches when inner is exactly (Col name).

eval (GroupCtx GroupedDataFrame
gdf) expr :: Expr a
expr@(Agg (FoldAgg Text
_ (Just a
seed) (a -> b -> a
f :: a -> b -> a)) (Col Text
name :: Expr b)) =
    Expr a
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b.
Show a =>
Expr a
-> Either DataFrameException b -> Either DataFrameException b
addContext Expr a
expr (Either DataFrameException (Value a)
 -> Either DataFrameException (Value a))
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
        case Text -> DataFrame -> Maybe Column
getColumn Text
name (GroupedDataFrame -> DataFrame
fullDataframe GroupedDataFrame
gdf) of
            Maybe Column
Nothing ->
                DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
                    Text -> Text -> [Text] -> DataFrameException
ColumnNotFoundException
                        Text
name
                        Text
""
                        (Map Text Int -> [Text]
forall k a. Map k a -> [k]
M.keys (Map Text Int -> [Text]) -> Map Text Int -> [Text]
forall a b. (a -> b) -> a -> b
$ DataFrame -> Map Text Int
columnIndices (DataFrame -> Map Text Int) -> DataFrame -> Map Text Int
forall a b. (a -> b) -> a -> b
$ GroupedDataFrame -> DataFrame
fullDataframe GroupedDataFrame
gdf)
            Just Column
col ->
                Column -> Value a
forall a. Columnable a => Column -> Value a
Flat (Column -> Value a)
-> Either DataFrameException Column
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall b acc.
(Columnable b, Columnable acc) =>
(acc -> b -> acc)
-> acc
-> Column
-> Vector Int
-> Int
-> Either DataFrameException Column
foldLinearGroups @b @a a -> b -> a
f a
seed Column
col (GroupedDataFrame -> Vector Int
rowToGroup GroupedDataFrame
gdf) (GroupedDataFrame -> Int
numGroups GroupedDataFrame
gdf)
-- Fast path: FoldAgg (seedless) on a bare Col in GroupCtx.

eval (GroupCtx GroupedDataFrame
gdf) expr :: Expr a
expr@(Agg (FoldAgg Text
_ Maybe a
Nothing (a -> b -> a
f :: a -> b -> a)) (Col Text
name :: Expr b)) =
    Expr a
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b.
Show a =>
Expr a
-> Either DataFrameException b -> Either DataFrameException b
addContext Expr a
expr (Either DataFrameException (Value a)
 -> Either DataFrameException (Value a))
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
        case TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) of
            Maybe (a :~: b)
Nothing ->
                DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
                    Text -> DataFrameException
InternalException
                        Text
"Type mismatch in seedless fold: \
                        \accumulator and element types must match"
            Just a :~: b
Refl ->
                case Text -> DataFrame -> Maybe Column
getColumn Text
name (GroupedDataFrame -> DataFrame
fullDataframe GroupedDataFrame
gdf) of
                    Maybe Column
Nothing ->
                        DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
                            Text -> Text -> [Text] -> DataFrameException
ColumnNotFoundException
                                Text
name
                                Text
""
                                (Map Text Int -> [Text]
forall k a. Map k a -> [k]
M.keys (Map Text Int -> [Text]) -> Map Text Int -> [Text]
forall a b. (a -> b) -> a -> b
$ DataFrame -> Map Text Int
columnIndices (DataFrame -> Map Text Int) -> DataFrame -> Map Text Int
forall a b. (a -> b) -> a -> b
$ GroupedDataFrame -> DataFrame
fullDataframe GroupedDataFrame
gdf)
                    Just Column
col ->
                        Column -> Value a
forall a. Columnable a => Column -> Value a
Flat (Column -> Value a) -> (Vector b -> Column) -> Vector b -> Value a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector b -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector
                            (Vector b -> Value a)
-> Either DataFrameException (Vector b)
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a.
Columnable a =>
(a -> a -> a)
-> Column
-> Vector Int
-> Vector Int
-> Either DataFrameException (Vector a)
foldl1DirectGroups @b a -> b -> a
b -> b -> b
f Column
col (GroupedDataFrame -> Vector Int
valueIndices GroupedDataFrame
gdf) (GroupedDataFrame -> Vector Int
offsets GroupedDataFrame
gdf)
-- Fast path: MergeAgg on a bare Col in GroupCtx.

eval
    (GroupCtx GroupedDataFrame
gdf)
    expr :: Expr a
expr@( Agg
                (MergeAgg Text
_ acc
seed (acc -> b -> acc
step :: acc -> b -> acc) acc -> acc -> acc
_ (acc -> a
finalize :: acc -> a))
                (Col Text
name :: Expr b)
            ) =
        Expr a
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b.
Show a =>
Expr a
-> Either DataFrameException b -> Either DataFrameException b
addContext Expr a
expr (Either DataFrameException (Value a)
 -> Either DataFrameException (Value a))
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
            case Text -> DataFrame -> Maybe Column
getColumn Text
name (GroupedDataFrame -> DataFrame
fullDataframe GroupedDataFrame
gdf) of
                Maybe Column
Nothing ->
                    DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
                        Text -> Text -> [Text] -> DataFrameException
ColumnNotFoundException
                            Text
name
                            Text
""
                            (Map Text Int -> [Text]
forall k a. Map k a -> [k]
M.keys (Map Text Int -> [Text]) -> Map Text Int -> [Text]
forall a b. (a -> b) -> a -> b
$ DataFrame -> Map Text Int
columnIndices (DataFrame -> Map Text Int) -> DataFrame -> Map Text Int
forall a b. (a -> b) -> a -> b
$ GroupedDataFrame -> DataFrame
fullDataframe GroupedDataFrame
gdf)
                Just Column
col ->
                    Column -> Value a
forall a. Columnable a => Column -> Value a
Flat
                        (Column -> Value a)
-> Either DataFrameException Column
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( forall b acc.
(Columnable b, Columnable acc) =>
(acc -> b -> acc)
-> acc
-> Column
-> Vector Int
-> Int
-> Either DataFrameException Column
foldLinearGroups @b acc -> b -> acc
step acc
seed Column
col (GroupedDataFrame -> Vector Int
rowToGroup GroupedDataFrame
gdf) (GroupedDataFrame -> Int
numGroups GroupedDataFrame
gdf)
                                Either DataFrameException Column
-> (Column -> Either DataFrameException Column)
-> Either DataFrameException Column
forall a b.
Either DataFrameException a
-> (a -> Either DataFrameException b)
-> Either DataFrameException b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (acc -> a) -> Column -> Either DataFrameException Column
forall b c.
(Columnable b, Columnable c) =>
(b -> c) -> Column -> Either DataFrameException Column
mapColumn acc -> a
finalize
                            )
-- Aggregation: CollectAgg ------------------------------------------------

eval Ctx
ctx expr :: Expr a
expr@(Agg (CollectAgg Text
_ (v b -> a
f :: v b -> a)) Expr b
inner) =
    Expr a
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b.
Show a =>
Expr a
-> Either DataFrameException b -> Either DataFrameException b
addContext Expr a
expr (Either DataFrameException (Value a)
 -> Either DataFrameException (Value a))
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$ do
        Value b
v <- forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval @b Ctx
ctx Expr b
inner
        case Value b
v of
            Scalar b
_ ->
                DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
                    Text -> DataFrameException
InternalException
                        Text
"Cannot apply a collection aggregation to a scalar"
            Flat Column
col ->
                a -> Value a
forall a. Columnable a => a -> Value a
Scalar (a -> Value a)
-> Either DataFrameException a
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (v :: * -> *) b a.
(Vector v b, Typeable v, Columnable b, Columnable a) =>
(v b -> a) -> Column -> Either DataFrameException a
applyCollect @v @b @a v b -> a
f Column
col
            Group Vector Column
gs ->
                Column -> Value a
forall a. Columnable a => Column -> Value a
Flat (Column -> Value a) -> (Vector a -> Column) -> Vector a -> Value a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector
                    (Vector a -> Value a)
-> Either DataFrameException (Vector a)
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Column -> Either DataFrameException a)
-> Vector Column -> Either DataFrameException (Vector a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM (forall (v :: * -> *) b a.
(Vector v b, Typeable v, Columnable b, Columnable a) =>
(v b -> a) -> Column -> Either DataFrameException a
applyCollect @v @b @a v b -> a
f) Vector Column
gs

-- Aggregation: FoldAgg with seed -----------------------------------------

eval Ctx
ctx expr :: Expr a
expr@(Agg (FoldAgg Text
_ (Just a
seed) (a -> b -> a
f :: a -> b -> a)) Expr b
inner) =
    Expr a
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b.
Show a =>
Expr a
-> Either DataFrameException b -> Either DataFrameException b
addContext Expr a
expr (Either DataFrameException (Value a)
 -> Either DataFrameException (Value a))
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$ do
        Value b
v <- forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval @b Ctx
ctx Expr b
inner
        case Value b
v of
            Scalar b
_ ->
                DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
                    Text -> DataFrameException
InternalException
                        Text
"Cannot apply a fold aggregation to a scalar"
            Flat Column
col ->
                a -> Value a
forall a. Columnable a => a -> Value a
Scalar (a -> Value a)
-> Either DataFrameException a
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a b.
(Columnable a, Columnable b) =>
(b -> a -> b) -> b -> Column -> Either DataFrameException b
foldlColumn @b @a a -> b -> a
f a
seed Column
col
            Group Vector Column
gs ->
                Column -> Value a
forall a. Columnable a => Column -> Value a
Flat (Column -> Value a) -> (Vector a -> Column) -> Vector a -> Value a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector
                    (Vector a -> Value a)
-> Either DataFrameException (Vector a)
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Column -> Either DataFrameException a)
-> Vector Column -> Either DataFrameException (Vector a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM (forall a b.
(Columnable a, Columnable b) =>
(b -> a -> b) -> b -> Column -> Either DataFrameException b
foldlColumn @b @a a -> b -> a
f a
seed) Vector Column
gs

-- Aggregation: MergeAgg --------------------------------------------------

eval
    Ctx
ctx
    expr :: Expr a
expr@( Agg
                (MergeAgg Text
_ acc
seed (acc -> b -> acc
step :: acc -> b -> acc) acc -> acc -> acc
_ (acc -> a
finalize :: acc -> a))
                (Expr b
inner :: Expr b)
            ) =
        Expr a
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b.
Show a =>
Expr a
-> Either DataFrameException b -> Either DataFrameException b
addContext Expr a
expr (Either DataFrameException (Value a)
 -> Either DataFrameException (Value a))
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$ do
            Value b
v <- forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval @b Ctx
ctx Expr b
inner
            case Value b
v of
                Scalar b
_ ->
                    DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
                        Text -> DataFrameException
InternalException
                            Text
"Cannot apply a merge aggregation to a scalar"
                Flat Column
col ->
                    a -> Value a
forall a. Columnable a => a -> Value a
Scalar (a -> Value a) -> (acc -> a) -> acc -> Value a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. acc -> a
finalize (acc -> Value a)
-> Either DataFrameException acc
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a b.
Columnable a =>
(b -> a -> b) -> b -> Column -> Either DataFrameException b
foldlColumnWith @b acc -> b -> acc
step acc
seed Column
col
                Group Vector Column
gs ->
                    Column -> Value a
forall a. Columnable a => Column -> Value a
Flat (Column -> Value a) -> (Vector a -> Column) -> Vector a -> Value a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector
                        (Vector a -> Value a)
-> Either DataFrameException (Vector a)
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Column -> Either DataFrameException a)
-> Vector Column -> Either DataFrameException (Vector a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM ((acc -> a)
-> Either DataFrameException acc -> Either DataFrameException a
forall a b.
(a -> b)
-> Either DataFrameException a -> Either DataFrameException b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap acc -> a
finalize (Either DataFrameException acc -> Either DataFrameException a)
-> (Column -> Either DataFrameException acc)
-> Column
-> Either DataFrameException a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b.
Columnable a =>
(b -> a -> b) -> b -> Column -> Either DataFrameException b
foldlColumnWith @b acc -> b -> acc
step acc
seed) Vector Column
gs

-- Aggregation: FoldAgg without seed (fold1) ------------------------------

eval Ctx
ctx expr :: Expr a
expr@(Agg (FoldAgg Text
_ Maybe a
Nothing (a -> b -> a
f :: a -> b -> a)) Expr b
inner) =
    Expr a
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b.
Show a =>
Expr a
-> Either DataFrameException b -> Either DataFrameException b
addContext Expr a
expr (Either DataFrameException (Value a)
 -> Either DataFrameException (Value a))
-> Either DataFrameException (Value a)
-> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
        case TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) of
            Maybe (a :~: b)
Nothing ->
                DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
                    Text -> DataFrameException
InternalException
                        Text
"Type mismatch in seedless fold: \
                        \accumulator and element types must match"
            Just a :~: b
Refl -> do
                Value b
v <- forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval @b Ctx
ctx Expr b
inner
                case Value b
v of
                    Scalar b
_ ->
                        DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Value a))
-> DataFrameException -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$
                            Text -> DataFrameException
InternalException
                                Text
"Cannot apply a fold aggregation to a scalar"
                    Flat Column
col ->
                        a -> Value a
forall a. Columnable a => a -> Value a
Scalar (a -> Value a)
-> Either DataFrameException a
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a.
Columnable a =>
(a -> a -> a) -> Column -> Either DataFrameException a
foldl1Column @a a -> a -> a
a -> b -> a
f Column
col
                    Group Vector Column
gs ->
                        Column -> Value a
forall a. Columnable a => Column -> Value a
Flat (Column -> Value a) -> (Vector a -> Column) -> Vector a -> Value a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector
                            (Vector a -> Value a)
-> Either DataFrameException (Vector a)
-> Either DataFrameException (Value a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Column -> Either DataFrameException a)
-> Vector Column -> Either DataFrameException (Vector a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM (forall a.
Columnable a =>
(a -> a -> a) -> Column -> Either DataFrameException a
foldl1Column @a a -> a -> a
a -> b -> a
f) Vector Column
gs

-------------------------------------------------------------------------------
-- Aggregation helpers
-------------------------------------------------------------------------------

{- | Apply a 'CollectAgg' function to a single column, extracting the
appropriate vector type and applying the aggregation function.
-}
applyCollect ::
    forall v b a.
    (VG.Vector v b, Typeable v, Columnable b, Columnable a) =>
    (v b -> a) -> Column -> Either DataFrameException a
applyCollect :: forall (v :: * -> *) b a.
(Vector v b, Typeable v, Columnable b, Columnable a) =>
(v b -> a) -> Column -> Either DataFrameException a
applyCollect v b -> a
f Column
col = v b -> a
f (v b -> a)
-> Either DataFrameException (v b) -> Either DataFrameException a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector @b @v Column
col

-------------------------------------------------------------------------------
-- Backward-compatible wrappers
-------------------------------------------------------------------------------

{- | Result of interpreting an expression in a grouped context.
Retained for backward compatibility with 'aggregate' and friends.
-}
data AggregationResult a
    = UnAggregated Column
    | Aggregated (TypedColumn a)

{- | Interpret an expression against a flat 'DataFrame', producing a
typed column.  This is the original top-level entry point; internally
it calls 'eval' and materialises the result.

NOTE: unlike the old implementation, 'Lit' values are no longer
eagerly broadcast.  The broadcast happens here, at the boundary,
via 'materialize'.
-}
interpret ::
    forall a.
    (Columnable a) =>
    DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)
interpret :: forall a.
Columnable a =>
DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)
interpret DataFrame
df Expr a
expr = do
    Value a
v <- Ctx -> Expr a -> Either DataFrameException (Value a)
forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval (DataFrame -> Ctx
FlatCtx DataFrame
df) Expr a
expr
    TypedColumn a -> Either DataFrameException (TypedColumn a)
forall a. a -> Either DataFrameException a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypedColumn a -> Either DataFrameException (TypedColumn a))
-> TypedColumn a -> Either DataFrameException (TypedColumn a)
forall a b. (a -> b) -> a -> b
$ Column -> TypedColumn a
forall a. Columnable a => Column -> TypedColumn a
TColumn (Column -> TypedColumn a) -> Column -> TypedColumn a
forall a b. (a -> b) -> a -> b
$ forall a. Columnable a => Int -> Value a -> Column
materialize @a ((Int, Int) -> Int
forall a b. (a, b) -> a
fst (DataFrame -> (Int, Int)
dataframeDimensions DataFrame
df)) Value a
v

{- | Interpret an expression against a 'GroupedDataFrame',
distinguishing aggregated results from bare column references.
Internally calls 'eval'.
-}
interpretAggregation ::
    forall a.
    (Columnable a) =>
    GroupedDataFrame ->
    Expr a ->
    Either DataFrameException (AggregationResult a)
interpretAggregation :: forall a.
Columnable a =>
GroupedDataFrame
-> Expr a -> Either DataFrameException (AggregationResult a)
interpretAggregation GroupedDataFrame
gdf Expr a
expr = do
    Value a
v <- Ctx -> Expr a -> Either DataFrameException (Value a)
forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval (GroupedDataFrame -> Ctx
GroupCtx GroupedDataFrame
gdf) Expr a
expr
    case Value a
v of
        Scalar a
a ->
            AggregationResult a
-> Either DataFrameException (AggregationResult a)
forall a b. b -> Either a b
Right (AggregationResult a
 -> Either DataFrameException (AggregationResult a))
-> AggregationResult a
-> Either DataFrameException (AggregationResult a)
forall a b. (a -> b) -> a -> b
$
                TypedColumn a -> AggregationResult a
forall a. TypedColumn a -> AggregationResult a
Aggregated (TypedColumn a -> AggregationResult a)
-> TypedColumn a -> AggregationResult a
forall a b. (a -> b) -> a -> b
$
                    Column -> TypedColumn a
forall a. Columnable a => Column -> TypedColumn a
TColumn (Column -> TypedColumn a) -> Column -> TypedColumn a
forall a b. (a -> b) -> a -> b
$
                        forall a. Columnable a => Int -> a -> Column
broadcastScalar @a (GroupedDataFrame -> Int
numGroups GroupedDataFrame
gdf) a
a
        Flat Column
col ->
            AggregationResult a
-> Either DataFrameException (AggregationResult a)
forall a b. b -> Either a b
Right (AggregationResult a
 -> Either DataFrameException (AggregationResult a))
-> AggregationResult a
-> Either DataFrameException (AggregationResult a)
forall a b. (a -> b) -> a -> b
$ TypedColumn a -> AggregationResult a
forall a. TypedColumn a -> AggregationResult a
Aggregated (TypedColumn a -> AggregationResult a)
-> TypedColumn a -> AggregationResult a
forall a b. (a -> b) -> a -> b
$ Column -> TypedColumn a
forall a. Columnable a => Column -> TypedColumn a
TColumn Column
col
        Group Vector Column
_ ->
            -- The Column payload is intentionally unused — the only
            -- call-site ('aggregate') immediately throws
            -- 'UnaggregatedException' on this constructor.
            AggregationResult a
-> Either DataFrameException (AggregationResult a)
forall a b. b -> Either a b
Right (AggregationResult a
 -> Either DataFrameException (AggregationResult a))
-> AggregationResult a
-> Either DataFrameException (AggregationResult a)
forall a b. (a -> b) -> a -> b
$ Column -> AggregationResult a
forall a. Column -> AggregationResult a
UnAggregated (Column -> AggregationResult a) -> Column -> AggregationResult a
forall a b. (a -> b) -> a -> b
$ forall a. Columnable a => Vector a -> Column
BoxedColumn @T.Text Vector Text
forall a. Vector a
V.empty