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

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 qualified Data.Vector.Unboxed.Mutable as VUM
import DataFrame.Errors
import DataFrame.Internal.Column
import DataFrame.Internal.DataFrame
import DataFrame.Internal.Expression
import qualified DataFrame.Internal.Grouping as G
import DataFrame.Internal.Types
import Type.Reflection (
    Typeable,
    typeRep,
 )

import Data.Int (Int16, Int32, Int64, Int8)

-- Specializations for common aggregation types to avoid dictionary overhead.
-- foldLinearGroups: mean accumulator
{-# SPECIALIZE foldLinearGroups ::
    (MeanAcc -> Double -> MeanAcc) ->
    MeanAcc ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (MeanAcc -> Float -> MeanAcc) ->
    MeanAcc ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (MeanAcc -> Int -> MeanAcc) ->
    MeanAcc ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (MeanAcc -> Int8 -> MeanAcc) ->
    MeanAcc ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (MeanAcc -> Int16 -> MeanAcc) ->
    MeanAcc ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (MeanAcc -> Int32 -> MeanAcc) ->
    MeanAcc ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (MeanAcc -> Int64 -> MeanAcc) ->
    MeanAcc ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
-- foldLinearGroups: count accumulator
{-# SPECIALIZE foldLinearGroups ::
    (Int -> Double -> Int) ->
    Int ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (Int -> Float -> Int) ->
    Int ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (Int -> Int -> Int) ->
    Int ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (Int -> Int8 -> Int) ->
    Int ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (Int -> Int16 -> Int) ->
    Int ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (Int -> Int32 -> Int) ->
    Int ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (Int -> Int64 -> Int) ->
    Int ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
-- foldLinearGroups: sum/min/max (acc == elem)
{-# SPECIALIZE foldLinearGroups ::
    (Double -> Double -> Double) ->
    Double ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (Float -> Float -> Float) ->
    Float ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (Int8 -> Int8 -> Int8) ->
    Int8 ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (Int16 -> Int16 -> Int16) ->
    Int16 ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (Int32 -> Int32 -> Int32) ->
    Int32 ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE foldLinearGroups ::
    (Int64 -> Int64 -> Int64) ->
    Int64 ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
    #-}

-- mapColumn: finalize
{-# SPECIALIZE mapColumn ::
    (MeanAcc -> Double) -> Column -> Either DataFrameException Column
    #-}
{-# SPECIALIZE mapColumn ::
    (Double -> Double) -> Column -> Either DataFrameException Column
    #-}
{-# SPECIALIZE mapColumn ::
    (Float -> Float) -> Column -> Either DataFrameException Column
    #-}
{-# SPECIALIZE mapColumn ::
    (Int -> Int) -> Column -> Either DataFrameException Column
    #-}

-- zipWithColumns: binary ops
{-# SPECIALIZE zipWithColumns ::
    (Double -> Double -> Double) ->
    Column ->
    Column ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE zipWithColumns ::
    (Float -> Float -> Float) ->
    Column ->
    Column ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE zipWithColumns ::
    (Int -> Int -> Int) -> Column -> Column -> Either DataFrameException Column
    #-}
{-# SPECIALIZE zipWithColumns ::
    (Int8 -> Int8 -> Int8) -> Column -> Column -> Either DataFrameException Column
    #-}
{-# SPECIALIZE zipWithColumns ::
    (Int16 -> Int16 -> Int16) ->
    Column ->
    Column ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE zipWithColumns ::
    (Int32 -> Int32 -> Int32) ->
    Column ->
    Column ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE zipWithColumns ::
    (Int64 -> Int64 -> Int64) ->
    Column ->
    Column ->
    Either DataFrameException Column
    #-}
-- Bool-returning binary comparators (hot path for Expr Bool used in
-- DecisionTree splits)
{-# SPECIALIZE zipWithColumns ::
    (Double -> Double -> Bool) ->
    Column ->
    Column ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE zipWithColumns ::
    (Float -> Float -> Bool) ->
    Column ->
    Column ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE zipWithColumns ::
    (Int -> Int -> Bool) ->
    Column ->
    Column ->
    Either DataFrameException Column
    #-}
{-# SPECIALIZE zipWithColumns ::
    (Bool -> Bool -> Bool) ->
    Column ->
    Column ->
    Either DataFrameException Column
    #-}

-- Bool-mapping unary ops (e.g. 'not')
{-# SPECIALIZE mapColumn ::
    (Bool -> Bool) -> Column -> Either DataFrameException Column
    #-}

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

instance (Show a) => Show (Value a) where
    show :: Value a -> String
show (Scalar a
v) = a -> String
forall a. Show a => a -> String
show a
v
    show (Flat Column
v) = Column -> String
forall a. Show a => a -> String
show Column
v
    show (Group Vector Column
v) = Vector Column -> String
forall a. Show a => a -> String
show Vector Column
v

-- | 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
_) =
    String -> Column
forall a. HasCallStack => String -> a
error String
"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
{-# INLINEABLE liftValue #-}

{- | Apply a binary function to two 'Value's. When one side is a 'Scalar' the
operation degenerates to 'liftValue', recovering the old @Binary op (Lit l) right@
special cases without explicit pattern matches.
-}
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
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"
{-# INLINEABLE liftValue2 #-}

-- | 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"
{-# INLINEABLE branchValue #-}

{- | 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 (String -> DataFrameException -> DataFrameException
enrichError (Expr a -> String
forall a. Show a => a -> String
show Expr a
expr))

enrichError :: String -> DataFrameException -> DataFrameException
enrichError :: String -> DataFrameException -> DataFrameException
enrichError String
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
    Maybe a
Nothing <|+> :: Maybe a -> Maybe a -> Maybe a
<|+> Maybe a
b = Maybe a
b
    Maybe a
a <|+> Maybe a
_ = Maybe a
a
enrichError String
_ 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
    PackedText Maybe Bitmap
_ PackedTextData
_ -> Column -> Vector Int -> Vector Int -> Vector Column
sliceGroups (Column -> Column
materializePacked Column
col) Vector Int
os Vector Int
indices
    BoxedColumn Maybe Bitmap
bm Vector a
vec ->
        let !sorted :: Vector a
sorted =
                Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
V.generate
                    (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
indices)
                    ((Vector a
vec Vector a -> Int -> a
forall a. Vector a -> Int -> a
`V.unsafeIndex`) (Int -> a) -> (Int -> Int) -> Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector Int
indices Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
`VU.unsafeIndex`))
         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 ->
                Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn
                    ((Bitmap -> Bitmap) -> Maybe Bitmap -> Maybe Bitmap
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> Int -> Bitmap -> Bitmap
bitmapSlice (Int -> Int
start Int
i) (Int -> Int
len Int
i)) Maybe Bitmap
bm)
                    (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 Maybe Bitmap
bm 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 ->
                Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn
                    ((Bitmap -> Bitmap) -> Maybe Bitmap -> Maybe Bitmap
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> Int -> Bitmap -> Bitmap
bitmapSlice (Int -> Int
start Int
i) (Int -> Int
len Int
i)) Maybe Bitmap
bm)
                    (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)
  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

-- | Build the inverse of a permutation vector.
invertPermutation :: VU.Vector Int -> VU.Vector Int
invertPermutation :: Vector Int -> Vector Int
invertPermutation Vector Int
perm = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
    let !n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
perm
    MVector s Int
inv <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
n
    (Int -> Int -> ST s ()) -> Vector Int -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
(Int -> a -> m b) -> Vector a -> m ()
VU.imapM_ ((Int -> Int -> ST s ()) -> Int -> Int -> ST s ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip (MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
inv)) Vector Int
perm
    MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return MVector s Int
inv
{-# INLINE invertPermutation #-}

-------------------------------------------------------------------------------
-- promoteColumnWith: unified numeric / text coercion for CastWith
-------------------------------------------------------------------------------

{- | Coerce a column to type @a@, then apply @onResult@ to each element; the handler
selects the mode (like @cast@, @castWithDefault@, or @castEither@). Handles Double/
Float/Int coercion and 'reads'-parses Text; other mismatches return 'Left'.
-}
promoteColumnWith ::
    forall a b.
    (Columnable a, Columnable b, Read a) =>
    (Either String a -> b) -> Column -> Either DataFrameException Column
promoteColumnWith :: forall a b.
(Columnable a, Columnable b, Read a) =>
(Either String a -> b)
-> Column -> Either DataFrameException Column
promoteColumnWith Either String a -> b
onResult Column
col
    | forall a. Columnable a => Column -> Bool
hasElemType @b Column
col = Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right Column
col
    | forall a. Columnable a => Column -> Bool
hasElemType @a Column
col = forall b c.
(Columnable b, Columnable c) =>
(b -> c) -> Column -> Either DataFrameException Column
mapColumn @a (Either String a -> b
onResult (Either String a -> b) -> (a -> Either String a) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Either String a
forall a b. b -> Either a b
Right) Column
col
    | Just Either DataFrameException Column
result <- forall a b.
(Columnable a, Columnable b) =>
(Either String a -> b)
-> Column -> Maybe (Either DataFrameException Column)
tryMaybeWrap @a @b Either String a -> b
onResult Column
col = Either DataFrameException Column
result
    | Bool
otherwise =
        case TypeRep a -> TypeRep Double -> Maybe (a :~: Double)
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 @Double) of
            Just a :~: Double
Refl -> (Either String Double -> b)
-> Column -> Either DataFrameException Column
forall b.
Columnable b =>
(Either String Double -> b)
-> Column -> Either DataFrameException Column
promoteToDoubleWith Either String a -> b
Either String Double -> b
onResult Column
col
            Maybe (a :~: Double)
Nothing ->
                case TypeRep a -> TypeRep Float -> Maybe (a :~: Float)
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 @Float) of
                    Just a :~: Float
Refl -> (Either String Float -> b)
-> Column -> Either DataFrameException Column
forall b.
Columnable b =>
(Either String Float -> b)
-> Column -> Either DataFrameException Column
promoteToFloatWith Either String a -> b
Either String Float -> b
onResult Column
col
                    Maybe (a :~: Float)
Nothing ->
                        case TypeRep a -> TypeRep Int -> Maybe (a :~: Int)
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 @Int) of
                            Just a :~: Int
Refl -> (Either String Int -> b)
-> Column -> Either DataFrameException Column
forall b.
Columnable b =>
(Either String Int -> b)
-> Column -> Either DataFrameException Column
promoteToIntWith Either String a -> b
Either String Int -> b
onResult Column
col
                            Maybe (a :~: Int)
Nothing -> forall a b.
(Columnable a, Columnable b, Read a) =>
(Either String a -> b)
-> Column -> Either DataFrameException Column
tryParseWith @a Either String a -> b
onResult Column
col

promoteToDoubleWith ::
    forall b.
    (Columnable b) =>
    (Either String Double -> b) -> Column -> Either DataFrameException Column
promoteToDoubleWith :: forall b.
Columnable b =>
(Either String Double -> b)
-> Column -> Either DataFrameException Column
promoteToDoubleWith Either String Double -> b
onResult Column
col = case Column
col of
    UnboxedColumn Maybe Bitmap
Nothing (Vector a
v :: VU.Vector c) ->
        case forall a. SBoolI (FloatingTypes a) => SBool (FloatingTypes a)
sFloating @c of
            SBool (FloatingTypes a)
STrue ->
                Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b
                        ((a -> b) -> Vector a -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Either String Double -> b
onResult (Either String Double -> b)
-> (a -> Either String Double) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Either String Double
forall a b. b -> Either a b
Right (Double -> Either String Double)
-> (a -> Double) -> a -> Either String Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac :: c -> Double)) (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
v))
            SBool (FloatingTypes a)
SFalse -> case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @c of
                SBool (IntegralTypes a)
STrue ->
                    Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b
                            ((a -> b) -> Vector a -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Either String Double -> b
onResult (Either String Double -> b)
-> (a -> Either String Double) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Either String Double
forall a b. b -> Either a b
Right (Double -> Either String Double)
-> (a -> Double) -> a -> Either String Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: c -> Double)) (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
v))
                SBool (IntegralTypes a)
SFalse -> forall src tgt.
(Typeable src, Typeable tgt) =>
Either DataFrameException Column
castMismatch @c @b
    UnboxedColumn (Just Bitmap
bm) (Vector a
v :: VU.Vector c) ->
        case forall a. SBoolI (FloatingTypes a) => SBool (FloatingTypes a)
sFloating @c of
            SBool (FloatingTypes a)
STrue ->
                Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b
                        ( Int -> (Int -> b) -> Vector b
forall a. Int -> (Int -> a) -> Vector a
V.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
v) ((Int -> b) -> Vector b) -> (Int -> b) -> Vector b
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                            if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i
                                then Either String Double -> b
onResult (Double -> Either String Double
forall a b. b -> Either a b
Right (a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
i) :: Double))
                                else Either String Double -> b
onResult (String -> Either String Double
forall a b. a -> Either a b
Left String
"null")
                        )
            SBool (FloatingTypes a)
SFalse -> case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @c of
                SBool (IntegralTypes a)
STrue ->
                    Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b
                            ( Int -> (Int -> b) -> Vector b
forall a. Int -> (Int -> a) -> Vector a
V.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
v) ((Int -> b) -> Vector b) -> (Int -> b) -> Vector b
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                                if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i
                                    then Either String Double -> b
onResult (Double -> Either String Double
forall a b. b -> Either a b
Right (a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
i) :: Double))
                                    else Either String Double -> b
onResult (String -> Either String Double
forall a b. a -> Either a b
Left String
"null")
                            )
                SBool (IntegralTypes a)
SFalse -> forall src tgt.
(Typeable src, Typeable tgt) =>
Either DataFrameException Column
castMismatch @c @b
    BoxedColumn Maybe Bitmap
_ Vector a
_ -> forall a b.
(Columnable a, Columnable b, Read a) =>
(Either String a -> b)
-> Column -> Either DataFrameException Column
tryParseWith @Double Either String Double -> b
onResult Column
col
    PackedText Maybe Bitmap
_ PackedTextData
_ -> (Either String Double -> b)
-> Column -> Either DataFrameException Column
forall b.
Columnable b =>
(Either String Double -> b)
-> Column -> Either DataFrameException Column
promoteToDoubleWith Either String Double -> b
onResult (Column -> Column
materializePacked Column
col)

promoteToFloatWith ::
    forall b.
    (Columnable b) =>
    (Either String Float -> b) -> Column -> Either DataFrameException Column
promoteToFloatWith :: forall b.
Columnable b =>
(Either String Float -> b)
-> Column -> Either DataFrameException Column
promoteToFloatWith Either String Float -> b
onResult Column
col = case Column
col of
    UnboxedColumn Maybe Bitmap
Nothing (Vector a
v :: VU.Vector c) ->
        case forall a. SBoolI (FloatingTypes a) => SBool (FloatingTypes a)
sFloating @c of
            SBool (FloatingTypes a)
STrue ->
                Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b
                        ((a -> b) -> Vector a -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Either String Float -> b
onResult (Either String Float -> b) -> (a -> Either String Float) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Float -> Either String Float
forall a b. b -> Either a b
Right (Float -> Either String Float)
-> (a -> Float) -> a -> Either String Float
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Float
forall a b. (Real a, Fractional b) => a -> b
realToFrac :: c -> Float)) (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
v))
            SBool (FloatingTypes a)
SFalse -> case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @c of
                SBool (IntegralTypes a)
STrue ->
                    Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b
                            ((a -> b) -> Vector a -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Either String Float -> b
onResult (Either String Float -> b) -> (a -> Either String Float) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Float -> Either String Float
forall a b. b -> Either a b
Right (Float -> Either String Float)
-> (a -> Float) -> a -> Either String Float
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: c -> Float)) (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
v))
                SBool (IntegralTypes a)
