{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}

{- | Mutable, growable column builders for high-throughput ingest. No
per-append @IORef@ traffic: hot counters live in an unboxed vector, payloads
double on demand, and validity is only materialized once a null is seen.
-}
module DataFrame.Internal.ColumnBuilder (
    ColumnBuilder (..),
    NumBuilder,
    IntBuilder,
    DoubleBuilder,
    TextBuilder,
    TextChunk (..),
    newIntBuilder,
    newDoubleBuilder,
    newNumBuilder,
    newTextBuilder,
    appendInt,
    appendDouble,
    appendNum,
    appendText,
    appendTextSlice,
    appendTextSliceFromPtr,
    freezeTextChunk,
    mergeColumns,
    mergeTextChunks,
) where

import qualified Data.Text as T
import qualified Data.Text.Array as A
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM

import Control.Monad (when)
import Control.Monad.ST (ST)
import Data.Bits (shiftR)
import Data.STRef
import Data.Text.Internal (Text (..))
import Data.Word (Word8)
import DataFrame.Internal.Column hiding (mergeColumns)
import DataFrame.Internal.ColumnMerge (
    TextChunk (..),
    mergeColumns,
    mergeTextChunks,
    packValidity,
 )
import Foreign.Ptr (Ptr)

{- | Operations shared by all column builders. A builder must not be used
again after 'freezeBuilder' (its storage is frozen in place, not copied).
-}
class ColumnBuilder b where
    -- | Append a null row (sentinel payload + invalid bit).
    appendNull :: b s -> ST s ()

    -- | Rows appended so far.
    builderLength :: b s -> ST s Int

    -- | Freeze into a fully-forced 'Column'; bitmap only when a null was seen.
    freezeBuilder :: b s -> ST s Column

-- Counter slots shared by the builders: rows, any-null flag, text bytes used.
cRows, cAnyNull, cBytes :: Int
cRows :: Int
cRows = Int
0
cAnyNull :: Int
cAnyNull = Int
1
cBytes :: Int
cBytes = Int
2

{- | Builder for unboxed numeric payloads ('Int', 'Double', ...). 'nbNull'
is the sentinel written into null slots (protected by the bitmap).
-}
data NumBuilder a s = NumBuilder
    { forall a s. NumBuilder a s -> a
nbNull :: !a
    , forall a s. NumBuilder a s -> MVector s Int
nbCounters :: !(VUM.MVector s Int)
    , forall a s. NumBuilder a s -> STRef s (NumArrays a s)
nbArrays :: !(STRef s (NumArrays a s))
    }

data NumArrays a s = NumArrays
    { forall a s. NumArrays a s -> MVector s a
naData :: !(VUM.MVector s a)
    , forall a s. NumArrays a s -> MVector s Word8
naValid :: !(VUM.MVector s Word8)
    }

type IntBuilder = NumBuilder Int

type DoubleBuilder = NumBuilder Double

-- | New numeric builder with a row-capacity hint and a null sentinel.
newNumBuilder :: (VU.Unbox a) => a -> Int -> ST s (NumBuilder a s)
newNumBuilder :: forall a s. Unbox a => a -> Int -> ST s (NumBuilder a s)
newNumBuilder a
nullValue Int
hint = do
    let cap :: Int
cap = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
16 Int
hint
    MVector s Int
counters <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
2 Int
0
    MVector s a
dat <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
cap
    MVector s Word8
