{-# LANGUAGE BangPatterns, CPP, DataKinds, EmptyCase,
             FlexibleInstances, PolyKinds, ScopedTypeVariables,
             TupleSections, TypeFamilies, TypeOperators,
             UndecidableInstances #-}
-- | Efficient in-memory (in-core) storage of tabular data.
module Frames.InCore where
import Control.Monad.Primitive
import Control.Monad.ST (runST)
import Data.Proxy
import Data.Text (Text)
import qualified Data.Vector as VB
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM
import qualified Data.Vector.Unboxed as VU
import Data.Vinyl (Rec(..))
import qualified Data.Vinyl as V
import Data.Vinyl.Functor (Compose(..), getCompose, ElField(..), (:.))
import Frames.Col
import Frames.Frame
import Frames.Rec
import Frames.RecF
#if __GLASGOW_HASKELL__ < 800
import GHC.Prim (RealWorld)
#endif
import GHC.TypeLits (KnownSymbol)
import GHC.Types (Symbol, Type)
import qualified Pipes as P
import qualified Pipes.Prelude as P

-- | The most efficient vector type for each column data type.
type family VectorFor t :: * -> *
type instance VectorFor Bool = VU.Vector
type instance VectorFor Int = VU.Vector
type instance VectorFor Float = VU.Vector
type instance VectorFor Double = VU.Vector
type instance VectorFor String = VB.Vector
type instance VectorFor Text = VB.Vector

-- | The mutable version of 'VectorFor' a particular type.
type VectorMFor a = VG.Mutable (VectorFor a)

-- | Since we stream into the in-memory representation, we use an
-- exponential growth strategy to resize arrays as more data is read
-- in. This is the initial capacity of each column.
initialCapacity :: Int
initialCapacity :: Int
initialCapacity = Int
128

-- | Mutable vector types for each column in a row.
type family VectorMs m rs where
  VectorMs m '[] = '[]
  VectorMs m (s :-> a ': rs) =
    s :-> VectorMFor a (PrimState m) a ': VectorMs m rs

-- | Immutable vector types for each column in a row.
type family Vectors rs where
  Vectors '[] = '[]
  Vectors (s :-> a ': rs) = s :-> VectorFor a a ': Vectors rs

-- | Tooling to allocate, grow, write to, freeze, and index into
-- records of vectors.
class RecVec (rs :: [(Symbol,Type)]) where
  allocRec   :: PrimMonad m
             => proxy rs -> Int -> m (Record (VectorMs m rs))
  freezeRec  :: PrimMonad m
             => proxy rs -> Int -> Record (VectorMs m rs)
             -> m (Record (Vectors rs))
  growRec    :: PrimMonad m
             => proxy rs -> Record (VectorMs m rs) -> m (Record (VectorMs m rs))
  writeRec   :: PrimMonad m
             => proxy rs -> Int -> Record (VectorMs m rs) -> Record rs -> m ()
  indexRec   :: proxy rs -> Int -> Record (Vectors rs) -> Record rs
  produceRec :: proxy rs -> Record (Vectors rs) -> V.Rec (((->) Int) :. ElField) rs

-- The use of the type families Vectors and VectorMs interferes with
-- GHC's pattern match exhaustiveness checker, so we write down dummy
-- cases to avoid warnings. This is probably a GHC bug as writing the
-- apparently missing pattern match causes a type checker error!

instance RecVec '[] where
  allocRec :: proxy '[] -> Int -> m (Record (VectorMs m '[]))
allocRec proxy '[]
_ Int
_ = Record '[] -> m (Record '[])
forall (m :: * -> *) a. Monad m => a -> m a
return Record '[]
forall u (a :: u -> *). Rec a '[]
V.RNil
  {-# INLINE allocRec #-}

  freezeRec :: proxy '[]
-> Int -> Record (VectorMs m '[]) -> m (Record (Vectors '[]))
freezeRec proxy '[]
_ Int
_ Record (VectorMs m '[])
V.RNil = Record '[] -> m (Record '[])
forall (m :: * -> *) a. Monad m => a -> m a
return Record '[]
forall u (a :: u -> *). Rec a '[]
V.RNil
#if __GLASGOW_HASKELL__ < 800
  freezeRec _ _ x = case x of