SFalse -> forall src tgt.
(Typeable src, Typeable tgt) =>
Either DataFrameException Column
castMismatch @c @b
    UnboxedColumn (Just Bitmap
bm) (Vector a
v :: VU.Vector c) ->
        case forall a. SBoolI (FloatingTypes a) => SBool (FloatingTypes a)
sFloating @c of
            SBool (FloatingTypes a)
STrue ->
                Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b
                        ( Int -> (Int -> b) -> Vector b
forall a. Int -> (Int -> a) -> Vector a
V.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
v) ((Int -> b) -> Vector b) -> (Int -> b) -> Vector b
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                            if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i
                                then Either String Float -> b
onResult (Float -> Either String Float
forall a b. b -> Either a b
Right (a -> Float
forall a b. (Real a, Fractional b) => a -> b
realToFrac (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
i) :: Float))
                                else Either String Float -> b
onResult (String -> Either String Float
forall a b. a -> Either a b
Left String
"null")
                        )
            SBool (FloatingTypes a)
SFalse -> case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @c of
                SBool (IntegralTypes a)
STrue ->
                    Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b
                            ( Int -> (Int -> b) -> Vector b
forall a. Int -> (Int -> a) -> Vector a
V.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
v) ((Int -> b) -> Vector b) -> (Int -> b) -> Vector b
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                                if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i
                                    then Either String Float -> b