val <- Int -> ST s (MVector (PrimState (ST s)) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
cap
    a -> MVector s Int -> STRef s (NumArrays a s) -> NumBuilder a s
forall a s.
a -> MVector s Int -> STRef s (NumArrays a s) -> NumBuilder a s
NumBuilder a
nullValue MVector s Int
counters (STRef s (NumArrays a s) -> NumBuilder a s)
-> ST s (STRef s (NumArrays a s)) -> ST s (NumBuilder a s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NumArrays a s -> ST s (STRef s (NumArrays a s))
forall a s. a -> ST s (STRef s a)
newSTRef (MVector s a -> MVector s Word8 -> NumArrays a s
forall a s. MVector s a -> MVector s Word8 -> NumArrays a s
NumArrays MVector s a
dat MVector s Word8
val)

newIntBuilder :: Int -> ST s (IntBuilder s)
newIntBuilder :: forall s. Int -> ST s (IntBuilder s)
newIntBuilder = Int -> Int -> ST s (NumBuilder Int s)
forall a s. Unbox a => a -> Int -> ST s (NumBuilder a s)
newNumBuilder Int
0

newDoubleBuilder :: Int -> ST s (DoubleBuilder s)
newDoubleBuilder :: forall s. Int -> ST s (DoubleBuilder s)
newDoubleBuilder = Double -> Int -> ST s (NumBuilder Double s)
forall a s. Unbox a => a -> Int -> ST s (NumBuilder a s)
newNumBuilder Double
0

appendNum :: (VU.Unbox a) => NumBuilder a s -> a -> ST s ()
appendNum :: forall a s. Unbox a => NumBuilder a s -> a -> ST s ()
appendNum NumBuilder a s
b !a
x = do
    Int
n <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (NumBuilder a s -> MVector s Int
forall a s. NumBuilder a s -> MVector s Int
nbCounters NumBuilder a s
b) Int
cRows
    Int
anyNull <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (NumBuilder a s -> MVector s Int
forall a s. NumBuilder a s -> MVector s Int
nbCounters NumBuilder a s
b) Int
cAnyNull
    NumArrays MVector s a
dat MVector s Word8
val <- NumBuilder a s -> Int -> ST s (NumArrays a s)
forall a s.
Unbox a =>
NumBuilder a s -> Int -> ST s (NumArrays a s)
reserveNum NumBuilder a s
b Int
n
    MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s a
MVector (PrimState (ST s)) a
dat Int
n a
x
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
anyNull Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState (ST s)) Word8 -> Int -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Word8
MVector (PrimState (ST s)) Word8
val Int
n Word8
1
    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 (NumBuilder a s -> MVector s Int
forall a s. NumBuilder a s -> MVector s Int
nbCounters NumBuilder a s
b) Int
cRows (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE appendNum #-}

appendInt :: IntBuilder s -> Int -> ST s ()
appendInt :: forall s. IntBuilder s -> Int -> ST s ()
appendInt = NumBuilder Int s -> Int -> ST s ()
forall a s. Unbox a => NumBuilder a s -> a -> ST s ()
appendNum
{-# INLINE appendInt #-}

appendDouble :: DoubleBuilder s -> Double -> ST s ()
appendDouble :: forall s. DoubleBuilder s -> Double -> ST s ()
appendDouble = NumBuilder Double s -> Double -> ST s ()
forall a s. Unbox a => NumBuilder a s -> a -> ST s ()
appendNum
{-# INLINE appendDouble #-}

-- Fetch the arrays, growing (doubling) first if row @n@ would not fit.
reserveNum :: (VU.Unbox a) => NumBuilder a s -> Int -> ST s (NumArrays a s)
reserveNum :: forall a s.
Unbox a =>
NumBuilder a s -> Int -> ST s (NumArrays a s)
reserveNum NumBuilder a s
b Int
n = do
    NumArrays a s
arrs <- STRef s (NumArrays a s) -> ST s (NumArrays a s)
forall s a. STRef s a -> ST s a
readSTRef (NumBuilder a s -> STRef s (NumArrays a s)
forall a s. NumBuilder a s -> STRef s (NumArrays a s)
nbArrays NumBuilder a s
b)
    if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< MVector s a -> Int
forall a s. Unbox a => MVector s a -> Int
VUM.length (NumArrays a s -> MVector s a
forall a s. NumArrays a s -> MVector s a
naData NumArrays a s
arrs) then NumArrays a s -> ST s (NumArrays a s)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NumArrays a s
arrs else NumBuilder a s -> NumArrays a s -> ST s (NumArrays a s)
forall a s.
Unbox a =>
NumBuilder a s -> NumArrays a s -> ST s (NumArrays a s)
growNum NumBuilder a s
b NumArrays a s
arrs
{-# INLINE reserveNum #-}

growNum ::
    (VU.Unbox a) => NumBuilder a s -> NumArrays a s -> ST s (NumArrays a s)
growNum :: forall a s.
Unbox a =>
NumBuilder a s -> NumArrays a s -> ST s (NumArrays a s)
growNum NumBuilder a s
b (NumArrays MVector s a
dat MVector s Word8
val) = do
    let cap :: Int
cap = MVector s a -> Int
forall a s. Unbox a => MVector s a -> Int
VUM.length MVector s a
dat
    MVector s a
dat' <- MVector (PrimState (ST s)) a
-> Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
VUM.unsafeGrow MVector s a
MVector (PrimState (ST s)) a
dat Int
cap
    MVector s Word8
val' <- MVector (PrimState (ST s)) Word8
-> Int -> ST s (MVector (PrimState (ST s)) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
VUM.unsafeGrow MVector s Word8
MVector (PrimState (ST s)) Word8
val Int
cap
    let arrs :: NumArrays a s
arrs = MVector s a -> MVector s Word8 -> NumArrays a s
forall a s. MVector s a -> MVector s Word8 -> NumArrays a s
NumArrays MVector s a
dat' MVector s Word8
val'
    STRef s (NumArrays a s) -> NumArrays a s -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef (NumBuilder a s -> STRef s (NumArrays a s)
forall a s. NumBuilder a s -> STRef s (NumArrays a s)
nbArrays NumBuilder a s
b) NumArrays a s
arrs
    NumArrays a s -> ST s (NumArrays a s)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NumArrays a s
arrs

instance (Columnable a, VU.Unbox a) => ColumnBuilder (NumBuilder a) where
    appendNull :: forall s. NumBuilder a s -> ST s ()
appendNull NumBuilder a s
b = do
        Int
n <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (NumBuilder a s -> MVector s Int
forall a s. NumBuilder a s -> MVector s Int
nbCounters NumBuilder a s
b) Int
cRows
        Int
anyNull <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (NumBuilder a s -> MVector s Int
forall a s. NumBuilder a s -> MVector s Int
nbCounters NumBuilder a s
b) Int
cAnyNull
        NumArrays MVector s a
dat MVector s Word8
val <- NumBuilder a s -> Int -> ST s (NumArrays a s)
forall a s.
Unbox a =>
NumBuilder a s -> Int -> ST s (NumArrays a s)
reserveNum NumBuilder a s
b Int
n
        MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s a
MVector (PrimState (ST s)) a
dat Int
n (NumBuilder a s -> a
forall a s. NumBuilder a s -> a
nbNull NumBuilder a s
b)
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
anyNull Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            MVector (PrimState (ST s)) Word8 -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> a -> m ()
VUM.set (Int -> Int -> MVector s Word8 -> MVector s Word8
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
VUM.slice Int
0 Int
n MVector s Word8
val) Word8
1
            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 (NumBuilder a s -> MVector s Int
forall a s. NumBuilder a s -> MVector s Int
nbCounters NumBuilder a s
b) Int
cAnyNull Int
1
        MVector (PrimState (ST s)) Word8 -> Int -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Word8
MVector (PrimState (ST s)) Word8
val Int
n Word8
0
        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 (NumBuilder a s -> MVector s Int
forall a s. NumBuilder a s -> MVector s Int
nbCounters NumBuilder a s
b) Int
cRows (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    {-# INLINE appendNull #-}

    builderLength :: forall s. NumBuilder a s -> ST s Int
builderLength NumBuilder a s
b = MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (NumBuilder a s -> MVector s Int
forall a s. NumBuilder a s -> MVector s Int
nbCounters NumBuilder a s
b) Int
cRows

    freezeBuilder :: forall s. NumBuilder a s -> ST s Column
freezeBuilder NumBuilder a s
b = do
        Int
n <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (NumBuilder a s -> MVector s Int
forall a s. NumBuilder a s -> MVector s Int
nbCounters NumBuilder a s
b) Int
cRows
        Int
anyNull <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (NumBuilder a s -> MVector s Int
forall a s. NumBuilder a s -> MVector s Int
nbCounters NumBuilder a s
b) Int
cAnyNull
        NumArrays MVector s a
dat MVector s Word8
val <- STRef s (NumArrays a s) -> ST s (NumArrays a s)
forall s a. STRef s a -> ST s a
readSTRef (NumBuilder a s -> STRef s (NumArrays a s)
forall a s. NumBuilder a s -> STRef s (NumArrays a s)
nbArrays NumBuilder a s
b)
        !Vector a
vs <- Int -> MVector s a -> ST s (Vector a)
forall a s. Unbox a => Int -> MVector s a -> ST s (Vector a)
freezeTrimmed Int
n MVector s a
dat
        if Int
anyNull Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
            then do
                !Bitmap
bm <- Int -> MVector s Word8 -> ST s Bitmap
forall s. Int -> MVector s Word8 -> ST s Bitmap
packValidity Int
n MVector s Word8
val
                Column -> ST s Column
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Column -> ST s Column) -> Column -> ST s Column
forall a b. (a -> b) -> a -> b
$! Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn (Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just Bitmap
bm) Vector a
vs
            else Column -> ST s Column
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Column -> ST s Column) -> Column -> ST s Column
forall a b. (a -> b) -> a -> b
$! Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
forall a. Maybe a
Nothing Vector a
vs

-- Zero-copy freeze; copies to exact size when slack exceeds a quarter of n.
freezeTrimmed :: (VU.Unbox a) => Int -> VUM.MVector s a -> ST s (VU.Vector a)
freezeTrimmed :: forall a s. Unbox a => Int -> MVector s a -> ST s (Vector a)
freezeTrimmed Int
n MVector s a
mv
    | MVector s a -> Int
forall a s. Unbox a => MVector s a -> Int
VUM.length MVector s a
mv Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
2 = MVector (PrimState (ST s)) a -> ST s (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze (Int -> Int -> MVector s a -> MVector s a
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
VUM.slice Int
0 Int
n MVector s a
mv)
    | Bool
otherwise = MVector (PrimState (ST s)) a -> ST s (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.freeze (Int -> Int -> MVector s a -> MVector s a
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
VUM.slice Int
0 Int
n MVector s a
mv)

{- | Builder for 'Text' columns. All field bytes go into one exponentially
grown byte array; rows are recorded as offsets, so an append is a memcpy
and freezing slices 'Text' values off the shared array without copying.
-}
data TextBuilder s = TextBuilder
    { forall s. TextBuilder s -> MVector s Int
tbCounters :: !(VUM.MVector s Int)
    , forall s. TextBuilder s -> STRef s (TextArrays s)
tbArrays :: !(STRef s (TextArrays s))
    }

data TextArrays s = TextArrays
    { forall s. TextArrays s -> MArray s
taBytes :: !(A.MArray s)
    , forall s. TextArrays s -> Int
taByteCap :: !Int
    , forall s. TextArrays s -> MVector s Int
taOffsets :: !(VUM.MVector s Int)
    -- ^ Row @i@ spans bytes @[offsets!i, offsets!(i+1))@.
    , forall s. TextArrays s -> MVector s Word8
taValid :: !(VUM.MVector s Word8)
    }

-- | New text builder with row-count and total-byte capacity hints.
newTextBuilder :: Int -> Int -> ST s (TextBuilder s)
newTextBuilder :: forall s. Int -> Int -> ST s (TextBuilder s)
newTextBuilder Int
rowHint Int
byteHint = do
    let rcap :: Int
rcap = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
16 Int
rowHint
        bcap :: Int
bcap = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
64 Int
byteHint
    MVector s Int
counters <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
3 Int
0
    MArray s
bytes <- Int -> ST s (MArray s)
forall s. Int -> ST s (MArray s)
A.new Int
bcap
    MVector s Int
offsets <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int
rcap Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    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
offsets Int
0 Int
0
    MVector s Word8
val <- Int -> ST s (MVector (PrimState (ST s)) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
rcap
    MVector s Int -> STRef s (TextArrays s) -> TextBuilder s
forall s. MVector s Int -> STRef s (TextArrays s) -> TextBuilder s
TextBuilder MVector s Int
counters (STRef s (TextArrays s) -> TextBuilder s)
-> ST s (STRef s (TextArrays s)) -> ST s (TextBuilder s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TextArrays s -> ST s (STRef s (TextArrays s))
forall a s. a -> ST s (STRef s a)
newSTRef (MArray s -> Int -> MVector s Int -> MVector s Word8 -> TextArrays s
forall s.
MArray s -> Int -> MVector s Int -> MVector s Word8 -> TextArrays s
TextArrays MArray s
bytes Int
bcap MVector s Int
offsets MVector s Word8
val)

-- | Append @len@ raw bytes at @off@ in @src@ as one field (one memcpy).
appendTextSlice :: TextBuilder s -> A.Array -> Int -> Int -> ST s ()
appendTextSlice :: forall s. TextBuilder s -> Array -> Int -> Int -> ST s ()
appendTextSlice TextBuilder s
b Array
src Int
off Int
len = do
    (Int
n, Int
pos, TextArrays s
arrs) <- TextBuilder s -> Int -> ST s (Int, Int, TextArrays s)
forall s. TextBuilder s -> Int -> ST s (Int, Int, TextArrays s)
reserveText TextBuilder s
b Int
len
    Int -> MArray s -> Int -> Array -> Int -> ST s ()
forall s. Int -> MArray s -> Int -> Array -> Int -> ST s ()
A.copyI Int
len (TextArrays s -> MArray s
forall s. TextArrays s -> MArray s
taBytes TextArrays s
arrs) Int
pos Array
src Int
off
    TextBuilder s -> TextArrays s -> Int -> Int -> ST s ()
forall s. TextBuilder s -> TextArrays s -> Int -> Int -> ST s ()
finishTextAppend TextBuilder s
b TextArrays s
arrs Int
n (Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len)
{-# INLINE appendTextSlice #-}

-- | 'appendTextSlice' from foreign memory (e.g. an mmapped file buffer).
appendTextSliceFromPtr :: TextBuilder s -> Ptr Word8 -> Int -> ST s ()
appendTextSliceFromPtr :: forall s. TextBuilder s -> Ptr Word8 -> Int -> ST s ()
appendTextSliceFromPtr TextBuilder s
b Ptr Word8
ptr Int
len = do
    (Int
n, Int
pos, TextArrays s
arrs) <- TextBuilder s -> Int -> ST s (Int, Int, TextArrays s)
forall s. TextBuilder s -> Int -> ST s (Int, Int, TextArrays s)
reserveText TextBuilder s
b Int
len
    MArray s -> Int -> Ptr Word8 -> Int -> ST s ()
forall s. MArray s -> Int -> Ptr Word8 -> Int -> ST s ()
A.copyFromPointer (TextArrays s -> MArray s
forall s. TextArrays s -> MArray s
taBytes TextArrays s
arrs) Int
pos Ptr Word8
ptr Int
len
    TextBuilder s -> TextArrays s -> Int -> Int -> ST s ()
forall s. TextBuilder s -> TextArrays s -> Int -> Int -> ST s ()
finishTextAppend TextBuilder s
b TextArrays s
arrs Int
n (Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len)
{-# INLINE appendTextSliceFromPtr #-}

-- | Append an already-decoded 'Text' (its bytes are UTF-8 already).
appendText :: TextBuilder s -> T.Text -> ST s ()
appendText :: forall s. TextBuilder s -> Text -> ST s ()
appendText TextBuilder s
b (Text Array
src Int
off Int
len) = TextBuilder s -> Array -> Int -> Int -> ST s ()
forall s. TextBuilder s -> Array -> Int -> Int -> ST s ()
appendTextSlice TextBuilder s
b Array
src Int
off Int
len
{-# INLINE appendText #-}

finishTextAppend :: TextBuilder s -> TextArrays s -> Int -> Int -> ST s ()
finishTextAppend :: forall s. TextBuilder s -> TextArrays s -> Int -> Int -> ST s ()
finishTextAppend TextBuilder s
b TextArrays s
arrs Int
n Int
endPos = do
    Int
anyNull <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (TextBuilder s -> MVector s Int
forall s. TextBuilder s -> MVector s Int
tbCounters TextBuilder s
b) Int
cAnyNull
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
anyNull Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState (ST s)) Word8 -> Int -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite (TextArrays s -> MVector s Word8
forall s. TextArrays s -> MVector s Word8
taValid TextArrays s
arrs) Int
n Word8
1
    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 (TextArrays s -> MVector s Int
forall s. TextArrays s -> MVector s Int
taOffsets TextArrays s
arrs) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
endPos
    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 (TextBuilder s -> MVector s Int
forall s. TextBuilder s -> MVector s Int
tbCounters TextBuilder s
b) Int
cRows (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    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 (TextBuilder s -> MVector s Int
forall s. TextBuilder s -> MVector s Int
tbCounters TextBuilder s
b) Int
cBytes Int
endPos
{-# INLINE finishTextAppend #-}

reserveText :: TextBuilder s -> Int -> ST s (Int, Int, TextArrays s)
reserveText :: forall s. TextBuilder s -> Int -> ST s (Int, Int, TextArrays s)
reserveText TextBuilder s
b Int
extra = do
    Int
n <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (TextBuilder s -> MVector s Int
forall s. TextBuilder s -> MVector s Int
tbCounters TextBuilder s
b) Int
cRows
    Int
pos <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (TextBuilder s -> MVector s Int
forall s. TextBuilder s -> MVector s Int
tbCounters TextBuilder s
b) Int
cBytes
    TextArrays s
arrs <- STRef s (TextArrays s) -> ST s (TextArrays s)
forall s a. STRef s a -> ST s a
readSTRef (TextBuilder s -> STRef s (TextArrays s)
forall s. TextBuilder s -> STRef s (TextArrays s)
tbArrays TextBuilder s
b)
    TextArrays s
arrs' <-
        if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< MVector s Word8 -> Int
forall a s. Unbox a => MVector s a -> Int
VUM.length (TextArrays s -> MVector s Word8
forall s. TextArrays s -> MVector s Word8
taValid TextArrays s
arrs) Bool -> Bool -> Bool
&& Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
extra Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= TextArrays s -> Int
forall s. TextArrays s -> Int
taByteCap TextArrays s
arrs
            then TextArrays s -> ST s (TextArrays s)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TextArrays s
arrs
            else TextBuilder s -> TextArrays s -> Int -> Int -> ST s (TextArrays s)
forall s.
TextBuilder s -> TextArrays s -> Int -> Int -> ST s (TextArrays s)
growText TextBuilder s
b TextArrays s
arrs (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
extra)
    (Int, Int, TextArrays s) -> ST s (Int, Int, TextArrays s)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
n, Int
pos, TextArrays s
arrs')
{-# INLINE reserveText #-}

growText :: TextBuilder s -> TextArrays s -> Int -> Int -> ST s (TextArrays s)
growText :: forall s.
TextBuilder s -> TextArrays s -> Int -> Int -> ST s (TextArrays s)
growText TextBuilder s
b (TextArrays MArray s
bytes Int
bcap MVector s Int
offsets MVector s Word8
val) Int
needRows Int
needBytes = do
    let rcap :: Int
rcap = MVector s Word8 -> Int
forall a s. Unbox a => MVector s a -> Int
VUM.length MVector s Word8
val
    (MVector s Int
offsets', MVector s Word8
val') <-
        if Int
needRows Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
rcap
            then do
                let rcap' :: Int
rcap' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
rcap) Int
needRows
                MVector s Int
o <- MVector (PrimState (ST s)) Int
-> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
VUM.unsafeGrow MVector s Int
MVector (PrimState (ST s)) Int
offsets (Int
rcap' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
rcap)
                MVector s Word8
v <- MVector (PrimState (ST s)) Word8
-> Int -> ST s (MVector (PrimState (ST s)) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
VUM.unsafeGrow MVector s Word8
MVector (PrimState (ST s)) Word8
val (Int
rcap' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
rcap)
                (MVector s Int, MVector s Word8)
-> ST s (MVector s Int, MVector s Word8)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MVector s Int
o, MVector s Word8
v)
            else (MVector s Int, MVector s Word8)
-> ST s (MVector s Int, MVector s Word8)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MVector s Int
offsets, MVector s Word8
val)
    (MArray s
bytes', Int
bcap') <-
        if Int
needBytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
bcap
            then do
                let cap' :: Int
cap' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
bcap) Int
needBytes
                MArray s
bs <- MArray s -> Int -> ST s (MArray s)
forall s. MArray s -> Int -> ST s (MArray s)
A.resizeM MArray s
bytes Int
cap'
                (MArray s, Int) -> ST s (MArray s, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MArray s
bs, Int
cap')
            else (MArray s, Int) -> ST s (MArray s, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MArray s
bytes, Int
bcap)
    let arrs :: TextArrays s
arrs = MArray s -> Int -> MVector s Int -> MVector s Word8 -> TextArrays s
forall s.
MArray s -> Int -> MVector s Int -> MVector s Word8 -> TextArrays s
TextArrays MArray s
bytes' Int
bcap' MVector s Int
offsets' MVector s Word8
val'
    STRef s (TextArrays s) -> TextArrays s -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef (TextBuilder s -> STRef s (TextArrays s)
forall s. TextBuilder s -> STRef s (TextArrays s)
tbArrays TextBuilder s
b) TextArrays s
arrs
    TextArrays s -> ST s (TextArrays s)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TextArrays s
arrs

{- | Freeze a 'TextBuilder' into a raw 'TextChunk' for byte-level merging
('mergeTextChunks'): no 'T.Text' values are created until chunks merge.
-}
freezeTextChunk :: TextBuilder s -> ST s TextChunk
freezeTextChunk :: forall s. TextBuilder s -> ST s TextChunk
freezeTextChunk TextBuilder s
b = do
    Int
n <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (TextBuilder s -> MVector s Int
forall s. TextBuilder s -> MVector s Int
tbCounters TextBuilder s
b) Int
cRows
    Int
anyNull <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (TextBuilder s -> MVector s Int
forall s. TextBuilder s -> MVector s Int
tbCounters TextBuilder s
b) Int
cAnyNull
    Int
used <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (TextBuilder s -> MVector s Int
forall s. TextBuilder s -> MVector s Int
tbCounters TextBuilder s
b) Int
cBytes
    TextArrays MArray s
bytes Int
bcap MVector s Int
offsets MVector s Word8
val <- STRef s (TextArrays s) -> ST s (TextArrays s)
forall s a. STRef s a -> ST s a
readSTRef (TextBuilder s -> STRef s (TextArrays s)
forall s. TextBuilder s -> STRef s (TextArrays s)
tbArrays TextBuilder s
b)
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
used Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
bcap) (MArray s -> Int -> ST s ()
forall s. MArray s -> Int -> ST s ()
A.shrinkM MArray s
bytes Int
used)
    Array
arr <- MArray s -> ST s Array
forall s. MArray s -> ST s Array
A.unsafeFreeze MArray s
bytes
    Vector Int
offs <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze (Int -> Int -> MVector s Int -> MVector s Int
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
VUM.slice Int
0 (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) MVector s Int
offsets)
    Maybe Bitmap
bm <-
        if Int
anyNull Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
            then Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Bitmap -> Maybe Bitmap) -> ST s Bitmap -> ST s (Maybe Bitmap)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> MVector s Word8 -> ST s Bitmap
forall s. Int -> MVector s Word8 -> ST s Bitmap
packValidity Int
n MVector s Word8
val
            else Maybe Bitmap -> ST s (Maybe Bitmap)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Bitmap
forall a. Maybe a
Nothing
    TextChunk -> ST s TextChunk
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array -> Int -> Vector Int -> Maybe Bitmap -> TextChunk
TextChunk Array
arr Int
used Vector Int
offs Maybe Bitmap
bm)

instance ColumnBuilder TextBuilder where
    appendNull :: forall s. TextBuilder s -> ST s ()
appendNull TextBuilder s
b = do
        (Int
n, Int
pos, TextArrays s
arrs) <- TextBuilder s -> Int -> ST s (Int, Int, TextArrays s)
forall s. TextBuilder s -> Int -> ST s (Int, Int, TextArrays s)
reserveText TextBuilder s
b Int
0
        Int
anyNull <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (TextBuilder s -> MVector s Int
forall s. TextBuilder s -> MVector s Int
tbCounters TextBuilder s
b) Int
cAnyNull
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
anyNull Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            MVector (PrimState (ST s)) Word8 -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> a -> m ()
VUM.set (Int -> Int -> MVector s Word8 -> MVector s Word8
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
VUM.slice Int
0 Int
n (TextArrays s -> MVector s Word8
forall s. TextArrays s -> MVector s Word8
taValid TextArrays s
arrs)) Word8
1
            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 (TextBuilder s -> MVector s Int
forall s. TextBuilder s -> MVector s Int
tbCounters TextBuilder s
b) Int
cAnyNull Int
1
        MVector (PrimState (ST s)) Word8 -> Int -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite (TextArrays s -> MVector s Word8
forall s. TextArrays s -> MVector s Word8
taValid TextArrays s
arrs) Int
n Word8
0
        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 (TextArrays s -> MVector s Int
forall s. TextArrays s -> MVector s Int
taOffsets TextArrays s
arrs) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
pos
        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 (TextBuilder s -> MVector s Int
forall s. TextBuilder s -> MVector s Int
tbCounters TextBuilder s
b) Int
cRows (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        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 (TextBuilder s -> MVector s Int
forall s. TextBuilder s -> MVector s Int
tbCounters TextBuilder s
b) Int
cBytes Int
pos
    {-# INLINE appendNull #-}

    builderLength :: forall s. TextBuilder s -> ST s Int
builderLength TextBuilder s
b = MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead (TextBuilder s -> MVector s Int
forall s. TextBuilder s -> MVector s Int
tbCounters TextBuilder s
b) Int
cRows

    freezeBuilder :: forall s. TextBuilder s -> ST s Column
freezeBuilder TextBuilder s
b = do
        TextChunk
chunk <- TextBuilder s -> ST s TextChunk
forall s. TextBuilder s -> ST s TextChunk
freezeTextChunk TextBuilder s
b
        Column -> ST s Column
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Column -> ST s Column) -> Column -> ST s Column
forall a b. (a -> b) -> a -> b
$! [TextChunk] -> Column
mergeTextChunks [TextChunk
chunk]