#endif
  {-# INLINE freezeRec #-}

  growRec :: proxy '[] -> Record (VectorMs m '[]) -> m (Record (VectorMs m '[]))
growRec proxy '[]
_ Record (VectorMs m '[])
V.RNil = Record '[] -> m (Record '[])
forall (m :: * -> *) a. Monad m => a -> m a
return Record '[]
forall u (a :: u -> *). Rec a '[]
V.RNil
#if __GLASGOW_HASKELL__ < 800
  growRec _ x = case x of
#endif
  {-# INLINE growRec #-}

  indexRec :: proxy '[] -> Int -> Record (Vectors '[]) -> Record '[]
indexRec proxy '[]
_ Int
_ Record (Vectors '[])
_ = Record '[]
forall u (a :: u -> *). Rec a '[]
V.RNil
  {-# INLINE indexRec #-}

  writeRec :: proxy '[] -> Int -> Record (VectorMs m '[]) -> Record '[] -> m ()
writeRec proxy '[]
_ Int
_ Record (VectorMs m '[])
V.RNil Record '[]
V.RNil = () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
#if __GLASGOW_HASKELL__ < 800
  writeRec _ _ x _ = case x of
#endif
  {-# INLINE writeRec #-}

  produceRec :: proxy '[] -> Record (Vectors '[]) -> Rec ((->) Int :. ElField) '[]
produceRec proxy '[]
_ Record (Vectors '[])
V.RNil = Rec ((->) Int :. ElField) '[]
forall u (a :: u -> *). Rec a '[]
V.RNil
#if __GLASGOW_HASKELL__ < 800
  produceRec _ x = case x of
#endif
  {-# INLINE produceRec #-}

instance forall s a rs.
  (VGM.MVector (VectorMFor a) a,
   VG.Vector (VectorFor a) a,
   KnownSymbol s, RecVec rs)
  => RecVec (s :-> a ': rs) where
  allocRec :: proxy ((s :-> a) : rs)
-> Int -> m (Record (VectorMs m ((s :-> a) : rs)))
allocRec proxy ((s :-> a) : rs)
_ Int
size = Mutable (VectorFor a) (PrimState m) a
-> Record (VectorMs m rs)
-> Record
     ((s :-> Mutable (VectorFor a) (PrimState m) a) : VectorMs m rs)
forall (s :: Symbol) a (rs :: [(Symbol, *)]).
KnownSymbol s =>
a -> Record rs -> Record ((s :-> a) : rs)
(&:) (Mutable (VectorFor a) (PrimState m) a
 -> Record (VectorMs m rs)
 -> Record
      ((s :-> Mutable (VectorFor a) (PrimState m) a) : VectorMs m rs))
-> m (Mutable (VectorFor a) (PrimState m) a)
-> m (Record (VectorMs m rs)
      -> Record
           ((s :-> Mutable (VectorFor a) (PrimState m) a) : VectorMs m rs))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m (Mutable (VectorFor a) (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
VGM.new Int
size m (Record (VectorMs m rs)
   -> Record
        ((s :-> Mutable (VectorFor a) (PrimState m) a) : VectorMs m rs))
-> m (Record (VectorMs m rs))
-> m (Record
        ((s :-> Mutable (VectorFor a) (PrimState m) a) : VectorMs m rs))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Proxy rs -> Int -> m (Record (VectorMs m rs))
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs -> Int -> m (Record (VectorMs m rs))
allocRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs) Int
size
  {-# INLINE allocRec #-}

  freezeRec :: proxy ((s :-> a) : rs)
-> Int
-> Record (VectorMs m ((s :-> a) : rs))
-> m (Record (Vectors ((s :-> a) : rs)))
freezeRec proxy ((s :-> a) : rs)
_ Int
n (Col x :& Rec ElField rs
xs) =
    VectorFor a a
-> Record (Vectors rs)
-> Record ((s :-> VectorFor a a) : Vectors rs)
forall (s :: Symbol) a (rs :: [(Symbol, *)]).
KnownSymbol s =>
a -> Record rs -> Record ((s :-> a) : rs)
(&:) (VectorFor a a
 -> Record (Vectors rs)
 -> Record ((s :-> VectorFor a a) : Vectors rs))
-> m (VectorFor a a)
-> m (Record (Vectors rs)
      -> Record ((s :-> VectorFor a a) : Vectors rs))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Mutable (VectorFor a) (PrimState m) a -> m (VectorFor a a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
VG.unsafeFreeze (Mutable (VectorFor a) (PrimState m) a -> m (VectorFor a a))
-> Mutable (VectorFor a) (PrimState m) a -> m (VectorFor a a)
forall a b. (a -> b) -> a -> b
$ Int
-> Int
-> Mutable (VectorFor a) (PrimState m) a
-> Mutable (VectorFor a) (PrimState m) a
forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
VGM.unsafeSlice Int
0 Int
n Mutable (VectorFor a) (PrimState m) a
x)
         m (Record (Vectors rs)
   -> Record ((s :-> VectorFor a a) : Vectors rs))
-> m (Record (Vectors rs))
-> m (Record ((s :-> VectorFor a a) : Vectors rs))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Proxy rs
-> Int -> Record (VectorMs m rs) -> m (Record (Vectors rs))
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs
-> Int -> Record (VectorMs m rs) -> m (Record (Vectors rs))
freezeRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs) Int
n Rec ElField rs
Record (VectorMs m rs)
xs
  freezeRec proxy ((s :-> a) : rs)
_ Int
_ Record (VectorMs m ((s :-> a) : rs))
_ = [Char] -> m (Record ((s :-> VectorFor a a) : Vectors rs))
forall a. HasCallStack => [Char] -> a
error [Char]
"Impossible case freezeRec"
  {-# INLINE freezeRec #-}

  growRec :: proxy ((s :-> a) : rs)
-> Record (VectorMs m ((s :-> a) : rs))
-> m (Record (VectorMs m ((s :-> a) : rs)))
growRec proxy ((s :-> a) : rs)
_ (Col x :& Rec ElField rs
xs) = Mutable (VectorFor a) (PrimState m) a
-> Rec ElField rs
-> Record ((s :-> Mutable (VectorFor a) (PrimState m) a) : rs)
forall (s :: Symbol) a (rs :: [(Symbol, *)]).
KnownSymbol s =>
a -> Record rs -> Record ((s :-> a) : rs)
(&:) (Mutable (VectorFor a) (PrimState m) a
 -> Rec ElField rs
 -> Record ((s :-> Mutable (VectorFor a) (PrimState m) a) : rs))
-> m (Mutable (VectorFor a) (PrimState m) a)
-> m (Rec ElField rs
      -> Record ((s :-> Mutable (VectorFor a) (PrimState m) a) : rs))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mutable (VectorFor a) (PrimState m) a
-> Int -> m (Mutable (VectorFor a) (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (v (PrimState m) a)
VGM.grow Mutable (VectorFor a) (PrimState m) a
x (Mutable (VectorFor a) (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
VGM.length Mutable (VectorFor a) (PrimState m) a
x)
                                              m (Rec ElField rs
   -> Record ((s :-> Mutable (VectorFor a) (PrimState m) a) : rs))
-> m (Rec ElField rs)
-> m (Record ((s :-> Mutable (VectorFor a) (PrimState m) a) : rs))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Proxy rs -> Record (VectorMs m rs) -> m (Record (VectorMs m rs))
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs -> Record (VectorMs m rs) -> m (Record (VectorMs m rs))
growRec (Proxy rs
forall k (t :: k). Proxy t
Proxy :: Proxy rs) Rec ElField rs
Record (VectorMs m rs)
xs
  growRec proxy ((s :-> a) : rs)
_ Record (VectorMs m ((s :-> a) : rs))
_ = [Char]
-> m (Record
        ((s :-> Mutable (VectorFor a) (PrimState m) a) : VectorMs m rs))
forall a. HasCallStack => [Char] -> a
error [Char]
"Impossible case growRec"
  {-# INLINE growRec #-}

  writeRec :: proxy ((s :-> a) : rs)
-> Int
-> Record (VectorMs m ((s :-> a) : rs))
-> Record ((s :-> a) : rs)
-> m ()
writeRec proxy ((s :-> a) : rs)
_ !Int
i !(Col v :& Rec ElField rs
vs) (Col x :& Rec ElField rs
xs) =
    Mutable (VectorFor a) (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite Mutable (VectorFor a) (PrimState m) a
v Int
i a
x m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Proxy rs -> Int -> Record (VectorMs m rs) -> Record rs -> m ()
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs -> Int -> Record (VectorMs m rs) -> Record rs -> m ()
writeRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs) Int
i Rec ElField rs
Record (VectorMs m rs)
vs Record rs
Rec ElField rs
xs
  writeRec proxy ((s :-> a) : rs)
_ Int
_ Record (VectorMs m ((s :-> a) : rs))
_ Record ((s :-> a) : rs)
_ = [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Impossible case writeRec"
  {-# INLINE writeRec #-}

  indexRec :: proxy ((s :-> a) : rs)
-> Int
-> Record (Vectors ((s :-> a) : rs))
-> Record ((s :-> a) : rs)
indexRec proxy ((s :-> a) : rs)
_ !Int
i !(Col x :& Rec ElField rs
xs) =
    VectorFor a a
x VectorFor a a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.! Int
i a -> Record rs -> Record ((s :-> a) : rs)
forall (s :: Symbol) a (rs :: [(Symbol, *)]).
KnownSymbol s =>
a -> Record rs -> Record ((s :-> a) : rs)
&: Proxy rs -> Int -> Record (Vectors rs) -> Record rs
forall (rs :: [(Symbol, *)]) (proxy :: [(Symbol, *)] -> *).
RecVec rs =>
proxy rs -> Int -> Record (Vectors rs) -> Record rs
indexRec (Proxy rs
forall k (t :: k). Proxy t
Proxy :: Proxy rs) Int
i Rec ElField rs
Record (Vectors rs)
xs
  indexRec proxy ((s :-> a) : rs)
_ Int
_ Record (Vectors ((s :-> a) : rs))
_ = [Char] -> Record ((s :-> a) : rs)
forall a. HasCallStack => [Char] -> a
error [Char]
"Impossible case indexRec"
  {-# INLINE indexRec #-}

  produceRec :: proxy ((s :-> a) : rs)
-> Record (Vectors ((s :-> a) : rs))
-> Rec ((->) Int :. ElField) ((s :-> a) : rs)
produceRec proxy ((s :-> a) : rs)
_ (Col v V.:& Rec ElField rs
vs) = (Int -> ElField (s :-> a)) -> Compose ((->) Int) ElField (s :-> a)
forall l k (f :: l -> *) (g :: k -> l) (x :: k).
f (g x) -> Compose f g x
Compose (a -> ElField (s :-> a)
forall (s :: Symbol) t. KnownSymbol s => t -> ElField '(s, t)
Field (a -> ElField (s :-> a)) -> (Int -> a) -> Int -> ElField (s :-> a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VectorFor a a
v VectorFor a a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.!)) Compose ((->) Int) ElField (s :-> a)
-> Rec ((->) Int :. ElField) rs
-> Rec ((->) Int :. ElField) ((s :-> a) : rs)
forall u (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& Proxy rs -> Record (Vectors rs) -> Rec ((->) Int :. ElField) rs
forall (rs :: [(Symbol, *)]) (proxy :: [(Symbol, *)] -> *).
RecVec rs =>
proxy rs -> Record (Vectors rs) -> Rec ((->) Int :. ElField) rs
produceRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs) Rec ElField rs
Record (Vectors rs)
vs
  produceRec proxy ((s :-> a) : rs)
_ Record (Vectors ((s :-> a) : rs))
_ = [Char] -> Rec ((->) Int :. ElField) ((s :-> a) : rs)
forall a. HasCallStack => [Char] -> a
error [Char]
"Impossible case produceRec"
  {-# INLINE produceRec #-}

-- | Stream a finite sequence of rows into an efficient in-memory
-- representation for further manipulation. Each column of the input
-- table will be stored optimally based on its type, making use of the
-- resulting generators a matter of indexing into a densely packed
-- representation. Returns the number of rows and a record of column
-- indexing functions. See 'toAoS' to convert the result to a 'Frame'
-- which provides an easier-to-use function that indexes into the
-- table in a row-major fashion.
inCoreSoA :: forall m rs. (PrimMonad m, RecVec rs)
          => P.Producer (Record rs) m ()
          -> m (Int, V.Rec (((->) Int) :. ElField) rs)
inCoreSoA :: Producer (Record rs) m () -> m (Int, Rec ((->) Int :. ElField) rs)
inCoreSoA Producer (Record rs) m ()
xs =
  do Record (VectorMs m rs)
mvs <- Proxy rs -> Int -> m (Record (VectorMs m rs))
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs -> Int -> m (Record (VectorMs m rs))
allocRec (Proxy rs
forall k (t :: k). Proxy t
Proxy :: Proxy rs) Int
initialCapacity
     let feed :: (Int, Int, Record (VectorMs m rs))
-> Record rs -> m (Int, Int, Record (VectorMs m rs))
feed (!Int
i, !Int
sz, !Record (VectorMs m rs)
mvs') Record rs
row
           | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sz = Proxy rs -> Record (VectorMs m rs) -> m (Record (VectorMs m rs))
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs -> Record (VectorMs m rs) -> m (Record (VectorMs m rs))
growRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs) Record (VectorMs m rs)
mvs'
                       m (Record (VectorMs m rs))
-> (Record (VectorMs m rs) -> m (Int, Int, Record (VectorMs m rs)))
-> m (Int, Int, Record (VectorMs m rs))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ((Int, Int, Record (VectorMs m rs))
 -> Record rs -> m (Int, Int, Record (VectorMs m rs)))
-> Record rs
-> (Int, Int, Record (VectorMs m rs))
-> m (Int, Int, Record (VectorMs m rs))
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Int, Int, Record (VectorMs m rs))
-> Record rs -> m (Int, Int, Record (VectorMs m rs))
feed Record rs
row ((Int, Int, Record (VectorMs m rs))
 -> m (Int, Int, Record (VectorMs m rs)))
-> (Record (VectorMs m rs) -> (Int, Int, Record (VectorMs m rs)))
-> Record (VectorMs m rs)
-> m (Int, Int, Record (VectorMs m rs))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
i, Int
szInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
2,)
           | Bool
otherwise = do Proxy rs -> Int -> Record (VectorMs m rs) -> Record rs -> m ()
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs -> Int -> Record (VectorMs m rs) -> Record rs -> m ()
writeRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs) Int
i Record (VectorMs m rs)
mvs' Record rs
row
                            (Int, Int, Record (VectorMs m rs))
-> m (Int, Int, Record (VectorMs m rs))
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1, Int
sz, Record (VectorMs m rs)
mvs')
         fin :: (Int, Int, Record (VectorMs m rs))
-> m (Int, Rec ((->) Int :. ElField) rs)
fin (Int
n,Int
_,Record (VectorMs m rs)
mvs') =
           do Record (Vectors rs)
vs <- Proxy rs
-> Int -> Record (VectorMs m rs) -> m (Record (Vectors rs))
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs
-> Int -> Record (VectorMs m rs) -> m (Record (Vectors rs))
freezeRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs) Int
n Record (VectorMs m rs)
mvs'
              (Int, Rec ((->) Int :. ElField) rs)
-> m (Int, Rec ((->) Int :. ElField) rs)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Int, Rec ((->) Int :. ElField) rs)
 -> m (Int, Rec ((->) Int :. ElField) rs))
-> (Rec ((->) Int :. ElField) rs
    -> (Int, Rec ((->) Int :. ElField) rs))
-> Rec ((->) Int :. ElField) rs
-> m (Int, Rec ((->) Int :. ElField) rs)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
n,) (Rec ((->) Int :. ElField) rs
 -> m (Int, Rec ((->) Int :. ElField) rs))
-> Rec ((->) Int :. ElField) rs
-> m (Int, Rec ((->) Int :. ElField) rs)
forall a b. (a -> b) -> a -> b
$ Proxy rs -> Record (Vectors rs) -> Rec ((->) Int :. ElField) rs
forall (rs :: [(Symbol, *)]) (proxy :: [(Symbol, *)] -> *).
RecVec rs =>
proxy rs -> Record (Vectors rs) -> Rec ((->) Int :. ElField) rs
produceRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs) Record (Vectors rs)
vs
     ((Int, Int, Record (VectorMs m rs))
 -> Record rs -> m (Int, Int, Record (VectorMs m rs)))
-> m (Int, Int, Record (VectorMs m rs))
-> ((Int, Int, Record (VectorMs m rs))
    -> m (Int, Rec ((->) Int :. ElField) rs))
-> Producer (Record rs) m ()
-> m (Int, Rec ((->) Int :. ElField) rs)
forall (m :: * -> *) x a b.
Monad m =>
(x -> a -> m x) -> m x -> (x -> m b) -> Producer a m () -> m b
P.foldM (Int, Int, Record (VectorMs m rs))
-> Record rs -> m (Int, Int, Record (VectorMs m rs))
feed ((Int, Int, Record (VectorMs m rs))
-> m (Int, Int, Record (VectorMs m rs))
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
0,Int
initialCapacity,Record (VectorMs m rs)
mvs)) (Int, Int, Record (VectorMs m rs))
-> m (Int, Rec ((->) Int :. ElField) rs)
fin Producer (Record rs) m ()
xs
{-# INLINE inCoreSoA #-}

-- | Stream a finite sequence of rows into an efficient in-memory
-- representation for further manipulation. Each column of the input
-- table will be stored optimally based on its type, making use of the
-- resulting generators a matter of indexing into a densely packed
-- representation. Returns a 'Frame' that provides a function to index
-- into the table.
inCoreAoS :: (PrimMonad m, RecVec rs)
          => P.Producer (Record rs) m () -> m (FrameRec rs)
inCoreAoS :: Producer (Record rs) m () -> m (FrameRec rs)
inCoreAoS = ((Int, Rec ((->) Int :. ElField) rs) -> FrameRec rs)
-> m (Int, Rec ((->) Int :. ElField) rs) -> m (FrameRec rs)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Int -> Rec ((->) Int :. ElField) rs -> FrameRec rs)
-> (Int, Rec ((->) Int :. ElField) rs) -> FrameRec rs
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Int -> Rec ((->) Int :. ElField) rs -> FrameRec rs
forall (rs :: [(Symbol, *)]).
Int -> Rec ((->) Int :. ElField) rs -> FrameRec rs
toAoS) (m (Int, Rec ((->) Int :. ElField) rs) -> m (FrameRec rs))
-> (Producer (Record rs) m ()
    -> m (Int, Rec ((->) Int :. ElField) rs))
-> Producer (Record rs) m ()
-> m (FrameRec rs)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Producer (Record rs) m () -> m (Int, Rec ((->) Int :. ElField) rs)
forall (m :: * -> *) (rs :: [(Symbol, *)]).
(PrimMonad m, RecVec rs) =>
Producer (Record rs) m () -> m (Int, Rec ((->) Int :. ElField) rs)
inCoreSoA

-- | Like 'inCoreAoS', but applies the provided function to the record
-- of columns before building the 'Frame'.
inCoreAoS' :: (PrimMonad m, RecVec rs)
           => (V.Rec ((->) Int :. ElField) rs -> V.Rec ((->) Int :. ElField) ss)
           -> P.Producer (Record rs) m () -> m (FrameRec ss)
inCoreAoS' :: (Rec ((->) Int :. ElField) rs -> Rec ((->) Int :. ElField) ss)
-> Producer (Record rs) m () -> m (FrameRec ss)
inCoreAoS' Rec ((->) Int :. ElField) rs -> Rec ((->) Int :. ElField) ss
f = ((Int, Rec ((->) Int :. ElField) rs) -> FrameRec ss)
-> m (Int, Rec ((->) Int :. ElField) rs) -> m (FrameRec ss)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Int -> Rec ((->) Int :. ElField) ss -> FrameRec ss)
-> (Int, Rec ((->) Int :. ElField) ss) -> FrameRec ss
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Int -> Rec ((->) Int :. ElField) ss -> FrameRec ss
forall (rs :: [(Symbol, *)]).
Int -> Rec ((->) Int :. ElField) rs -> FrameRec rs
toAoS ((Int, Rec ((->) Int :. ElField) ss) -> FrameRec ss)
-> ((Int, Rec ((->) Int :. ElField) rs)
    -> (Int, Rec ((->) Int :. ElField) ss))
-> (Int, Rec ((->) Int :. ElField) rs)
-> FrameRec ss
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Rec ((->) Int :. ElField) rs)
-> (Int, Rec ((->) Int :. ElField) ss)
aux) (m (Int, Rec ((->) Int :. ElField) rs) -> m (FrameRec ss))
-> (Producer (Record rs) m ()
    -> m (Int, Rec ((->) Int :. ElField) rs))
-> Producer (Record rs) m ()
-> m (FrameRec ss)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Producer (Record rs) m () -> m (Int, Rec ((->) Int :. ElField) rs)
forall (m :: * -> *) (rs :: [(Symbol, *)]).
(PrimMonad m, RecVec rs) =>
Producer (Record rs) m () -> m (Int, Rec ((->) Int :. ElField) rs)
inCoreSoA
  where aux :: (Int, Rec ((->) Int :. ElField) rs)
-> (Int, Rec ((->) Int :. ElField) ss)
aux (Int
x,Rec ((->) Int :. ElField) rs
y) = (Int
x, Rec ((->) Int :. ElField) rs -> Rec ((->) Int :. ElField) ss
f Rec ((->) Int :. ElField) rs
y)

-- | Convert a structure-of-arrays to an array-of-structures. This can
-- simplify usage of an in-memory representation.
toAoS :: Int -> V.Rec ((->) Int :. ElField) rs -> FrameRec rs
toAoS :: Int -> Rec ((->) Int :. ElField) rs -> FrameRec rs
toAoS Int
n = Int -> (Int -> Record rs) -> FrameRec rs
forall r. Int -> (Int -> r) -> Frame r
Frame Int
n ((Int -> Record rs) -> FrameRec rs)
-> (Rec ((->) Int :. ElField) rs -> Int -> Record rs)
-> Rec ((->) Int :. ElField) rs
-> FrameRec rs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (x :: (Symbol, *)).
 (:.) ((->) Int) ElField x -> Int -> ElField x)
-> Rec ((->) Int :. ElField) rs -> Int -> Record rs
forall u (h :: * -> *) (f :: u -> *) (g :: u -> *) (rs :: [u]).
Applicative h =>
(forall (x :: u). f x -> h (g x)) -> Rec f rs -> h (Rec g rs)
rtraverse forall l (f :: l -> *) k (g :: k -> l) (x :: k).
Compose f g x -> f (g x)
forall (x :: (Symbol, *)).
(:.) ((->) Int) ElField x -> Int -> ElField x
getCompose
{-# INLINE toAoS #-}

-- | Stream a finite sequence of rows into an efficient in-memory
-- representation for further manipulation. Each column of the input
-- table will be stored optimally based on its type, making use of the
-- resulting generator a matter of indexing into a densely packed
-- representation.
inCore :: forall m n rs. (PrimMonad m, RecVec rs, Monad n)
       => P.Producer (Record rs) m () -> m (P.Producer (Record rs) n ())
inCore :: Producer (Record rs) m () -> m (Producer (Record rs) n ())
inCore Producer (Record rs) m ()
xs =
  do Record (VectorMs m rs)
mvs <- Proxy rs -> Int -> m (Record (VectorMs m rs))
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs -> Int -> m (Record (VectorMs m rs))
allocRec (Proxy rs
forall k (t :: k). Proxy t
Proxy :: Proxy rs) Int
initialCapacity
     let feed :: (Int, Int, Record (VectorMs m rs))
-> Record rs -> m (Int, Int, Record (VectorMs m rs))
feed (!Int
i,!Int
sz,!Record (VectorMs m rs)
mvs') Record rs
row
              | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sz = Proxy rs -> Record (VectorMs m rs) -> m (Record (VectorMs m rs))
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs -> Record (VectorMs m rs) -> m (Record (VectorMs m rs))
growRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs) Record (VectorMs m rs)
mvs'
                          m (Record (VectorMs m rs))
-> (Record (VectorMs m rs) -> m (Int, Int, Record (VectorMs m rs)))
-> m (Int, Int, Record (VectorMs m rs))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ((Int, Int, Record (VectorMs m rs))
 -> Record rs -> m (Int, Int, Record (VectorMs m rs)))
-> Record rs
-> (Int, Int, Record (VectorMs m rs))
-> m (Int, Int, Record (VectorMs m rs))
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Int, Int, Record (VectorMs m rs))
-> Record rs -> m (Int, Int, Record (VectorMs m rs))
feed Record rs
row ((Int, Int, Record (VectorMs m rs))
 -> m (Int, Int, Record (VectorMs m rs)))
-> (Record (VectorMs m rs) -> (Int, Int, Record (VectorMs m rs)))
-> Record (VectorMs m rs)
-> m (Int, Int, Record (VectorMs m rs))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
i, Int
szInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
2,)
              | Bool
otherwise = do Proxy rs -> Int -> Record (VectorMs m rs) -> Record rs -> m ()
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs -> Int -> Record (VectorMs m rs) -> Record rs -> m ()
writeRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs) Int
i Record (VectorMs m rs)
mvs' Record rs
row
                               (Int, Int, Record (VectorMs m rs))
-> m (Int, Int, Record (VectorMs m rs))
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1, Int
sz, Record (VectorMs m rs)
mvs')
         fin :: (Int, Int, Record (VectorMs m rs)) -> m (Producer (Record rs) n ())
fin (Int
n,Int
_,Record (VectorMs m rs)
mvs') =
           do Record (Vectors rs)
vs <- Proxy rs
-> Int -> Record (VectorMs m rs) -> m (Record (Vectors rs))
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs
-> Int -> Record (VectorMs m rs) -> m (Record (Vectors rs))
freezeRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs) Int
n Record (VectorMs m rs)
mvs'
              let spool :: Int -> Producer (Record rs) n ()
spool !Int
i
                    | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = () -> Producer (Record rs) n ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                    | Bool
otherwise = Record rs -> Producer (Record rs) n ()
forall (m :: * -> *) a x' x. Functor m => a -> Proxy x' x () a m ()
P.yield (Proxy rs -> Int -> Record (Vectors rs) -> Record rs
forall (rs :: [(Symbol, *)]) (proxy :: [(Symbol, *)] -> *).
RecVec rs =>
proxy rs -> Int -> Record (Vectors rs) -> Record rs
indexRec Proxy rs
forall k (t :: k). Proxy t
Proxy Int
i Record (Vectors rs)
vs) Producer (Record rs) n ()
-> Producer (Record rs) n () -> Producer (Record rs) n ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Producer (Record rs) n ()
spool (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
              Producer (Record rs) n () -> m (Producer (Record rs) n ())
forall (m :: * -> *) a. Monad m => a -> m a
return (Producer (Record rs) n () -> m (Producer (Record rs) n ()))
-> Producer (Record rs) n () -> m (Producer (Record rs) n ())
forall a b. (a -> b) -> a -> b
$ Int -> Producer (Record rs) n ()
spool Int
0
     ((Int, Int, Record (VectorMs m rs))
 -> Record rs -> m (Int, Int, Record (VectorMs m rs)))
-> m (Int, Int, Record (VectorMs m rs))
-> ((Int, Int, Record (VectorMs m rs))
    -> m (Producer (Record rs) n ()))
-> Producer (Record rs) m ()
-> m (Producer (Record rs) n ())
forall (m :: * -> *) x a b.
Monad m =>
(x -> a -> m x) -> m x -> (x -> m b) -> Producer a m () -> m b
P.foldM (Int, Int, Record (VectorMs m rs))
-> Record rs -> m (Int, Int, Record (VectorMs m rs))
feed ((Int, Int, Record (VectorMs m rs))
-> m (Int, Int, Record (VectorMs m rs))
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
0,Int
initialCapacity,Record (VectorMs m rs)
mvs)) (Int, Int, Record (VectorMs m rs)) -> m (Producer (Record rs) n ())
fin Producer (Record rs) m ()
xs
{-# INLINE inCore #-}

-- | Build a 'Frame' from a collection of 'Record's using efficient
-- column-based storage.
toFrame :: (P.Foldable f, RecVec rs) => f (Record rs) -> Frame (Record rs)
toFrame :: f (Record rs) -> Frame (Record rs)
toFrame f (Record rs)
xs = (forall s. ST s (Frame (Record rs))) -> Frame (Record rs)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Frame (Record rs))) -> Frame (Record rs))
-> (forall s. ST s (Frame (Record rs))) -> Frame (Record rs)
forall a b. (a -> b) -> a -> b
$ Producer (Record rs) (ST s) () -> ST s (Frame (Record rs))
forall (m :: * -> *) (rs :: [(Symbol, *)]).
(PrimMonad m, RecVec rs) =>
Producer (Record rs) m () -> m (FrameRec rs)
inCoreAoS (f (Record rs) -> Producer (Record rs) (ST s) ()
forall (m :: * -> *) (f :: * -> *) a x' x.
(Functor m, Foldable f) =>
f a -> Proxy x' x () a m ()
P.each f (Record rs)
xs)
{-# INLINE toFrame #-}

-- | Keep only those rows of a 'FrameRec' that satisfy a predicate.
filterFrame :: RecVec rs => (Record rs -> Bool) -> FrameRec rs -> FrameRec rs
filterFrame :: (Record rs -> Bool) -> FrameRec rs -> FrameRec rs
filterFrame Record rs -> Bool
p FrameRec rs
f = (forall s. ST s (FrameRec rs)) -> FrameRec rs
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (FrameRec rs)) -> FrameRec rs)
-> (forall s. ST s (FrameRec rs)) -> FrameRec rs
forall a b. (a -> b) -> a -> b
$ Producer (Record rs) (ST s) () -> ST s (FrameRec rs)
forall (m :: * -> *) (rs :: [(Symbol, *)]).
(PrimMonad m, RecVec rs) =>
Producer (Record rs) m () -> m (FrameRec rs)
inCoreAoS (Producer (Record rs) (ST s) () -> ST s (FrameRec rs))
-> Producer (Record rs) (ST s) () -> ST s (FrameRec rs)
forall a b. (a -> b) -> a -> b
$ FrameRec rs -> Producer (Record rs) (ST s) ()
forall (m :: * -> *) (f :: * -> *) a x' x.
(Functor m, Foldable f) =>
f a -> Proxy x' x () a m ()
P.each FrameRec rs
f Producer (Record rs) (ST s) ()
-> Proxy () (Record rs) () (Record rs) (ST s) ()
-> Producer (Record rs) (ST s) ()
forall (m :: * -> *) a' a b r c' c.
Functor m =>
Proxy a' a () b m r -> Proxy () b c' c m r -> Proxy a' a c' c m r
P.>-> (Record rs -> Bool)
-> Proxy () (Record rs) () (Record rs) (ST s) ()
forall (m :: * -> *) a r. Functor m => (a -> Bool) -> Pipe a a m r
P.filter Record rs -> Bool
p
{-# INLINE filterFrame #-}

-- | Process a stream of 'Record's into a stream of 'Frame's that each
-- contains no more than the given number of records.
produceFrameChunks :: forall rs m. (RecVec rs, PrimMonad m)
                   => Int
                   -> P.Producer (Record rs) m ()
                   -> P.Producer (FrameRec rs) m ()
produceFrameChunks :: Int -> Producer (Record rs) m () -> Producer (FrameRec rs) m ()
produceFrameChunks Int
chunkSize = Producer (Record rs) m () -> Producer (FrameRec rs) m ()
go
  where go :: Producer (Record rs) m () -> Producer (FrameRec rs) m ()
go Producer (Record rs) m ()
src = do Record (VectorMs m rs)
mutVecs <- m (Record (VectorMs m rs))
-> Proxy X () () (FrameRec rs) m (Record (VectorMs m rs))
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
P.lift (Proxy rs -> Int -> m (Record (VectorMs m rs))
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs -> Int -> m (Record (VectorMs m rs))
allocRec (Proxy rs
forall k (t :: k). Proxy t
Proxy :: Proxy rs) Int
chunkSize)
                    Producer (Record rs) m ()
-> Record (VectorMs m rs) -> Int -> Producer (FrameRec rs) m ()
goChunk Producer (Record rs) m ()
src Record (VectorMs m rs)
mutVecs Int
0
        goChunk :: Producer (Record rs) m ()
-> Record (VectorMs m rs) -> Int -> Producer (FrameRec rs) m ()
goChunk Producer (Record rs) m ()
src Record (VectorMs m rs)
mutVecs !Int
i
          | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
chunkSize =
              do FrameRec rs
chunk <- m (FrameRec rs) -> Proxy X () () (FrameRec rs) m (FrameRec rs)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
P.lift (Int -> Record (VectorMs m rs) -> m (FrameRec rs)
freezeFrame Int
i Record (VectorMs m rs)
mutVecs)
                 FrameRec rs -> Producer (FrameRec rs) m ()
forall (m :: * -> *) a x' x. Functor m => a -> Proxy x' x () a m ()
P.yield FrameRec rs
chunk
                 Producer (Record rs) m () -> Producer (FrameRec rs) m ()
go Producer (Record rs) m ()
src
          | Bool
otherwise =
            do Either () (Record rs, Producer (Record rs) m ())
maybeRow <- m (Either () (Record rs, Producer (Record rs) m ()))
-> Proxy
     X
     ()
     ()
     (FrameRec rs)
     m
     (Either () (Record rs, Producer (Record rs) m ()))
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
P.lift (Producer (Record rs) m ()
-> m (Either () (Record rs, Producer (Record rs) m ()))
forall (m :: * -> *) a r.
Monad m =>
Producer a m r -> m (Either r (a, Producer a m r))
P.next Producer (Record rs) m ()
src)
               case Either () (Record rs, Producer (Record rs) m ())
maybeRow of
                 Left ()
_ -> do
                   m (FrameRec rs) -> Proxy X () () (FrameRec rs) m (FrameRec rs)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
P.lift (Int -> Record (VectorMs m rs) -> m (FrameRec rs)
freezeFrame Int
i Record (VectorMs m rs)
mutVecs) Proxy X () () (FrameRec rs) m (FrameRec rs)
-> (FrameRec rs -> Producer (FrameRec rs) m ())
-> Producer (FrameRec rs) m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= FrameRec rs -> Producer (FrameRec rs) m ()
forall (m :: * -> *) a x' x. Functor m => a -> Proxy x' x () a m ()
P.yield
                 Right (Record rs
r,Producer (Record rs) m ()
src') -> do
                   m () -> Producer (FrameRec rs) m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
P.lift (Proxy rs -> Int -> Record (VectorMs m rs) -> Record rs -> m ()
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs -> Int -> Record (VectorMs m rs) -> Record rs -> m ()
writeRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs) Int
i Record (VectorMs m rs)
mutVecs Record rs
r)
                   Producer (Record rs) m ()
-> Record (VectorMs m rs) -> Int -> Producer (FrameRec rs) m ()
goChunk Producer (Record rs) m ()
src' Record (VectorMs m rs)
mutVecs (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
        freezeFrame :: Int -> Record (VectorMs m rs) -> m (FrameRec rs)
        freezeFrame :: Int -> Record (VectorMs m rs) -> m (FrameRec rs)
freezeFrame Int
n =
          (Record (Vectors rs) -> FrameRec rs)
-> m (Record (Vectors rs)) -> m (FrameRec rs)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> Rec ((->) Int :. ElField) rs -> FrameRec rs
forall (rs :: [(Symbol, *)]).
Int -> Rec ((->) Int :. ElField) rs -> FrameRec rs
toAoS Int
n (Rec ((->) Int :. ElField) rs -> FrameRec rs)
-> (Record (Vectors rs) -> Rec ((->) Int :. ElField) rs)
-> Record (Vectors rs)
-> FrameRec rs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proxy rs -> Record (Vectors rs) -> Rec ((->) Int :. ElField) rs
forall (rs :: [(Symbol, *)]) (proxy :: [(Symbol, *)] -> *).
RecVec rs =>
proxy rs -> Record (Vectors rs) -> Rec ((->) Int :. ElField) rs
produceRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs))
          (m (Record (Vectors rs)) -> m (FrameRec rs))