onResult (Float -> Either String Float
forall a b. b -> Either a b
Right (a -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
i) :: Float))
                                    else Either String Float -> b
onResult (String -> Either String Float
forall a b. a -> Either a b
Left String
"null")
                            )
                SBool (IntegralTypes a)
SFalse -> forall src tgt.
(Typeable src, Typeable tgt) =>
Either DataFrameException Column
castMismatch @c @b
    BoxedColumn Maybe Bitmap
_ Vector a
_ -> forall a b.
(Columnable a, Columnable b, Read a) =>
(Either String a -> b)
-> Column -> Either DataFrameException Column
tryParseWith @Float Either String Float -> b
onResult Column
col
    PackedText Maybe Bitmap
_ PackedTextData
_ -> (Either String Float -> b)
-> Column -> Either DataFrameException Column
forall b.
Columnable b =>
(Either String Float -> b)
-> Column -> Either DataFrameException Column
promoteToFloatWith Either String Float -> b
onResult (Column -> Column
materializePacked Column
col)

promoteToIntWith ::
    forall b.
    (Columnable b) =>
    (Either String Int -> b) -> Column -> Either DataFrameException Column
promoteToIntWith :: forall b.
Columnable b =>
(Either String Int -> b)
-> Column -> Either DataFrameException Column
promoteToIntWith Either String Int -> b
onResult Column
col = case Column
col of
    UnboxedColumn Maybe Bitmap
Nothing (Vector a
v :: VU.Vector c) ->
        case forall a. SBoolI (FloatingTypes a) => SBool (FloatingTypes a)
sFloating @c of
            SBool (FloatingTypes a)
STrue ->
                Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b
                        ((a -> b) -> Vector a -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Either String Int -> b
onResult (Either String Int -> b) -> (a -> Either String Int) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Either String Int
forall a b. b -> Either a b
Right (Int -> Either String Int) -> (a -> Int) -> a -> Either String Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Double -> Int) -> (a -> Double) -> a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac :: c -> Double))) (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
v))
            SBool (FloatingTypes a)
SFalse -> case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @c of
                SBool (IntegralTypes a)
STrue ->
                    Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b
                            ((a -> b) -> Vector a -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Either String Int -> b
onResult (Either String Int -> b) -> (a -> Either String Int) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Either String Int
forall a b. b -> Either a b
Right (Int -> Either String Int) -> (a -> Int) -> a -> Either String Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: c -> Int)) (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
v))
                SBool (IntegralTypes a)
SFalse -> forall src tgt.
(Typeable src, Typeable tgt) =>
Either DataFrameException Column
castMismatch @c @b
    UnboxedColumn (Just Bitmap
bm) (Vector a
v :: VU.Vector c) ->
        case forall a. SBoolI (FloatingTypes a) => SBool (FloatingTypes a)
sFloating @c of
            SBool (FloatingTypes a)
STrue ->
                Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b
                        ( Int -> (Int -> b) -> Vector b
forall a. Int -> (Int -> a) -> Vector a
V.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
v) ((Int -> b) -> Vector b) -> (Int -> b) -> Vector b
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                            if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i
                                then Either String Int -> b
onResult (Int -> Either String Int
forall a b. b -> Either a b
Right (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
i) :: Double)))
                                else Either String Int -> b
onResult (String -> Either String Int
forall a b. a -> Either a b
Left String
"null")
                        )
            SBool (FloatingTypes a)
SFalse -> case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @c of
                SBool (IntegralTypes a)
STrue ->
                    Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b
                            ( Int -> (Int -> b) -> Vector b
forall a. Int -> (Int -> a) -> Vector a
V.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
v) ((Int -> b) -> Vector b) -> (Int -> b) -> Vector b
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                                if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i
                                    then Either String Int -> b
onResult (Int -> Either String Int
forall a b. b -> Either a b
Right (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
i) :: Int))
                                    else Either String Int -> b
onResult (String -> Either String Int
forall a b. a -> Either a b
Left String
"null")
                            )
                SBool (IntegralTypes a)
SFalse -> forall src tgt.
(Typeable src, Typeable tgt) =>
Either DataFrameException Column
castMismatch @c @b
    BoxedColumn Maybe Bitmap
_ Vector a
_ -> forall a b.
(Columnable a, Columnable b, Read a) =>
(Either String a -> b)
-> Column -> Either DataFrameException Column
tryParseWith @Int Either String Int -> b
onResult Column
col
    PackedText Maybe Bitmap
_ PackedTextData
_ -> (Either String Int -> b)
-> Column -> Either DataFrameException Column
forall b.
Columnable b =>
(Either String Int -> b)
-> Column -> Either DataFrameException Column
promoteToIntWith Either String Int -> b
onResult (Column -> Column
materializePacked Column
col)

-- | Single parse primitive: apply @onResult@ to the result of 'reads'.
parseWith :: (Read a) => (Either String a -> b) -> String -> b
parseWith :: forall a b. Read a => (Either String a -> b) -> String -> b
parseWith Either String a -> b
f String
s = case ReadS a
forall a. Read a => ReadS a
reads String
s of
    [(a
x, String
"")] -> Either String a -> b
f (a -> Either String a
forall a b. b -> Either a b
Right a
x)
    [(a, String)]