-> (Record (VectorMs m rs) -> m (Record (Vectors rs)))
-> Record (VectorMs m rs)
-> m (FrameRec rs)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proxy rs
-> Int -> Record (VectorMs m rs) -> m (Record (Vectors rs))
forall (rs :: [(Symbol, *)]) (m :: * -> *)
       (proxy :: [(Symbol, *)] -> *).
(RecVec rs, PrimMonad m) =>
proxy rs
-> Int -> Record (VectorMs m rs) -> m (Record (Vectors rs))
freezeRec (Proxy rs
forall k (t :: k). Proxy t
Proxy::Proxy rs) Int
n
{-# INLINABLE produceFrameChunks #-}

-- | Split a 'Frame' into chunks of no more than the given number of
-- records. The underlying memory is shared with the original 'Frame'.
frameChunks :: Int -> FrameRec rs -> [FrameRec rs]
frameChunks :: Int -> FrameRec rs -> [FrameRec rs]
frameChunks Int
chunkSize FrameRec rs
whole = (Int -> FrameRec rs) -> [Int] -> [FrameRec rs]
forall a b. (a -> b) -> [a] -> [b]
map Int -> FrameRec rs
aux [ Int
0, Int
chunkSize .. FrameRec rs -> Int
forall r. Frame r -> Int
frameLength FrameRec rs
whole Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 ]
  where aux :: Int -> FrameRec rs
aux Int
i = Int -> (Int -> Record rs) -> FrameRec rs
forall r. Int -> (Int -> r) -> Frame r
Frame (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (FrameRec rs -> Int
forall r. Frame r -> Int
frameLength FrameRec rs
whole Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) Int
chunkSize)
                      (FrameRec rs -> Int -> Record rs
forall r. Frame r -> Int -> r
frameRow FrameRec rs
whole (Int -> Record rs) -> (Int -> Int) -> Int -> Record rs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i))
{-# INLINABLE frameChunks #-}