_ -> case ReadS a
forall a. Read a => ReadS a
reads (ShowS
forall a. Show a => a -> String
show String
s) of
        [(a
x, String
"")] -> Either String a -> b
f (a -> Either String a
forall a b. b -> Either a b
Right a
x)
        [(a, String)]
_ -> Either String a -> b
f (String -> Either String a
forall a b. a -> Either a b
Left String
s)

tryParseWith ::
    forall a b.
    (Columnable a, Columnable b, Read a) =>
    (Either String a -> b) -> Column -> Either DataFrameException Column
tryParseWith :: forall a b.
(Columnable a, Columnable b, Read a) =>
(Either String a -> b)
-> Column -> Either DataFrameException Column
tryParseWith Either String a -> b
onResult Column
col = case Column
col of
    PackedText Maybe Bitmap
_ PackedTextData
_ -> (Either String a -> b)
-> Column -> Either DataFrameException Column
forall a b.
(Columnable a, Columnable b, Read a) =>
(Either String a -> b)
-> Column -> Either DataFrameException Column
tryParseWith Either String a -> b
onResult (Column -> Column
materializePacked Column
col)
    BoxedColumn Maybe Bitmap
bm (Vector a
v :: V.Vector c) ->
        case TypeRep a -> TypeRep String -> Maybe (a :~: String)
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 @c) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @String) of
            Just a :~: String
Refl -> case Maybe Bitmap
bm of
                Maybe Bitmap
Nothing -> Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b (Vector b -> Column) -> Vector b -> Column
forall a b. (a -> b) -> a -> b
$ (String -> b) -> Vector String -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
V.map ((Either String a -> b) -> String -> b
forall a b. Read a => (Either String a -> b) -> String -> b
parseWith Either String a -> b
onResult) Vector a
Vector String
v
                Just Bitmap
bitmap ->
                    Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b (Vector b -> Column) -> Vector b -> Column
forall a b. (a -> b) -> a -> b
$
                            (Int -> String -> b) -> Vector String -> Vector b
forall a b. (Int -> a -> b) -> Vector a -> Vector b
V.imap
                                ( \Int
i String
x ->
                                    if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bitmap Int
i then (Either String a -> b) -> String -> b
forall a b. Read a => (Either String a -> b) -> String -> b
parseWith Either String a -> b
onResult String
x else Either String a -> b
onResult (String -> Either String a
forall a b. a -> Either a b
Left String
"null")
                                )
                                Vector a
Vector String
v
            Maybe (a :~: String)
Nothing ->
                case TypeRep a -> TypeRep Text -> Maybe (a :~: Text)
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 @c) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @T.Text) of
                    Just a :~: Text
Refl -> case Maybe Bitmap
bm of
                        Maybe Bitmap
Nothing -> Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b (Vector b -> Column) -> Vector b -> Column
forall a b. (a -> b) -> a -> b
$ (Text -> b) -> Vector Text -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
V.map ((Either String a -> b) -> String -> b
forall a b. Read a => (Either String a -> b) -> String -> b
parseWith Either String a -> b
onResult (String -> b) -> (Text -> String) -> Text -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack) Vector a
Vector Text
v
                        Just Bitmap
bitmap ->
                            Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b (Vector b -> Column) -> Vector b -> Column
forall a b. (a -> b) -> a -> b
$
                                    (Int -> Text -> b) -> Vector Text -> Vector b
forall a b. (Int -> a -> b) -> Vector a -> Vector b
V.imap
                                        ( \Int
i Text
x ->
                                            if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bitmap Int
i
                                                then (Either String a -> b) -> String -> b
forall a b. Read a => (Either String a -> b) -> String -> b
parseWith Either String a -> b
onResult (Text -> String
T.unpack Text
x)
                                                else Either String a -> b
onResult (String -> Either String a
forall a b. a -> Either a b
Left String
"null")
                                        )
                                        Vector a
Vector Text
v
                    Maybe (a :~: Text)
Nothing -> forall src tgt.
(Typeable src, Typeable tgt) =>
Either DataFrameException Column
castMismatch @c @b
    UnboxedColumn Maybe Bitmap
bm (Vector a
v :: VU.Vector c) -> case Maybe Bitmap
bm of
        Maybe Bitmap
Nothing -> Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b (Vector b -> Column) -> Vector b -> Column
forall a b. (a -> b) -> a -> b
$ (a -> b) -> Vector a -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
V.map ((Either String a -> b) -> String -> b
forall a b. Read a => (Either String a -> b) -> String -> b
parseWith Either String a -> b
onResult (String -> b) -> (a -> String) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> String
forall a. Show a => a -> String
show) (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
V.convert Vector a
v)
        Just Bitmap
bitmap ->
            Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b (Vector b -> Column) -> Vector b -> Column
forall a b. (a -> b) -> a -> b
$
                    (Int -> a -> b) -> Vector a -> Vector b
forall a b. (Int -> a -> b) -> Vector a -> Vector b
V.imap
                        ( \Int
i a
x ->
                            if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bitmap Int
i
                                then (Either String a -> b) -> String -> b
forall a b. Read a => (Either String a -> b) -> String -> b
parseWith Either String a -> b
onResult (a -> String
forall a. Show a => a -> String
show a
x)
                                else Either String a -> b
onResult (String -> Either String a
forall a b. a -> Either a b
Left String
"null")
                        )
                        (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
V.convert Vector a
v)

{- | When output type @b@ is @Maybe c@ (or @Maybe (Maybe c)@) and the column stores
plain @c@, wrap each element in 'Just' (the double-Maybe case collapses to a single
@Maybe c@). Returns 'Nothing' when neither condition holds.
-}
tryMaybeWrap ::
    forall a b.
    (Columnable a, Columnable b) =>
    (Either String a -> b) -> Column -> Maybe (Either DataFrameException Column)
tryMaybeWrap :: forall a b.
(Columnable a, Columnable b) =>
(Either String a -> b)
-> Column -> Maybe (Either DataFrameException Column)
tryMaybeWrap Either String a -> b
_onResult Column
col = case Column
col of
    UnboxedColumn Maybe Bitmap
Nothing (Vector a
v :: VU.Vector c) ->
        let wrapped :: Vector (Maybe a)
wrapped = (a -> Maybe a) -> Vector a -> Vector (Maybe a)
forall a b. (a -> b) -> Vector a -> Vector b
V.map a -> Maybe a
forall a. a -> Maybe a
Just (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
v) :: V.Vector (Maybe c)
         in case TypeRep b -> TypeRep (Maybe a) -> Maybe (b :~: Maybe a)
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 @b) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @(Maybe c)) of
                Just b :~: Maybe a
Refl -> Either DataFrameException Column
-> Maybe (Either DataFrameException Column)
forall a. a -> Maybe a
Just (Either DataFrameException Column
 -> Maybe (Either DataFrameException Column))
-> Either DataFrameException Column
-> Maybe (Either DataFrameException Column)
forall a b. (a -> b) -> a -> b
$ Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b Vector b
Vector (Maybe a)
wrapped
                Maybe (b :~: Maybe a)
Nothing ->
                    case TypeRep b
-> TypeRep (Maybe (Maybe a)) -> Maybe (b :~: Maybe (Maybe a))
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 @b) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @(Maybe (Maybe c))) of
                        Just b :~: Maybe (Maybe a)
_ -> Either DataFrameException Column
-> Maybe (Either DataFrameException Column)
forall a. a -> Maybe a
Just (Either DataFrameException Column
 -> Maybe (Either DataFrameException Column))
-> Either DataFrameException Column
-> Maybe (Either DataFrameException Column)
forall a b. (a -> b) -> a -> b
$ Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @(Maybe c) Vector (Maybe a)
wrapped
                        Maybe (b :~: Maybe (Maybe a))
Nothing -> Maybe (Either DataFrameException Column)
forall a. Maybe a
Nothing
    BoxedColumn Maybe Bitmap
Nothing (Vector a
v :: V.Vector c) ->
        let wrapped :: Vector (Maybe a)
wrapped = (a -> Maybe a) -> Vector a -> Vector (Maybe a)
forall a b. (a -> b) -> Vector a -> Vector b
V.map a -> Maybe a
forall a. a -> Maybe a
Just Vector a
v :: V.Vector (Maybe c)
         in case TypeRep b -> TypeRep (Maybe a) -> Maybe (b :~: Maybe a)
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 @b) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @(Maybe c)) of
                Just b :~: Maybe a
Refl -> Either DataFrameException Column
-> Maybe (Either DataFrameException Column)
forall a. a -> Maybe a
Just (Either DataFrameException Column
 -> Maybe (Either DataFrameException Column))
-> Either DataFrameException Column
-> Maybe (Either DataFrameException Column)
forall a b. (a -> b) -> a -> b
$ Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @b Vector b
Vector (Maybe a)
wrapped
                Maybe (b :~: Maybe a)
Nothing ->
                    case TypeRep b
-> TypeRep (Maybe (Maybe a)) -> Maybe (b :~: Maybe (Maybe a))
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 @b) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @(Maybe (Maybe c))) of
                        Just b :~: Maybe (Maybe a)
_ -> Either DataFrameException Column
-> Maybe (Either DataFrameException Column)
forall a. a -> Maybe a
Just (Either DataFrameException Column
 -> Maybe (Either DataFrameException Column))
-> Either DataFrameException Column
-> Maybe (Either DataFrameException Column)
forall a b. (a -> b) -> a -> b
$ Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (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 @(Maybe c) Vector (Maybe a)
wrapped
                        Maybe (b :~: Maybe (Maybe a))
Nothing -> Maybe (Either DataFrameException Column)
forall a. Maybe a
Nothing
    Column
_ -> Maybe (Either DataFrameException Column)
forall a. Maybe a
Nothing

castMismatch ::
    forall src tgt.
    (Typeable src, Typeable tgt) =>
    Either DataFrameException Column
castMismatch :: forall src tgt.
(Typeable src, Typeable tgt) =>
Either DataFrameException Column
castMismatch =
    DataFrameException -> Either DataFrameException Column
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException Column)
-> DataFrameException -> Either DataFrameException Column
forall a b. (a -> b) -> a -> b
$
        TypeErrorContext tgt src -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
            MkTypeErrorContext
                { userType :: Either String (TypeRep tgt)
userType = TypeRep tgt -> Either String (TypeRep tgt)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @tgt)
                , expectedType :: Either String (TypeRep src)
expectedType = TypeRep src -> Either String (TypeRep src)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @src)
                , callingFunctionName :: Maybe String
callingFunctionName = String -> Maybe String
forall a. a -> Maybe a
Just String
"cast"
                , errorColumnName :: Maybe String
errorColumnName = Maybe String
forall a. Maybe a
Nothing
                }

-------------------------------------------------------------------------------
-- 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)
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
ColumnsNotFoundException [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
            | forall a. Columnable a => Column -> Bool
hasElemType @a 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)
            | Bool
otherwise ->
                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
$
                    TypeErrorContext a () -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                        ( MkTypeErrorContext
                            { userType :: Either String (TypeRep a)
userType = TypeRep a -> Either String (TypeRep a)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a)
                            , expectedType :: Either String (TypeRep ())
expectedType = String -> Either String (TypeRep ())
forall a b. a -> Either a b
Left (Column -> String
columnTypeString Column
c)
                            , errorColumnName :: Maybe String
errorColumnName = String -> Maybe String
forall a. a -> Maybe a
Just (Text -> String
T.unpack Text
name)
                            , callingFunctionName :: Maybe String
callingFunctionName = String -> Maybe String
forall a. a -> Maybe a
Just String
"col"
                            } ::
                            TypeErrorContext a ()
                        )
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
ColumnsNotFoundException
                    [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
            | forall a. Columnable a => Column -> Bool
hasElemType @a 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)))
            | Bool
otherwise ->
                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
$
                    TypeErrorContext a () -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                        ( MkTypeErrorContext
                            { userType :: Either String (TypeRep a)
userType = TypeRep a -> Either String (TypeRep a)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a)
                            , expectedType :: Either String (TypeRep ())
expectedType = String -> Either String (TypeRep ())
forall a b. a -> Either a b
Left (Column -> String
columnTypeString Column
c)
                            , errorColumnName :: Maybe String
errorColumnName = String -> Maybe String
forall a. a -> Maybe a
Just (Text -> String
T.unpack Text
name)
                            , callingFunctionName :: Maybe String
callingFunctionName = String -> Maybe String
forall a. a -> Maybe a
Just String
"col"
                            } ::
                            TypeErrorContext a ()
                        )
eval (FlatCtx DataFrame
df) (CastWith Text
name Text
_tag Either String a1 -> a
onResult) =
    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
ColumnsNotFoundException [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 -> 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
<$> (Either String a1 -> a)
-> Column -> Either DataFrameException Column
forall a b.
(Columnable a, Columnable b, Read a) =>
(Either String a -> b)
-> Column -> Either DataFrameException Column
promoteColumnWith Either String a1 -> a
onResult Column
c
eval (GroupCtx GroupedDataFrame
gdf) (CastWith Text
name Text
_tag Either String a1 -> a
onResult) =
    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
ColumnsNotFoundException
                    [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 -> do
            Column
promoted <- (Either String a1 -> a)
-> Column -> Either DataFrameException Column
forall a b.
(Columnable a, Columnable b, Read a) =>
(Either String a -> b)
-> Column -> Either DataFrameException Column
promoteColumnWith Either String a1 -> a
onResult Column
c
            Value a -> Either DataFrameException (Value a)
forall a b. b -> Either a b
Right (Value a -> Either DataFrameException (Value a))
-> Value a -> Either DataFrameException (Value a)
forall a b. (a -> b) -> a -> b
$ Vector Column -> Value a
forall a. Columnable a => Vector Column -> Value a
Group (Column -> Vector Int -> Vector Int -> Vector Column
sliceGroups Column
promoted (GroupedDataFrame -> Vector Int
offsets GroupedDataFrame
gdf) (GroupedDataFrame -> Vector Int
valueIndices GroupedDataFrame
gdf))
eval Ctx
ctx (CastExprWith Text
_tag Either String a1 -> a
onResult (Expr src
inner :: Expr src)) = do
    Value src
v <- forall a.
Columnable a =>
Ctx -> Expr a -> Either DataFrameException (Value a)
eval @src Ctx
ctx Expr src
inner
    case Value src
v of
        Scalar src
s ->
            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
<$> (Either String a1 -> a)
-> Column -> Either DataFrameException Column
forall a b.
(Columnable a, Columnable b, Read a) =>
(Either String a -> b)
-> Column -> Either DataFrameException Column
promoteColumnWith Either String a1 -> a
onResult (forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @src [src
s])
        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
<$> (Either String a1 -> a)
-> Column -> Either DataFrameException Column
forall a b.
(Columnable a, Columnable b, Read a) =>
(Either String a -> b)
-> Column -> Either DataFrameException Column
promoteColumnWith Either String a1 -> a
onResult Column
col
        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 ((Either String a1 -> a)
-> Column -> Either DataFrameException Column
forall a b.
(Columnable a, Columnable b, Read a) =>
(Either String a -> b)
-> Column -> Either DataFrameException Column
promoteColumnWith Either String a1 -> a
onResult) Vector Column
gs
eval Ctx
ctx expr :: Expr a
expr@(Unary op b a
op (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
    (b -> a) -> Value b -> Either DataFrameException (Value a)
forall b a.
(Columnable b, Columnable a) =>
(b -> a) -> Value b -> Either DataFrameException (Value a)
liftValue (op b a -> b -> a
forall a b. op a b -> a -> b
forall (op :: * -> * -> *) a b. UnaryOp op => op a b -> a -> b
unaryFn op b a
op) Value b
v
eval Ctx
ctx expr :: Expr a
expr@(Binary op c b a
op (Expr c
left :: Expr c) (Expr b
right :: 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 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 (op c b a -> c -> b -> a
forall a b c. op a b c -> a -> b -> c
forall (op :: * -> * -> * -> *) a b c.
BinaryOp op =>
op a b c -> a -> b -> c
binaryFn op c b a
op) Value c
l Value b
r
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
eval (FlatCtx DataFrame
df) expr :: Expr a
expr@(Over [Text]
keys Expr a
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
    let gdf :: GroupedDataFrame
gdf = [Text] -> DataFrame -> GroupedDataFrame
G.groupBy [Text]
keys DataFrame
df
    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
inner
    case Value a
v of
        Scalar a
s ->
            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
s)
        Flat Column
groupCol ->
            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 (Vector Int -> Column -> Column
atIndicesStable (GroupedDataFrame -> Vector Int
rowToGroup GroupedDataFrame
gdf) Column
groupCol))
        Group Vector Column
groupCols -> do
            Column
sorted <- (Column -> Column -> Either DataFrameException Column)
-> Vector Column -> Either DataFrameException Column
forall (m :: * -> *) a.
Monad m =>
(a -> a -> m a) -> Vector a -> m a
V.fold1M' Column -> Column -> Either DataFrameException Column
concatColumns Vector Column
groupCols
            let inv :: Vector Int
inv = Vector Int -> Vector Int
invertPermutation (GroupedDataFrame -> Vector Int
valueIndices GroupedDataFrame
gdf)
            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 (Vector Int -> Column -> Column
atIndicesStable Vector Int
inv Column
sorted))
eval (GroupCtx GroupedDataFrame
_) expr :: Expr a
expr@(Over [Text]
_ Expr a
_) =
    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
$
        DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left
            ( Text -> DataFrameException
InternalException
                Text
"Over (window function) is not supported inside a grouped context"
            )
-- 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
ColumnsNotFoundException
                        [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
ColumnsNotFoundException
                                [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 a.
Columnable a =>
(a -> a -> a)
-> Column
-> Vector Int
-> Vector Int
-> Either DataFrameException Column
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
ColumnsNotFoundException
                            [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
                            )
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
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
x -> Value a -> Either DataFrameException (Value a)
forall a b. b -> Either a b
Right (Ctx -> a -> (a -> b -> a) -> b -> Value a
forall acc b.
Columnable acc =>
Ctx -> acc -> (acc -> b -> acc) -> b -> Value acc
broadcastFold Ctx
ctx a
seed a -> b -> a
f b
x)
            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
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
x -> case Ctx -> acc -> (acc -> b -> acc) -> b -> Value acc
forall acc b.
Columnable acc =>
Ctx -> acc -> (acc -> b -> acc) -> b -> Value acc
broadcastFold Ctx
ctx acc
seed acc -> b -> acc
step b
x of
                    Scalar acc
acc -> 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 (acc -> a
finalize acc
acc))
                    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
<$> forall b c.
(Columnable b, Columnable c) =>
(b -> c) -> Column -> Either DataFrameException Column
mapColumn @acc @a acc -> a
finalize Column
col
                    Group Vector Column
_ ->
                        DataFrameException -> Either DataFrameException (Value a)
forall a b. a -> Either a b
Left
                            ( Text -> DataFrameException
InternalException
                                Text
"broadcastFold unexpectedly produced a Group value"
                            )
                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, Columnable b) =>
(b -> a -> b) -> b -> Column -> Either DataFrameException b
foldlColumn @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, Columnable b) =>
(b -> a -> b) -> b -> Column -> Either DataFrameException b
foldlColumn @b acc -> b -> acc
step acc
seed) Vector Column
gs
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
"fold1 requires at least one element"
                    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

broadcastFold ::
    forall acc b.
    (Columnable acc) =>
    Ctx -> acc -> (acc -> b -> acc) -> b -> Value acc
broadcastFold :: forall acc b.
Columnable acc =>
Ctx -> acc -> (acc -> b -> acc) -> b -> Value acc
broadcastFold (FlatCtx DataFrame
df) acc
seed acc -> b -> acc
step b
x =
    let n :: Int
n = (Int, Int) -> Int
forall a b. (a, b) -> a
fst (DataFrame -> (Int, Int)
dataframeDimensions DataFrame
df)
     in acc -> Value acc
forall a. Columnable a => a -> Value a
Scalar (Int -> (acc -> b -> acc) -> acc -> b -> acc
forall acc b. Int -> (acc -> b -> acc) -> acc -> b -> acc
iterateStep Int
n acc -> b -> acc
step acc
seed b
x)
broadcastFold (GroupCtx GroupedDataFrame
gdf) acc
seed acc -> b -> acc
step b
x =
    let offs :: Vector Int
offs = GroupedDataFrame -> Vector Int
offsets GroupedDataFrame
gdf
        ng :: Int
ng = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
offs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        results :: Vector acc
results =
            Int -> (Int -> acc) -> Vector acc
forall a. Int -> (Int -> a) -> Vector a
V.generate Int
ng ((Int -> acc) -> Vector acc) -> (Int -> acc) -> Vector acc
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                let sz :: Int
sz = Vector Int
offs Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector Int
offs Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
i
                 in Int -> (acc -> b -> acc) -> acc -> b -> acc
forall acc b. Int -> (acc -> b -> acc) -> acc -> b -> acc
iterateStep Int
sz acc -> b -> acc
step acc
seed b
x
     in Column -> Value acc
forall a. Columnable a => Column -> Value a
Flat (Vector acc -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector Vector acc
results)

iterateStep :: Int -> (acc -> b -> acc) -> acc -> b -> acc
iterateStep :: forall acc b. Int -> (acc -> b -> acc) -> acc -> b -> acc
iterateStep Int
n acc -> b -> acc
step = Int -> acc -> b -> acc
go Int
n
  where
    go :: Int -> acc -> b -> acc
go Int
0 !acc
acc b
_ = acc
acc
    go Int
k !acc
acc b
x = Int -> acc -> b -> acc
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (acc -> b -> acc
step acc
acc b
x) b
x

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

{- | 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.
Calls 'eval' then 'materialize'; 'Lit' values are broadcast here at the boundary
rather than eagerly.
-}
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
_ ->
            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 => Maybe Bitmap -> Vector a -> Column
BoxedColumn @T.Text Maybe Bitmap
forall a. Maybe a
Nothing Vector Text
forall a. Vector a
V.empty