{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

module DataFrame.Internal.Column where

import qualified Data.Text as T
import qualified Data.Vector as VB
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Mutable as VBM
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM

import Control.Exception (throw)
import Control.Monad (forM_, when)
import Control.Monad.ST (ST, runST)
import Data.Bits (
    complement,
    popCount,
    setBit,
    shiftL,
    shiftR,
    testBit,
    (.&.),
 )
import Data.Kind (Type)
import Data.Maybe
import Data.Type.Equality (TestEquality (..))
import Data.Word (Word8)
import DataFrame.Errors
import DataFrame.Internal.PackedText (
    PackedTextData (..),
    packedGather,
    packedIndexText,
    packedLength,
    packedSlice,
    packedTake,
    sliceEqBytes,
 )
import DataFrame.Internal.Types
import System.IO.Unsafe (unsafePerformIO)
import System.Random
import Type.Reflection

-- | A bit-packed validity bitmap. Bit @i@ = 1 means row @i@ is valid (not null).
type Bitmap = VU.Vector Word8

{- | Type-erased column GADT. Pattern-matching on the constructor recovers the
representation; nullability is an optional bit-packed 'Bitmap' (@Nothing@ = no
nulls, @Just bm@ = bit @i@ set iff row @i@ is valid).
-}
data Column where
    BoxedColumn :: (Columnable a) => Maybe Bitmap -> VB.Vector a -> Column
    UnboxedColumn ::
        (Columnable a, VU.Unbox a) => Maybe Bitmap -> VU.Vector a -> Column
    -- Efficient intermediate formats.
    -- Bit-packed Text: shared UTF-8 byte buffer + row offsets + optional bitmap;
    -- Text is materialized on demand. Only CSV ingest emits this; user-built
    -- Text columns stay 'BoxedColumn'.
    PackedText :: Maybe Bitmap -> {-# UNPACK #-} !PackedTextData -> Column
    -- A join's same-named non-key column pair ('mergeColumns'): both sides
    -- keep their native (packed/dict/unboxed) representation; per-row 'These'
    -- values only materialize on element access ('materializeMerged').
    MergedColumn :: !Column -> !Column -> Column

{- | A mutable companion struct to dataframe columns.

Used mostly as an intermediate structure for I/O.
-}
data MutableColumn where
    MBoxedColumn :: (Columnable a) => VBM.IOVector a -> MutableColumn
    MUnboxedColumn :: (Columnable a, VU.Unbox a) => VUM.IOVector a -> MutableColumn

-- ---------------------------------------------------------------------------
-- Bitmap helpers
-- ---------------------------------------------------------------------------

-- | Test whether row @i@ is valid (not null) in a bitmap.
bitmapTestBit :: Bitmap -> Int -> Bool
bitmapTestBit :: Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i = Word8 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit (Bitmap -> Int -> Word8
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Bitmap
bm (Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3)) (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
7)
{-# INLINE bitmapTestBit #-}

-- | Build a fully-valid bitmap for @n@ rows (all bits set).
allValidBitmap :: Int -> Bitmap
allValidBitmap :: Int -> Bitmap
allValidBitmap Int
n =
    let bytes :: Int
bytes = (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3
        lastBits :: Int
lastBits = Int
n Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
7
        full :: Bitmap
full = Int -> Word8 -> Bitmap
forall a. Unbox a => Int -> a -> Vector a
VU.replicate (Int
bytes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Word8
0xFF
        lastByte :: Word8
lastByte = if Int
lastBits Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Word8
0xFF else (Word8
1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
lastBits) Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
- Word8
1
     in if Int
bytes Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Bitmap
forall a. Unbox a => Vector a
VU.empty else Bitmap -> Word8 -> Bitmap
forall a. Unbox a => Vector a -> a -> Vector a
VU.snoc Bitmap
full Word8
lastByte
{-# INLINE allValidBitmap #-}

{- | Build a bitmap from a @VU.Vector Word8@ validity vector
(1 = valid, 0 = null), as produced by Arrow / Parquet decoders.
-}
buildBitmapFromValid :: VU.Vector Word8 -> Bitmap
buildBitmapFromValid :: Bitmap -> Bitmap
buildBitmapFromValid Bitmap
valid =
    let n :: Int
n = Bitmap -> Int
forall a. Unbox a => Vector a -> Int
VU.length Bitmap
valid
        bytes :: Int
bytes = (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3
     in Int -> (Int -> Word8) -> Bitmap
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
bytes ((Int -> Word8) -> Bitmap) -> (Int -> Word8) -> Bitmap
forall a b. (a -> b) -> a -> b
$ \Int
b ->
            let base :: Int
base = Int
b Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
3
                setBitIf :: Word8 -> Int -> Word8
setBitIf Word8
acc Int
bit =
                    let idx :: Int
idx = Int
base Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bit
                     in if Int
idx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n Bool -> Bool -> Bool
&& Bitmap -> Int -> Word8
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Bitmap
valid Int
idx Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0
                            then Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
setBit Word8
acc Int
bit
                            else Word8
acc
             in (Word8 -> Int -> Word8) -> Word8 -> [Int] -> Word8
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Word8 -> Int -> Word8
setBitIf (Word8
0 :: Word8) [Int
0 .. Int
7]

{- | Build a bitmap from a list of null-row indices.
@nullIdxs@ are the positions that are NULL.
-}
buildBitmapFromNulls :: Int -> [Int] -> Bitmap
buildBitmapFromNulls :: Int -> [Int] -> Bitmap
buildBitmapFromNulls Int
n [Int]
nullIdxs =
    let base :: Bitmap
base = Int -> Bitmap
allValidBitmap Int
n
     in (forall s. MVector s Word8 -> ST s ()) -> Bitmap -> Bitmap
forall a.
Unbox a =>
(forall s. MVector s a -> ST s ()) -> Vector a -> Vector a
VU.modify
            ( \MVector s Word8
mv ->
                [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int]
nullIdxs ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
                    let byteIdx :: Int
byteIdx = Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3
                        bitIdx :: Int
bitIdx = Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
7
                    Word8
v <- MVector (PrimState (ST s)) Word8 -> Int -> ST s Word8
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector s Word8
MVector (PrimState (ST s)) Word8
mv Int
byteIdx
                    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
mv Int
byteIdx (Word8 -> Int -> Word8
clearBit8 Word8
v Int
bitIdx)
            )
            Bitmap
base
  where
    clearBit8 :: Word8 -> Int -> Word8
    clearBit8 :: Word8 -> Int -> Word8
clearBit8 Word8
b Int
bit = Word8
b Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8 -> Word8
forall a. Bits a => a -> a
complement (Word8
1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
bit)

-- | Slice a bitmap for rows @[start .. start+len-1]@.
bitmapSlice :: Int -> Int -> Bitmap -> Bitmap
bitmapSlice :: Int -> Int -> Bitmap -> Bitmap
bitmapSlice Int
start Int
len Bitmap
bm
    | Int
start Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
7 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 =
        let startByte :: Int
startByte = Int
start Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3
            bytes :: Int
bytes = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min ((Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3) (Bitmap -> Int
forall a. Unbox a => Vector a -> Int
VU.length Bitmap
bm Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
startByte)
         in Int -> Int -> Bitmap -> Bitmap
forall a. Unbox a => Int -> Int -> Vector a -> Vector a
VU.slice Int
startByte Int
bytes Bitmap
bm
    | Bool
otherwise =
        let n :: Int
n = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
len (Bitmap -> Int
forall a. Unbox a => Vector a -> Int
VU.length Bitmap
bm Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start)
         in Bitmap -> Bitmap
buildBitmapFromValid (Bitmap -> Bitmap) -> Bitmap -> Bitmap
forall a b. (a -> b) -> a -> b
$
                Int -> (Int -> Word8) -> Bitmap
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n ((Int -> Word8) -> Bitmap) -> (Int -> Word8) -> Bitmap
forall a b. (a -> b) -> a -> b
$
                    \Int
i -> if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) then Word8
1 else Word8
0

-- | Concatenate two bitmaps covering @n1@ and @n2@ rows respectively.
bitmapConcat :: Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat :: Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat Int
n1 Bitmap
bm1 Int
n2 Bitmap
bm2 =
    Bitmap -> Bitmap
buildBitmapFromValid (Bitmap -> Bitmap) -> Bitmap -> Bitmap
forall a b. (a -> b) -> a -> b
$
        Int -> (Int -> Word8) -> Bitmap
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate (Int
n1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n2) ((Int -> Word8) -> Bitmap) -> (Int -> Word8) -> Bitmap
forall a b. (a -> b) -> a -> b
$ \Int
i ->
            if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n1
                then if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm1 Int
i then Word8
1 else Word8
0
                else if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm2 (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n1) then Word8
1 else Word8
0

-- | Combine two bitmaps with AND (both must be valid for result to be valid).
mergeBitmaps :: Bitmap -> Bitmap -> Bitmap
mergeBitmaps :: Bitmap -> Bitmap -> Bitmap
mergeBitmaps = (Word8 -> Word8 -> Word8) -> Bitmap -> Bitmap -> Bitmap
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
VU.zipWith Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
(.&.)

{- | Materialize a nullable column from @VB.Vector (Maybe a)@; picks 'UnboxedColumn'
when @a@ is unboxable, else 'BoxedColumn'. Always attaches a bitmap so the column
reads as nullable even with no 'Nothing' values.
-}
fromMaybeVec :: forall a. (Columnable a) => VB.Vector (Maybe a) -> Column
fromMaybeVec :: forall a. Columnable a => Vector (Maybe a) -> Column
fromMaybeVec Vector (Maybe a)
v = case forall a. SBoolI (Unboxable a) => SBool (Unboxable a)
sUnbox @a of
    SBool (Unboxable a)
STrue -> Vector (Maybe a) -> Column
forall a. (Columnable a, Unbox a) => Vector (Maybe a) -> Column
fromMaybeVecUnboxed Vector (Maybe a)
v
    SBool (Unboxable a)
SFalse ->
        let n :: Int
n = Vector (Maybe a) -> Int
forall a. Vector a -> Int
VB.length Vector (Maybe a)
v
            nullIdxs :: [Int]
nullIdxs = [Int
i | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Maybe a -> Bool
forall a. Maybe a -> Bool
isNothing (Vector (Maybe a) -> Int -> Maybe a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector (Maybe a)
v Int
i)]
            bm :: Bitmap
bm = if [Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
nullIdxs then Int -> Bitmap
allValidBitmap Int
n else Int -> [Int] -> Bitmap
buildBitmapFromNulls Int
n [Int]
nullIdxs
            dat :: Vector a
dat = (Maybe a -> a) -> Vector (Maybe a) -> Vector a
forall a b. (a -> b) -> Vector a -> Vector b
VB.map (a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> a
forall a. [Char] -> a
errorWithoutStackTrace [Char]
"fromMaybeVec: Nothing slot")) Vector (Maybe a)
v
         in Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn (Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just Bitmap
bm) Vector a
dat

{- | Materialize a nullable 'UnboxedColumn' to @VB.Vector (Maybe a)@ using runST.
Always attaches a bitmap so the column is recognized as nullable even when
no 'Nothing' values are present (preserves the Maybe type marker).
-}
fromMaybeVecUnboxed ::
    forall a. (Columnable a, VU.Unbox a) => VB.Vector (Maybe a) -> Column
fromMaybeVecUnboxed :: forall a. (Columnable a, Unbox a) => Vector (Maybe a) -> Column
fromMaybeVecUnboxed Vector (Maybe a)
v =
    let n :: Int
n = Vector (Maybe a) -> Int
forall a. Vector a -> Int
VB.length Vector (Maybe a)
v
        nullIdxs :: [Int]
nullIdxs = [Int
i | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Maybe a -> Bool
forall a. Maybe a -> Bool
isNothing (Vector (Maybe a) -> Int -> Maybe a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector (Maybe a)
v Int
i)]
        bm :: Bitmap
bm = if [Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
nullIdxs then Int -> Bitmap
allValidBitmap Int
n else Int -> [Int] -> Bitmap
buildBitmapFromNulls Int
n [Int]
nullIdxs
        dat :: Vector a
dat = (forall s. ST s (Vector a)) -> Vector a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector a)) -> Vector a)
-> (forall s. ST s (Vector a)) -> Vector a
forall a b. (a -> b) -> a -> b
$ do
            MVector s a
mv <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
n
            Vector (Maybe a) -> (Int -> Maybe a -> ST s ()) -> ST s ()
forall (m :: * -> *) (v :: * -> *) a b.
(Monad m, Vector v a) =>
v a -> (Int -> a -> m b) -> m ()
VG.iforM_ Vector (Maybe a)
v ((Int -> Maybe a -> ST s ()) -> ST s ())
-> (Int -> Maybe a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i Maybe a
mx -> Maybe a -> (a -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe a
mx (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
mv Int
i)
            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 MVector s a
MVector (PrimState (ST s)) a
mv
     in 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
dat

-- | Whether row @i@ is null, respecting the bitmap.
columnElemIsNull :: Column -> Int -> Bool
columnElemIsNull :: Column -> Int -> Bool
columnElemIsNull (BoxedColumn (Just Bitmap
bm) Vector a
_) Int
i = Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i)
columnElemIsNull (UnboxedColumn (Just Bitmap
bm) Vector a
_) Int
i = Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i)
columnElemIsNull (PackedText (Just Bitmap
bm) PackedTextData
_) Int
i = Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i)
columnElemIsNull Column
_ Int
_ = Bool
False

-- | Return the 'Maybe Bitmap' from a column.
columnBitmap :: Column -> Maybe Bitmap
columnBitmap :: Column -> Maybe Bitmap
columnBitmap (BoxedColumn Maybe Bitmap
bm Vector a
_) = Maybe Bitmap
bm
columnBitmap (UnboxedColumn Maybe Bitmap
bm Vector a
_) = Maybe Bitmap
bm
columnBitmap (PackedText Maybe Bitmap
bm PackedTextData
_) = Maybe Bitmap
bm
columnBitmap (MergedColumn Column
_ Column
_) = Maybe Bitmap
forall a. Maybe a
Nothing

{- | Decode a 'PackedText' into a @BoxedColumn Text@ (bit-identical to
materializing at freeze). Identity on every other column.
-}
materializePacked :: Column -> Column
materializePacked :: Column -> Column
materializePacked (PackedText Maybe Bitmap
bm PackedTextData
p) =
    Maybe Bitmap -> Vector Text -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
bm (Int -> (Int -> Text) -> Vector Text
forall a. Int -> (Int -> a) -> Vector a
VB.generate (PackedTextData -> Int
packedLength PackedTextData
p) (PackedTextData -> Int -> Text
packedIndexText PackedTextData
p))
materializePacked Column
c = Column
c
{-# INLINE materializePacked #-}

-- | Whether a column is a 'PackedText'.
isPackedText :: Column -> Bool
isPackedText :: Column -> Bool
isPackedText (PackedText Maybe Bitmap
_ PackedTextData
_) = Bool
True
isPackedText Column
_ = Bool
False
{-# INLINE isPackedText #-}

-- | Whether a column is a 'MergedColumn'.
isMergedColumn :: Column -> Bool
isMergedColumn :: Column -> Bool
isMergedColumn (MergedColumn Column
_ Column
_) = Bool
True
isMergedColumn Column
_ = Bool
False
{-# INLINE isMergedColumn #-}

{- | 'MergedColumn' defers element construction, so forcing must still surface
the one deferred error — a row null on both sides — inside strict IO/executor
boundaries. O(rows) bitmap walk, no allocation; both-null needs a bitmap on
each side, so anything else passes immediately.
-}
checkMergedNoBothNull :: Column -> Column -> ()
checkMergedNoBothNull :: Column -> Column -> ()
checkMergedNoBothNull Column
a Column
b = case (Column -> Maybe Bitmap
columnBitmap Column
a, Column -> Maybe Bitmap
columnBitmap Column
b) of
    (Just Bitmap
ba, Just Bitmap
bb) ->
        let !n :: Int
n = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Column -> Int
columnLength Column
a) (Column -> Int
columnLength Column
b)
            go :: Int -> ()
go !Int
i
                | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = ()
                | Bitmap -> Int -> Bool
bitmapTestBit Bitmap
ba Int
i Bool -> Bool -> Bool
|| Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bb Int
i = Int -> ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                | Bool
otherwise = [Char] -> ()
forall a. HasCallStack => [Char] -> a
error [Char]
"mergeColumns: both null"
         in Int -> ()
go Int
0
    (Maybe Bitmap, Maybe Bitmap)
_ -> ()

-- ---------------------------------------------------------------------------
-- End bitmap helpers
-- ---------------------------------------------------------------------------

{- | A wrapper around the type-erased 'Column' carrying a phantom element type,
used to type-check expressions. The phantom is not guaranteed to match the
underlying vector's type.
-}
data TypedColumn a where
    TColumn :: (Columnable a) => Column -> TypedColumn a

instance (Eq a) => Eq (TypedColumn a) where
    (==) :: (Eq a) => TypedColumn a -> TypedColumn a -> Bool
    == :: Eq a => TypedColumn a -> TypedColumn a -> Bool
(==) (TColumn Column
a) (TColumn Column
b) = Column
a Column -> Column -> Bool
forall a. Eq a => a -> a -> Bool
== Column
b

-- | Gets the underlying value from a TypedColumn.
unwrapTypedColumn :: TypedColumn a -> Column
unwrapTypedColumn :: forall a. TypedColumn a -> Column
unwrapTypedColumn (TColumn Column
value) = Column
value

-- | Gets the underlying vector from a TypedColumn.
vectorFromTypedColumn :: TypedColumn a -> VB.Vector a
vectorFromTypedColumn :: forall a. TypedColumn a -> Vector a
vectorFromTypedColumn (TColumn Column
value) = (DataFrameException -> Vector a)
-> (Vector a -> Vector a)
-> Either DataFrameException (Vector a)
-> Vector a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either DataFrameException -> Vector a
forall a e. Exception e => e -> a
throw Vector a -> Vector a
forall a. a -> a
id (Column -> Either DataFrameException (Vector a)
forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector Column
value)

-- | Checks if a column contains missing values (has a bitmap).
hasMissing :: Column -> Bool
hasMissing :: Column -> Bool
hasMissing (BoxedColumn (Just Bitmap
_) Vector a
_) = Bool
True
hasMissing (UnboxedColumn (Just Bitmap
_) Vector a
_) = Bool
True
hasMissing (PackedText (Just Bitmap
_) PackedTextData
_) = Bool
True
hasMissing Column
_ = Bool
False

-- | Checks if a column contains only missing values.
allMissing :: Column -> Bool
allMissing :: Column -> Bool
allMissing (BoxedColumn (Just Bitmap
bm) Vector a
col) = (Word8 -> Bool) -> Bitmap -> Bool
forall a. Unbox a => (a -> Bool) -> Vector a -> Bool
VU.all (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0) Bitmap
bm Bool -> Bool -> Bool
&& Bool -> Bool
not (Vector a -> Bool
forall a. Vector a -> Bool
VB.null Vector a
col)
allMissing (UnboxedColumn (Just Bitmap
bm) Vector a
col) = (Word8 -> Bool) -> Bitmap -> Bool
forall a. Unbox a => (a -> Bool) -> Vector a -> Bool
VU.all (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0) Bitmap
bm Bool -> Bool -> Bool
&& Bool -> Bool
not (Vector a -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector a
col)
allMissing (PackedText (Just Bitmap
bm) PackedTextData
p) = (Word8 -> Bool) -> Bitmap -> Bool
forall a. Unbox a => (a -> Bool) -> Vector a -> Bool
VU.all (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0) Bitmap
bm Bool -> Bool -> Bool
&& PackedTextData -> Int
packedLength PackedTextData
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
allMissing Column
_ = Bool
False

-- | Checks if a column contains numeric values.
isNumeric :: Column -> Bool
isNumeric :: Column -> Bool
isNumeric c :: Column
c@(MergedColumn Column
_ Column
_) = Column -> Bool
isNumeric (Column -> Column
mergedHead Column
c)
isNumeric (UnboxedColumn Maybe Bitmap
_ (Vector a
_vec :: VU.Vector a)) = case forall a. SBoolI (Numeric a) => SBool (Numeric a)
sNumeric @a of
    SBool (Numeric a)
STrue -> Bool
True
    SBool (Numeric a)
_ -> Bool
False
isNumeric (BoxedColumn Maybe Bitmap
_ (Vector a
_vec :: VB.Vector a)) = case TypeRep a -> TypeRep Integer -> Maybe (a :~: Integer)
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 @Integer) of
    Maybe (a :~: Integer)
Nothing -> Bool
False
    Just a :~: Integer
Refl -> Bool
True
isNumeric (PackedText Maybe Bitmap
_ PackedTextData
_) = Bool
False

{- | Whether the column stores element type @a@. For nullable columns, also
'True' when @a = Maybe b@ and the column stores @b@ internally.
-}
hasElemType :: forall a. (Columnable a) => Column -> Bool
hasElemType :: forall a. Columnable a => Column -> Bool
hasElemType = \case
    BoxedColumn Maybe Bitmap
bm (Vector a
_column :: VB.Vector b) -> Maybe Bitmap -> TypeRep a -> Bool
forall b. Maybe Bitmap -> TypeRep b -> Bool
checkBoxed Maybe Bitmap
bm (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b)
    UnboxedColumn Maybe Bitmap
bm (Vector a
_column :: VU.Vector b) -> Maybe Bitmap -> TypeRep a -> Bool
forall b. Maybe Bitmap -> TypeRep b -> Bool
checkUnboxed Maybe Bitmap
bm (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b)
    PackedText Maybe Bitmap
bm PackedTextData
_ -> Maybe Bitmap -> TypeRep Text -> Bool
forall b. Maybe Bitmap -> TypeRep b -> Bool
checkBoxed Maybe Bitmap
bm (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @T.Text)
    c :: Column
c@(MergedColumn Column
_ Column
_) -> forall a. Columnable a => Column -> Bool
hasElemType @a (Column -> Column
mergedHead Column
c)
  where
    directMatch :: forall (b :: Type). TypeRep b -> Bool
    directMatch :: forall b. TypeRep b -> Bool
directMatch = Maybe (a :~: b) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (a :~: b) -> Bool)
-> (TypeRep b -> Maybe (a :~: b)) -> TypeRep b -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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)
    checkMaybe :: forall (b :: Type). TypeRep b -> Bool
    checkMaybe :: forall b. TypeRep b -> Bool
checkMaybe TypeRep b
tb = case forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a of
        App TypeRep a
tMaybe TypeRep b
tInner -> case TypeRep a -> TypeRep Maybe -> Maybe (a :~~: Maybe)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tMaybe (forall {k} (a :: k). Typeable a => TypeRep a
forall (a :: * -> *). Typeable a => TypeRep a
typeRep @Maybe) of
            Just a :~~: Maybe
HRefl -> Maybe (b :~: b) -> Bool
forall a. Maybe a -> Bool
isJust (TypeRep b -> TypeRep b -> Maybe (b :~: b)
forall (a :: k1) (b :: k1).
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 TypeRep b
tInner TypeRep b
TypeRep b
tb)
            Maybe (a :~~: Maybe)
Nothing -> Bool
False
        TypeRep a
_ -> Bool
False
    checkBoxed :: forall (b :: Type). Maybe Bitmap -> TypeRep b -> Bool
    checkBoxed :: forall b. Maybe Bitmap -> TypeRep b -> Bool
checkBoxed Maybe Bitmap
bm TypeRep b
tb = TypeRep b -> Bool
forall b. TypeRep b -> Bool
directMatch TypeRep b
tb Bool -> Bool -> Bool
|| (Maybe Bitmap -> Bool
forall a. Maybe a -> Bool
isJust Maybe Bitmap
bm Bool -> Bool -> Bool
&& TypeRep b -> Bool
forall b. TypeRep b -> Bool
checkMaybe TypeRep b
tb)
    checkUnboxed :: forall (b :: Type). Maybe Bitmap -> TypeRep b -> Bool
    checkUnboxed :: forall b. Maybe Bitmap -> TypeRep b -> Bool
checkUnboxed Maybe Bitmap
bm TypeRep b
tb = TypeRep b -> Bool
forall b. TypeRep b -> Bool
directMatch TypeRep b
tb Bool -> Bool -> Bool
|| (Maybe Bitmap -> Bool
forall a. Maybe a -> Bool
isJust Maybe Bitmap
bm Bool -> Bool -> Bool
&& TypeRep b -> Bool
forall b. TypeRep b -> Bool
checkMaybe TypeRep b
tb)

-- | An internal/debugging function to get the column type of a column.
columnVersionString :: Column -> String
columnVersionString :: Column -> [Char]
columnVersionString Column
column = case Column
column of
    BoxedColumn Maybe Bitmap
Nothing Vector a
_ -> [Char]
"Boxed"
    BoxedColumn (Just Bitmap
_) Vector a
_ -> [Char]
"NullableBoxed"
    UnboxedColumn Maybe Bitmap
Nothing Vector a
_ -> [Char]
"Unboxed"
    UnboxedColumn (Just Bitmap
_) Vector a
_ -> [Char]
"NullableUnboxed"
    PackedText Maybe Bitmap
Nothing PackedTextData
_ -> [Char]
"Boxed"
    PackedText (Just Bitmap
_) PackedTextData
_ -> [Char]
"NullableBoxed"
    MergedColumn Column
_ Column
_ -> Column -> [Char]
columnVersionString (Column -> Column
mergedHead Column
column)

{- | An internal/debugging function to get the type stored in the outermost vector
of a column.
-}
columnTypeString :: Column -> String
columnTypeString :: Column -> [Char]
columnTypeString Column
column = case Column
column of
    BoxedColumn Maybe Bitmap
Nothing (Vector a
_ :: VB.Vector a) -> TypeRep a -> [Char]
forall a. Show a => a -> [Char]
show (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a)
    BoxedColumn (Just Bitmap
_) (Vector a
_ :: VB.Vector a) -> forall a. Typeable a => [Char]
forall {k} (a :: k). Typeable a => [Char]
showMaybeType @a
    UnboxedColumn Maybe Bitmap
Nothing (Vector a
_ :: VU.Vector a) -> TypeRep a -> [Char]
forall a. Show a => a -> [Char]
show (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a)
    UnboxedColumn (Just Bitmap
_) (Vector a
_ :: VU.Vector a) -> forall a. Typeable a => [Char]
forall {k} (a :: k). Typeable a => [Char]
showMaybeType @a
    PackedText Maybe Bitmap
Nothing PackedTextData
_ -> TypeRep Text -> [Char]
forall a. Show a => a -> [Char]
show (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @T.Text)
    PackedText (Just Bitmap
_) PackedTextData
_ -> forall a. Typeable a => [Char]
forall {k} (a :: k). Typeable a => [Char]
showMaybeType @T.Text
    MergedColumn Column
_ Column
_ -> Column -> [Char]
columnTypeString (Column -> Column
mergedHead Column
column)
  where
    showMaybeType :: forall a. (Typeable a) => String
    showMaybeType :: forall {k} (a :: k). Typeable a => [Char]
showMaybeType =
        let s :: [Char]
s = TypeRep a -> [Char]
forall a. Show a => a -> [Char]
show (forall (a :: k). Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a)
         in [Char]
"Maybe " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ if Char
' ' Char -> [Char] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Char]
s then [Char]
"(" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
s [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")" else [Char]
s

instance (Show a) => Show (TypedColumn a) where
    show :: (Show a) => TypedColumn a -> String
    show :: Show a => TypedColumn a -> [Char]
show (TColumn Column
col) = Column -> [Char]
forall a. Show a => a -> [Char]
show Column
col

{- | Force evaluation of all elements in a column. Replacement for the removed
@instance NFData Column@; used by the IO and lazy-executor strict paths.
-}
forceColumn :: Column -> ()
forceColumn :: Column -> ()
forceColumn (BoxedColumn Maybe Bitmap
Nothing (Vector a
v :: VB.Vector a)) = (() -> a -> ()) -> () -> Vector a -> ()
forall a b. (a -> b -> a) -> a -> Vector b -> a
VB.foldl' ((a -> ()) -> () -> a -> ()
forall a b. a -> b -> a
const (a -> () -> ()
forall a b. a -> b -> b
`seq` ())) () Vector a
v
forceColumn (BoxedColumn (Just Bitmap
bm) (Vector a
v :: VB.Vector a)) =
    let n :: Int
n = Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
v
        go :: Int -> ()
go !Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = ()
            | Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i = Vector a -> Int -> a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector a
v Int
i a -> () -> ()
forall a b. a -> b -> b
`seq` Int -> ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            | Bool
otherwise = Int -> ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
     in Int -> ()
go Int
0
forceColumn (UnboxedColumn Maybe Bitmap
_ Vector a
v) = Vector a
v Vector a -> () -> ()
forall a b. a -> b -> b
`seq` ()
forceColumn (PackedText Maybe Bitmap
_ (PackedTextData Array
arr PackedOffsets
offs Maybe PackedSel
sel Bool
_)) = Array
arr Array -> () -> ()
forall a b. a -> b -> b
`seq` PackedOffsets
offs PackedOffsets -> () -> ()
forall a b. a -> b -> b
`seq` Maybe PackedSel
sel Maybe PackedSel -> () -> ()
forall a b. a -> b -> b
`seq` ()
forceColumn (MergedColumn Column
a Column
b) =
    Column -> ()
forceColumn Column
a () -> () -> ()
forall a b. a -> b -> b
`seq` Column -> ()
forceColumn Column
b () -> () -> ()
forall a b. a -> b -> b
`seq` Column -> Column -> ()
checkMergedNoBothNull Column
a Column
b

instance Show Column where
    show :: Column -> String
    show :: Column -> [Char]
show c :: Column
c@(MergedColumn Column
_ Column
_) = Column -> [Char]
forall a. Show a => a -> [Char]
show (Column -> Column
materializeMerged Column
c)
    show (BoxedColumn Maybe Bitmap
Nothing Vector a
column) = Vector a -> [Char]
forall a. Show a => a -> [Char]
show Vector a
column
    show (BoxedColumn (Just Bitmap
bm) Vector a
column) =
        let n :: Int
n = Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
column
            elems :: [[Char]]
elems =
                [ if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i then a -> [Char]
forall a. Show a => a -> [Char]
show (Vector a -> Int -> a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector a
column Int
i) else [Char]
"null"
                | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                ]
         in [Char]
"[" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ ([Char] -> [Char] -> [Char]) -> [Char] -> [[Char]] -> [Char]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\[Char]
acc [Char]
e -> if [Char] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Char]
acc then [Char]
e else [Char]
acc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"," [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
e) [Char]
"" [[Char]]
elems [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"]"
    show (UnboxedColumn Maybe Bitmap
Nothing Vector a
column) = Vector a -> [Char]
forall a. Show a => a -> [Char]
show Vector a
column
    show (UnboxedColumn (Just Bitmap
bm) Vector a
column) =
        let n :: Int
n = Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
column
            elems :: [[Char]]
elems =
                [ if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i then a -> [Char]
forall a. Show a => a -> [Char]
show (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
column Int
i) else [Char]
"null"
                | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                ]
         in [Char]
"[" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ ([Char] -> [Char] -> [Char]) -> [Char] -> [[Char]] -> [Char]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\[Char]
acc [Char]
e -> if [Char] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Char]
acc then [Char]
e else [Char]
acc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"," [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
e) [Char]
"" [[Char]]
elems [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"]"
    show c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
_) = Column -> [Char]
forall a. Show a => a -> [Char]
show (Column -> Column
materializePacked Column
c)

{- | Compare two nullable boxed columns element by element, skipping null slots.
Uses a manual loop to avoid stream fusion forcing null-slot error thunks.
-}
eqBoxedCols ::
    (Eq a) => Maybe Bitmap -> VB.Vector a -> Maybe Bitmap -> VB.Vector a -> Bool
eqBoxedCols :: forall a.
Eq a =>
Maybe Bitmap -> Vector a -> Maybe Bitmap -> Vector a -> Bool
eqBoxedCols Maybe Bitmap
bm1 Vector a
a Maybe Bitmap
bm2 Vector a
b
    | Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
b = Bool
False
    | Bool
otherwise = Int -> Bool
go Int
0
  where
    !n :: Int
n = Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
a
    go :: Int -> Bool
go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Bool
True
        | Bool
nullA Bool -> Bool -> Bool
|| Bool
nullB = (Bool
nullA Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
nullB) Bool -> Bool -> Bool
&& Int -> Bool
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Vector a -> Int -> a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector a
a Int
i a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a -> Int -> a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector a
b Int
i = Int -> Bool
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = Bool
False
      where
        nullA :: Bool
nullA = Bool -> (Bitmap -> Bool) -> Maybe Bitmap -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (\Bitmap
bm -> Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i)) Maybe Bitmap
bm1
        nullB :: Bool
nullB = Bool -> (Bitmap -> Bool) -> Maybe Bitmap -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (\Bitmap
bm -> Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i)) Maybe Bitmap
bm2
{-# INLINE eqBoxedCols #-}

instance Eq Column where
    (==) :: Column -> Column -> Bool
    == :: Column -> Column -> Bool
(==) (BoxedColumn Maybe Bitmap
bm1 (Vector a
a :: VB.Vector t1)) (BoxedColumn Maybe Bitmap
bm2 (Vector a
b :: VB.Vector t2)) =
        case TypeRep a -> TypeRep a -> Maybe (a :~: 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 @t1) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @t2) of
            Maybe (a :~: a)
Nothing -> Bool
False
            Just a :~: a
Refl -> Maybe Bitmap -> Vector a -> Maybe Bitmap -> Vector a -> Bool
forall a.
Eq a =>
Maybe Bitmap -> Vector a -> Maybe Bitmap -> Vector a -> Bool
eqBoxedCols Maybe Bitmap
bm1 Vector a
a Maybe Bitmap
bm2 Vector a
Vector a
b
    (==) (UnboxedColumn Maybe Bitmap
bm1 (Vector a
a :: VU.Vector t1)) (UnboxedColumn Maybe Bitmap
bm2 (Vector a
b :: VU.Vector t2)) =
        case TypeRep a -> TypeRep a -> Maybe (a :~: 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 @t1) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @t2) of
            Maybe (a :~: a)
Nothing -> Bool
False
            Just a :~: a
Refl ->
                Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
b
                    Bool -> Bool -> Bool
&& Vector Bool -> Bool
VU.and
                        ( (Int -> a -> Bool) -> Vector a -> Vector Bool
forall a b.
(Unbox a, Unbox b) =>
(Int -> a -> b) -> Vector a -> Vector b
VU.imap
                            ( \Int
i a
x ->
                                let nullA :: Bool
nullA = Bool -> (Bitmap -> Bool) -> Maybe Bitmap -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (\Bitmap
bm -> Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i)) Maybe Bitmap
bm1
                                    nullB :: Bool
nullB = Bool -> (Bitmap -> Bool) -> Maybe Bitmap -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (\Bitmap
bm -> Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i)) Maybe Bitmap
bm2
                                 in if Bool
nullA Bool -> Bool -> Bool
|| Bool
nullB then Bool
nullA Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
nullB else a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
b Int
i
                            )
                            Vector a
Vector a
a
                        )
    (==) lhs :: Column
lhs@(MergedColumn Column
_ Column
_) Column
rhs = Column -> Column
materializeMerged Column
lhs Column -> Column -> Bool
forall a. Eq a => a -> a -> Bool
== Column
rhs
    (==) Column
lhs rhs :: Column
rhs@(MergedColumn Column
_ Column
_) = Column
lhs Column -> Column -> Bool
forall a. Eq a => a -> a -> Bool
== Column -> Column
materializeMerged Column
rhs
    (==) (PackedText Maybe Bitmap
bm1 PackedTextData
p1) (PackedText Maybe Bitmap
bm2 PackedTextData
p2) = Maybe Bitmap
-> PackedTextData -> Maybe Bitmap -> PackedTextData -> Bool
eqPackedCols Maybe Bitmap
bm1 PackedTextData
p1 Maybe Bitmap
bm2 PackedTextData
p2
    (==) lhs :: Column
lhs@(PackedText Maybe Bitmap
_ PackedTextData
_) Column
rhs = Column -> Column
materializePacked Column
lhs Column -> Column -> Bool
forall a. Eq a => a -> a -> Bool
== Column
rhs
    (==) Column
lhs rhs :: Column
rhs@(PackedText Maybe Bitmap
_ PackedTextData
_) = Column
lhs Column -> Column -> Bool
forall a. Eq a => a -> a -> Bool
== Column -> Column
materializePacked Column
rhs
    (==) Column
_ Column
_ = Bool
False

{- | Byte-slice equality of two packed-text columns, skipping null slots
(a null compares equal only to a null), mirroring 'eqBoxedCols'.
-}
eqPackedCols ::
    Maybe Bitmap -> PackedTextData -> Maybe Bitmap -> PackedTextData -> Bool
eqPackedCols :: Maybe Bitmap
-> PackedTextData -> Maybe Bitmap -> PackedTextData -> Bool
eqPackedCols Maybe Bitmap
bm1 PackedTextData
p1 Maybe Bitmap
bm2 PackedTextData
p2
    | PackedTextData -> Int
packedLength PackedTextData
p1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= PackedTextData -> Int
packedLength PackedTextData
p2 = Bool
False
    | Bool
otherwise = Int -> Bool
go Int
0
  where
    !n :: Int
n = PackedTextData -> Int
packedLength PackedTextData
p1
    go :: Int -> Bool
go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Bool
True
        | Bool
nullA Bool -> Bool -> Bool
|| Bool
nullB = (Bool
nullA Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
nullB) Bool -> Bool -> Bool
&& Int -> Bool
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise =
            let (Array
a1, Int
o1, Int
l1) = PackedTextData -> Int -> (Array, Int, Int)
packedSlice PackedTextData
p1 Int
i
                (Array
a2, Int
o2, Int
l2) = PackedTextData -> Int -> (Array, Int, Int)
packedSlice PackedTextData
p2 Int
i
             in Array -> Int -> Int -> Array -> Int -> Int -> Bool
sliceEqBytes Array
a1 Int
o1 Int
l1 Array
a2 Int
o2 Int
l2 Bool -> Bool -> Bool
&& Int -> Bool
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      where
        nullA :: Bool
nullA = Bool -> (Bitmap -> Bool) -> Maybe Bitmap -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (\Bitmap
bm -> Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i)) Maybe Bitmap
bm1
        nullB :: Bool
nullB = Bool -> (Bitmap -> Bool) -> Maybe Bitmap -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (\Bitmap
bm -> Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i)) Maybe Bitmap
bm2
{-# INLINE eqPackedCols #-}

{- | A class for converting a vector to a column of the appropriate type.
Given each Rep we tell the `toColumnRep` function which Column type to pick.
-}
class ColumnifyRep (r :: Rep) a where
    toColumnRep :: VB.Vector a -> Column

-- | Constraint synonym for what we can put into columns.
type Columnable a =
    ( Columnable' a
    , ColumnifyRep (KindOf a) a
    , UnboxIf a
    , IntegralIf a
    , FloatingIf a
    , SBoolI (Unboxable a)
    , SBoolI (Numeric a)
    , SBoolI (IntegralTypes a)
    , SBoolI (FloatingTypes a)
    )

instance
    (Columnable a, VU.Unbox a) =>
    ColumnifyRep 'RUnboxed a
    where
    toColumnRep :: (Columnable a, VUM.Unbox a) => VB.Vector a -> Column
    toColumnRep :: (Columnable a, Unbox a) => Vector a -> Column
toColumnRep Vector a
v = 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 -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VU.convert Vector a
v)

instance
    (Columnable a) =>
    ColumnifyRep 'RBoxed a
    where
    toColumnRep :: (Columnable a) => VB.Vector a -> Column
    toColumnRep :: Columnable a => Vector a -> Column
toColumnRep = Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
forall a. Maybe a
Nothing

instance
    (Columnable a) =>
    ColumnifyRep 'RNullableBoxed (Maybe a)
    where
    toColumnRep :: (Columnable a) => VB.Vector (Maybe a) -> Column
    toColumnRep :: Columnable a => Vector (Maybe a) -> Column
toColumnRep = Vector (Maybe a) -> Column
forall a. Columnable a => Vector (Maybe a) -> Column
fromMaybeVec

{- | O(n) Convert a vector to a column. Automatically picks the best representation of a vector to store the underlying data in.

__Examples:__

@
> import qualified Data.Vector as V
> fromVector (VB.fromList [(1 :: Int), 2, 3, 4])
[1,2,3,4]
@
-}
fromVector ::
    forall a.
    (Columnable a, ColumnifyRep (KindOf a) a) =>
    VB.Vector a -> Column
fromVector :: forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector = forall (r :: Rep) a. ColumnifyRep r a => Vector a -> Column
toColumnRep @(KindOf a)

{- | O(n) Convert an unboxed vector to a column. This avoids the extra conversion if you already have the data in an unboxed vector.

__Examples:__

@
> import qualified Data.Vector.Unboxed as V
> fromUnboxedVector (VB.fromList [(1 :: Int), 2, 3, 4])
[1,2,3,4]
@
-}
fromUnboxedVector ::
    forall a. (Columnable a, VU.Unbox a) => VU.Vector a -> Column
fromUnboxedVector :: forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector = Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
forall a. Maybe a
Nothing

{- | O(n) Convert a list to a column. Automatically picks the best representation of a vector to store the underlying data in.

__Examples:__

@
> fromList [(1 :: Int), 2, 3, 4]
[1,2,3,4]
@
-}
fromList ::
    forall a.
    (Columnable a, ColumnifyRep (KindOf a) a) =>
    [a] -> Column
fromList :: forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList = forall (r :: Rep) a. ColumnifyRep r a => Vector a -> Column
toColumnRep @(KindOf a) (Vector a -> Column) -> ([a] -> Vector a) -> [a] -> Column
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Vector a
forall a. [a] -> Vector a
VB.fromList

{- | O(n) Create a column of random elements within a range.

Takes a random number generator, a length, and a lower and upper bound for the random values.

__Examples:__

@
> import System.Random (mkStdGen)
> mkRandom (mkStdGen 42) 4 0 10
[4,2,6,5]
@
-}
mkRandom ::
    (RandomGen g, Columnable a, ColumnifyRep (KindOf a) a, UniformRange a) =>
    g -> Int -> a -> a -> Column
mkRandom :: forall g a.
(RandomGen g, Columnable a, ColumnifyRep (KindOf a) a,
 UniformRange a) =>
g -> Int -> a -> a -> Column
mkRandom g
pureGen Int
k a
lo a
hi = [a] -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList ([a] -> Column) -> [a] -> Column
forall a b. (a -> b) -> a -> b
$ g -> Int -> [a]
go g
pureGen Int
k
  where
    go :: g -> Int -> [a]
go g
_g Int
0 = []
    go g
g Int
n =
        let
            (!a
v, !g
g') = (a, a) -> g -> (a, g)
forall a g. (UniformRange a, RandomGen g) => (a, a) -> g -> (a, g)
uniformR (a
lo, a
hi) g
g
         in
            a
v a -> [a] -> [a]
forall a. a -> [a] -> [a]
: g -> Int -> [a]
go g
g' (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

-- An internal helper for type errors
throwTypeMismatch ::
    forall (a :: Type) (b :: Type).
    (Typeable a, Typeable b) => Either DataFrameException Column
throwTypeMismatch :: forall a b.
(Typeable a, Typeable b) =>
Either DataFrameException Column
throwTypeMismatch =
    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 b a -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
            MkTypeErrorContext
                { userType :: Either [Char] (TypeRep b)
userType = TypeRep b -> Either [Char] (TypeRep b)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b)
                , expectedType :: Either [Char] (TypeRep a)
expectedType = TypeRep a -> Either [Char] (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)
                , callingFunctionName :: Maybe [Char]
callingFunctionName = Maybe [Char]
forall a. Maybe a
Nothing
                , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                }

-- | An internal function to map a function over the values of a column.
mapColumn ::
    forall b c.
    (Columnable b, Columnable c) =>
    (b -> c) -> Column -> Either DataFrameException Column
mapColumn :: forall b c.
(Columnable b, Columnable c) =>
(b -> c) -> Column -> Either DataFrameException Column
mapColumn b -> c
f = \case
    BoxedColumn Maybe Bitmap
bm (Vector a
col :: VB.Vector a) -> Maybe Bitmap -> Vector a -> Either DataFrameException Column
forall a.
Columnable a =>
Maybe Bitmap -> Vector a -> Either DataFrameException Column
runBoxed Maybe Bitmap
bm Vector a
col
    UnboxedColumn Maybe Bitmap
bm (Vector a
col :: VU.Vector a) -> Maybe Bitmap -> Vector a -> Either DataFrameException Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Either DataFrameException Column
runUnboxed Maybe Bitmap
bm Vector a
col
    c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
_) -> (b -> c) -> Column -> Either DataFrameException Column
forall b c.
(Columnable b, Columnable c) =>
(b -> c) -> Column -> Either DataFrameException Column
mapColumn b -> c
f (Column -> Column
materializePacked Column
c)
    c :: Column
c@(MergedColumn Column
_ Column
_) -> (b -> c) -> Column -> Either DataFrameException Column
forall b c.
(Columnable b, Columnable c) =>
(b -> c) -> Column -> Either DataFrameException Column
mapColumn b -> c
f (Column -> Column
materializeMerged Column
c)
  where
    runBoxed ::
        forall a.
        (Columnable a) =>
        Maybe Bitmap -> VB.Vector a -> Either DataFrameException Column
    runBoxed :: forall a.
Columnable a =>
Maybe Bitmap -> Vector a -> Either DataFrameException Column
runBoxed Maybe Bitmap
bm Vector a
col = 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 a)) of
        Just b :~: Maybe a
Refl ->
            let !n :: Int
n = Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
col
             in 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
$ case forall a. SBoolI (Unboxable a) => SBool (Unboxable a)
sUnbox @c of
                    SBool (Unboxable c)
STrue -> Maybe Bitmap -> Vector c -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector c -> Column) -> Vector c -> Column
forall a b. (a -> b) -> a -> b
$
                        Int -> (Int -> c) -> Vector c
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n ((Int -> c) -> Vector c) -> (Int -> c) -> Vector c
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                            b -> c
f
                                ( if Bool -> (Bitmap -> Bool) -> Maybe Bitmap -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (Bitmap -> Int -> Bool
`bitmapTestBit` Int
i) Maybe Bitmap
bm
                                    then a -> Maybe a
forall a. a -> Maybe a
Just (Vector a -> Int -> a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector a
col Int
i)
                                    else b
Maybe a
forall a. Maybe a
Nothing
                                )
                    SBool (Unboxable c)
SFalse -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector @c (Vector c -> Column) -> Vector c -> Column
forall a b. (a -> b) -> a -> b
$
                        Int -> (Int -> c) -> Vector c
forall a. Int -> (Int -> a) -> Vector a
VB.generate Int
n ((Int -> c) -> Vector c) -> (Int -> c) -> Vector c
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                            b -> c
f
                                ( if Bool -> (Bitmap -> Bool) -> Maybe Bitmap -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (Bitmap -> Int -> Bool
`bitmapTestBit` Int
i) Maybe Bitmap
bm
                                    then a -> Maybe a
forall a. a -> Maybe a
Just (Vector a -> Int -> a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector a
col Int
i)
                                    else b
Maybe a
forall a. Maybe a
Nothing
                                )
        Maybe (b :~: Maybe a)
Nothing -> 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
            Just a :~: b
Refl ->
                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
$ case forall a. SBoolI (Unboxable a) => SBool (Unboxable a)
sUnbox @c of
                    SBool (Unboxable c)
STrue -> Maybe Bitmap -> Vector c -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
bm (Int -> (Int -> c) -> Vector c
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
col) (b -> c
f (b -> c) -> (Int -> b) -> Int -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector b -> Int -> b
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector b
Vector a
col))
                    SBool (Unboxable c)
SFalse -> case Maybe Bitmap
bm of
                        Maybe Bitmap
Nothing -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector @c ((b -> c) -> Vector b -> Vector c
forall a b. (a -> b) -> Vector a -> Vector b
VB.map b -> c
f Vector b
Vector a
col)
                        Just Bitmap
_ -> Maybe Bitmap -> Vector c -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
bm ((b -> c) -> Vector b -> Vector c
forall a b. (a -> b) -> Vector a -> Vector b
VB.map b -> c
f Vector b
Vector a
col)
            Maybe (a :~: b)
Nothing -> forall a b.
(Typeable a, Typeable b) =>
Either DataFrameException Column
throwTypeMismatch @a @b

    runUnboxed ::
        forall a.
        (Columnable a, VU.Unbox a) =>
        Maybe Bitmap -> VU.Vector a -> Either DataFrameException Column
    runUnboxed :: forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Either DataFrameException Column
runUnboxed Maybe Bitmap
bm Vector a
col = 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 a)) of
        Just b :~: Maybe a
Refl ->
            let !n :: Int
n = Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
col
             in 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
$ case forall a. SBoolI (Unboxable a) => SBool (Unboxable a)
sUnbox @c of
                    SBool (Unboxable c)
STrue -> Maybe Bitmap -> Vector c -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector c -> Column) -> Vector c -> Column
forall a b. (a -> b) -> a -> b
$
                        Int -> (Int -> c) -> Vector c
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n ((Int -> c) -> Vector c) -> (Int -> c) -> Vector c
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                            b -> c
f
                                ( if Bool -> (Bitmap -> Bool) -> Maybe Bitmap -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (Bitmap -> Int -> Bool
`bitmapTestBit` Int
i) Maybe Bitmap
bm
                                    then a -> Maybe a
forall a. a -> Maybe a
Just (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
col Int
i)
                                    else b
Maybe a
forall a. Maybe a
Nothing
                                )
                    SBool (Unboxable c)
SFalse -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector @c (Vector c -> Column) -> Vector c -> Column
forall a b. (a -> b) -> a -> b
$
                        Int -> (Int -> c) -> Vector c
forall a. Int -> (Int -> a) -> Vector a
VB.generate Int
n ((Int -> c) -> Vector c) -> (Int -> c) -> Vector c
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                            b -> c
f
                                ( if Bool -> (Bitmap -> Bool) -> Maybe Bitmap -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (Bitmap -> Int -> Bool
`bitmapTestBit` Int
i) Maybe Bitmap
bm
                                    then a -> Maybe a
forall a. a -> Maybe a
Just (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
col Int
i)
                                    else b
Maybe a
forall a. Maybe a
Nothing
                                )
        Maybe (b :~: Maybe a)
Nothing -> 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
            Just a :~: b
Refl -> 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
$ case forall a. SBoolI (Unboxable a) => SBool (Unboxable a)
sUnbox @c of
                SBool (Unboxable c)
STrue -> Maybe Bitmap -> Vector c -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
bm ((b -> c) -> Vector b -> Vector c
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map b -> c
f Vector b
Vector a
col)
                SBool (Unboxable c)
SFalse -> case Maybe Bitmap
bm of
                    Maybe Bitmap
Nothing -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector @c (Int -> (Int -> c) -> Vector c
forall a. Int -> (Int -> a) -> Vector a
VB.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
col) (b -> c
f (b -> c) -> (Int -> b) -> Int -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector b -> Int -> b
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector b
Vector a
col))
                    Just Bitmap
_ -> Maybe Bitmap -> Vector c -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
bm (Int -> (Int -> c) -> Vector c
forall a. Int -> (Int -> a) -> Vector a
VB.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
col) (b -> c
f (b -> c) -> (Int -> b) -> Int -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector b -> Int -> b
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector b
Vector a
col))
            Maybe (a :~: b)
Nothing -> forall a b.
(Typeable a, Typeable b) =>
Either DataFrameException Column
throwTypeMismatch @a @b
{-# INLINEABLE mapColumn #-}

-- | Applies a function that returns an unboxed result to an unboxed vector, storing the result in a column.
imapColumn ::
    forall b c.
    (Columnable b, Columnable c) =>
    (Int -> b -> c) -> Column -> Either DataFrameException Column
imapColumn :: forall b c.
(Columnable b, Columnable c) =>
(Int -> b -> c) -> Column -> Either DataFrameException Column
imapColumn Int -> b -> c
f = \case
    BoxedColumn Maybe Bitmap
bm (Vector a
col :: VB.Vector a) -> Maybe Bitmap -> Vector a -> Either DataFrameException Column
forall a.
Columnable a =>
Maybe Bitmap -> Vector a -> Either DataFrameException Column
runBoxed Maybe Bitmap
bm Vector a
col
    UnboxedColumn Maybe Bitmap
bm (Vector a
col :: VU.Vector a) -> Maybe Bitmap -> Vector a -> Either DataFrameException Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Either DataFrameException Column
runUnboxed Maybe Bitmap
bm Vector a
col
    c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
_) -> (Int -> b -> c) -> Column -> Either DataFrameException Column
forall b c.
(Columnable b, Columnable c) =>
(Int -> b -> c) -> Column -> Either DataFrameException Column
imapColumn Int -> b -> c
f (Column -> Column
materializePacked Column
c)
    c :: Column
c@(MergedColumn Column
_ Column
_) -> (Int -> b -> c) -> Column -> Either DataFrameException Column
forall b c.
(Columnable b, Columnable c) =>
(Int -> b -> c) -> Column -> Either DataFrameException Column
imapColumn Int -> b -> c
f (Column -> Column
materializeMerged Column
c)
  where
    runBoxed ::
        forall a.
        (Columnable a) =>
        Maybe Bitmap -> VB.Vector a -> Either DataFrameException Column
    runBoxed :: forall a.
Columnable a =>
Maybe Bitmap -> Vector a -> Either DataFrameException Column
runBoxed Maybe Bitmap
bm Vector a
col = 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
        Just a :~: b
Refl -> 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
$ case forall a. SBoolI (Unboxable a) => SBool (Unboxable a)
sUnbox @c of
            SBool (Unboxable c)
STrue ->
                Maybe Bitmap -> Vector c -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn
                    Maybe Bitmap
bm
                    (Int -> (Int -> c) -> Vector c
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
col) (\Int
i -> Int -> b -> c
f Int
i (Vector b -> Int -> b
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector b
Vector a
col Int
i)))
            SBool (Unboxable c)
SFalse -> Maybe Bitmap -> Vector c -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
bm ((Int -> b -> c) -> Vector b -> Vector c
forall a b. (Int -> a -> b) -> Vector a -> Vector b
VB.imap Int -> b -> c
f Vector b
Vector a
col)
        Maybe (a :~: b)
Nothing -> forall a b.
(Typeable a, Typeable b) =>
Either DataFrameException Column
throwTypeMismatch @a @b

    runUnboxed ::
        forall a.
        (Columnable a, VU.Unbox a) =>
        Maybe Bitmap -> VU.Vector a -> Either DataFrameException Column
    runUnboxed :: forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Either DataFrameException Column
runUnboxed Maybe Bitmap
bm Vector a
col = 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
        Just a :~: b
Refl -> 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
$ case forall a. SBoolI (Unboxable a) => SBool (Unboxable a)
sUnbox @c of
            SBool (Unboxable c)
STrue -> Maybe Bitmap -> Vector c -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
bm ((Int -> b -> c) -> Vector b -> Vector c
forall a b.
(Unbox a, Unbox b) =>
(Int -> a -> b) -> Vector a -> Vector b
VU.imap Int -> b -> c
f Vector b
Vector a
col)
            SBool (Unboxable c)
SFalse -> Maybe Bitmap -> Vector c -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
bm ((Int -> b -> c) -> Vector b -> Vector c
forall a b. (Int -> a -> b) -> Vector a -> Vector b
VB.imap Int -> b -> c
f (Vector b -> Vector b
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector b
Vector a
col))
        Maybe (a :~: b)
Nothing -> forall a b.
(Typeable a, Typeable b) =>
Either DataFrameException Column
throwTypeMismatch @a @b

-- | O(1) Gets the number of elements in the column.
columnLength :: Column -> Int
columnLength :: Column -> Int
columnLength (MergedColumn Column
a Column
b) = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Column -> Int
columnLength Column
a) (Column -> Int
columnLength Column
b)
columnLength (BoxedColumn Maybe Bitmap
_ Vector a
xs) = Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
xs
columnLength (UnboxedColumn Maybe Bitmap
_ Vector a
xs) = Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs
columnLength (PackedText Maybe Bitmap
_ PackedTextData
p) = PackedTextData -> Int
packedLength PackedTextData
p
{-# INLINE columnLength #-}

-- | O(n) Gets the number of non-null elements in the column.
numElements :: Column -> Int
numElements :: Column -> Int
numElements (MergedColumn Column
a Column
b) = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Column -> Int
columnLength Column
a) (Column -> Int
columnLength Column
b)
numElements (BoxedColumn Maybe Bitmap
Nothing Vector a
xs) = Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
xs
numElements (BoxedColumn (Just Bitmap
bm) Vector a
_xs) = (Int -> Word8 -> Int) -> Int -> Bitmap -> Int
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl' (\Int
acc Word8
b -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Word8 -> Int
forall a. Bits a => a -> Int
popCount Word8
b) Int
0 Bitmap
bm
numElements (UnboxedColumn Maybe Bitmap
Nothing Vector a
xs) = Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs
numElements (UnboxedColumn (Just Bitmap
bm) Vector a
_xs) = (Int -> Word8 -> Int) -> Int -> Bitmap -> Int
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl' (\Int
acc Word8
b -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Word8 -> Int
forall a. Bits a => a -> Int
popCount Word8
b) Int
0 Bitmap
bm
numElements (PackedText Maybe Bitmap
Nothing PackedTextData
p) = PackedTextData -> Int
packedLength PackedTextData
p
numElements (PackedText (Just Bitmap
bm) PackedTextData
_p) = (Int -> Word8 -> Int) -> Int -> Bitmap -> Int
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl' (\Int
acc Word8
b -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Word8 -> Int
forall a. Bits a => a -> Int
popCount Word8
b) Int
0 Bitmap
bm
{-# INLINE numElements #-}

-- | O(n) Takes the first n values of a column.
takeColumn :: Int -> Column -> Column
takeColumn :: Int -> Column -> Column
takeColumn Int
n (MergedColumn Column
a Column
b) = Column -> Column -> Column
MergedColumn (Int -> Column -> Column
takeColumn Int
n Column
a) (Int -> Column -> Column
takeColumn Int
n Column
b)
takeColumn Int
n (BoxedColumn Maybe Bitmap
bm Vector a
xs) =
    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
0 Int
n) Maybe Bitmap
bm) (Int -> Vector a -> Vector a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
VG.take Int
n Vector a
xs)
takeColumn Int
n (UnboxedColumn Maybe Bitmap
bm Vector a
xs) =
    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
0 Int
n) Maybe Bitmap
bm) (Int -> Vector a -> Vector a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
VG.take Int
n Vector a
xs)
takeColumn Int
n (PackedText Maybe Bitmap
bm PackedTextData
p) =
    Maybe Bitmap -> PackedTextData -> Column
PackedText ((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
0 Int
n) Maybe Bitmap
bm) (Int -> PackedTextData -> PackedTextData
packedTake Int
n PackedTextData
p)
{-# INLINE takeColumn #-}

-- | O(n) Takes the last n values of a column.
takeLastColumn :: Int -> Column -> Column
takeLastColumn :: Int -> Column -> Column
takeLastColumn Int
n Column
column = Int -> Int -> Column -> Column
sliceColumn (Column -> Int
columnLength Column
column Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) Int
n Column
column
{-# INLINE takeLastColumn #-}

-- | O(n) Takes n values after a given column index.
sliceColumn :: Int -> Int -> Column -> Column
sliceColumn :: Int -> Int -> Column -> Column
sliceColumn Int
start Int
n (MergedColumn Column
a Column
b) =
    Column -> Column -> Column
MergedColumn (Int -> Int -> Column -> Column
sliceColumn Int
start Int
n Column
a) (Int -> Int -> Column -> Column
sliceColumn Int
start Int
n Column
b)
sliceColumn Int
start Int
n (BoxedColumn Maybe Bitmap
bm Vector a
xs) =
    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
start Int
n) Maybe Bitmap
bm) (Int -> Int -> Vector a -> Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
Int -> Int -> v a -> v a
VG.slice Int
start Int
n Vector a
xs)
sliceColumn Int
start Int
n (UnboxedColumn Maybe Bitmap
bm Vector a
xs) =
    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
start Int
n) Maybe Bitmap
bm) (Int -> Int -> Vector a -> Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
Int -> Int -> v a -> v a
VG.slice Int
start Int
n Vector a
xs)
sliceColumn Int
start Int
n c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
_) = Int -> Int -> Column -> Column
sliceColumn Int
start Int
n (Column -> Column
materializePacked Column
c)
{-# INLINE sliceColumn #-}

-- | O(n) Selects the elements at a given set of indices. Does not change the order.
atIndicesStable :: VU.Vector Int -> Column -> Column
atIndicesStable :: Vector Int -> Column -> Column
atIndicesStable Vector Int
indexes (BoxedColumn Maybe Bitmap
bm Vector a
column) =
    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
            ( \Bitmap
bm0 ->
                Bitmap -> Bitmap
buildBitmapFromValid (Bitmap -> Bitmap) -> Bitmap -> Bitmap
forall a b. (a -> b) -> a -> b
$
                    (Int -> Word8) -> Vector Int -> Bitmap
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (\Int
i -> if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm0 Int
i then Word8
1 else Word8
0) Vector Int
indexes
            )
            Maybe Bitmap
bm
        )
        ( Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
VB.generate
            (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
indexes)
            ((Vector a
column Vector a -> Int -> a
forall a. Vector a -> Int -> a
`VB.unsafeIndex`) (Int -> a) -> (Int -> Int) -> Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector Int
indexes Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
`VU.unsafeIndex`))
        )
atIndicesStable Vector Int
indexes (UnboxedColumn Maybe Bitmap
bm Vector a
column) =
    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
            ( \Bitmap
bm0 ->
                Bitmap -> Bitmap
buildBitmapFromValid (Bitmap -> Bitmap) -> Bitmap -> Bitmap
forall a b. (a -> b) -> a -> b
$
                    (Int -> Word8) -> Vector Int -> Bitmap
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (\Int
i -> if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm0 Int
i then Word8
1 else Word8
0) Vector Int
indexes
            )
            Maybe Bitmap
bm
        )
        (Vector a -> Vector Int -> Vector a
forall a. Unbox a => Vector a -> Vector Int -> Vector a
VU.unsafeBackpermute Vector a
column Vector Int
indexes)
atIndicesStable Vector Int
indexes (MergedColumn Column
a Column
b) =
    Column -> Column -> Column
MergedColumn (Vector Int -> Column -> Column
atIndicesStable Vector Int
indexes Column
a) (Vector Int -> Column -> Column
atIndicesStable Vector Int
indexes Column
b)
atIndicesStable Vector Int
indexes (PackedText Maybe Bitmap
bm PackedTextData
p) =
    Maybe Bitmap -> PackedTextData -> Column
PackedText
        ( (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
            ( \Bitmap
bm0 ->
                Bitmap -> Bitmap
buildBitmapFromValid (Bitmap -> Bitmap) -> Bitmap -> Bitmap
forall a b. (a -> b) -> a -> b
$
                    (Int -> Word8) -> Vector Int -> Bitmap
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (\Int
i -> if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm0 Int
i then Word8
1 else Word8
0) Vector Int
indexes
            )
            Maybe Bitmap
bm
        )
        (Vector Int -> PackedTextData -> PackedTextData
packedGather Vector Int
indexes PackedTextData
p)
{-# INLINE atIndicesStable #-}

{- | Like 'atIndicesStable' but treats negative indices as null.
Keeps the index vector fully unboxed (no @VB.Vector (Maybe Int)@).
-}
gatherWithSentinel :: VU.Vector Int -> Column -> Column
gatherWithSentinel :: Vector Int -> Column -> Column
gatherWithSentinel Vector Int
indices c :: Column
c@(MergedColumn Column
_ Column
_) =
    Vector Int -> Column -> Column
gatherWithSentinel Vector Int
indices (Column -> Column
materializeMerged Column
c)
gatherWithSentinel Vector Int
indices Column
col =
    let !n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
indices
        newBm :: Bitmap
newBm = Bitmap -> Bitmap
buildBitmapFromValid (Bitmap -> Bitmap) -> Bitmap -> Bitmap
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Word8) -> Bitmap
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n ((Int -> Word8) -> Bitmap) -> (Int -> Word8) -> Bitmap
forall a b. (a -> b) -> a -> b
$ \Int
i ->
            if Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
indices Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then Word8
0 else Word8
1
     in case Column
col of
            PackedText Maybe Bitmap
srcBm PackedTextData
p ->
                let bm :: Maybe Bitmap
bm = case Maybe Bitmap
srcBm of
                        Maybe Bitmap
Nothing -> Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just Bitmap
newBm
                        Just Bitmap
sb ->
                            Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just
                                ( Bitmap -> Bitmap -> Bitmap
mergeBitmaps
                                    Bitmap
newBm
                                    ( Bitmap -> Bitmap
buildBitmapFromValid (Bitmap -> Bitmap) -> Bitmap -> Bitmap
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Word8) -> Bitmap
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n ((Int -> Word8) -> Bitmap) -> (Int -> Word8) -> Bitmap
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                                        let idx :: Int
idx = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
indices Int
i
                                         in if Int
idx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Bitmap -> Int -> Bool
bitmapTestBit Bitmap
sb Int
idx then Word8
1 else Word8
0
                                    )
                                )
                 in Maybe Bitmap -> PackedTextData -> Column
PackedText Maybe Bitmap
bm (Vector Int -> PackedTextData -> PackedTextData
packedGather Vector Int
indices PackedTextData
p)
            BoxedColumn Maybe Bitmap
srcBm Vector a
v ->
                let dat :: Vector a
dat = Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
VB.generate Int
n ((Int -> a) -> Vector a) -> (Int -> a) -> Vector a
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                        let !idx :: Int
idx = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
indices Int
i
                         in if Int
idx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then Vector a -> Int -> a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector a
v Int
0 else Vector a -> Int -> a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector a
v Int
idx
                    bm :: Maybe Bitmap
bm = case Maybe Bitmap
srcBm of
                        Maybe Bitmap
Nothing -> Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just Bitmap
newBm
                        Just Bitmap
sb ->
                            Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just
                                ( Bitmap -> Bitmap -> Bitmap
mergeBitmaps
                                    Bitmap
newBm
                                    ( Bitmap -> Bitmap
buildBitmapFromValid (Bitmap -> Bitmap) -> Bitmap -> Bitmap
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Word8) -> Bitmap
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n ((Int -> Word8) -> Bitmap) -> (Int -> Word8) -> Bitmap
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                                        let idx :: Int
idx = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
indices Int
i
                                         in if Int
idx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Bitmap -> Int -> Bool
bitmapTestBit Bitmap
sb Int
idx then Word8
1 else Word8
0
                                    )
                                )
                 in Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
bm Vector a
dat
            UnboxedColumn Maybe Bitmap
srcBm Vector a
v ->
                let dat :: Vector a
dat = (forall s. ST s (Vector a)) -> Vector a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector a)) -> Vector a)
-> (forall s. ST s (Vector a)) -> Vector a
forall a b. (a -> b) -> a -> b
$ do
                        MVector s a
mv <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
n
                        Vector Int -> (Int -> Int -> ST s ()) -> ST s ()
forall (m :: * -> *) (v :: * -> *) a b.
(Monad m, Vector v a) =>
v a -> (Int -> a -> m b) -> m ()
VG.iforM_ Vector Int
indices ((Int -> Int -> ST s ()) -> ST s ())
-> (Int -> Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i Int
idx ->
                            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
idx Int -> Int -> Bool
forall a. Ord 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)) 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
mv Int
i (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
idx)
                        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 MVector s a
MVector (PrimState (ST s)) a
mv
                    bm :: Maybe Bitmap
bm = case Maybe Bitmap
srcBm of
                        Maybe Bitmap
Nothing -> Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just Bitmap
newBm
                        Just Bitmap
sb ->
                            Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just
                                ( Bitmap -> Bitmap -> Bitmap
mergeBitmaps
                                    Bitmap
newBm
                                    ( Bitmap -> Bitmap
buildBitmapFromValid (Bitmap -> Bitmap) -> Bitmap -> Bitmap
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Word8) -> Bitmap
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n ((Int -> Word8) -> Bitmap) -> (Int -> Word8) -> Bitmap
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                                        let idx :: Int
idx = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
indices Int
i
                                         in if Int
idx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Bitmap -> Int -> Bool
bitmapTestBit Bitmap
sb Int
idx then Word8
1 else Word8
0
                                    )
                                )
                 in Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
bm Vector a
dat
{-# INLINE gatherWithSentinel #-}

-- | Internal helper to get indices in a boxed vector.
getIndices :: VU.Vector Int -> VB.Vector a -> VB.Vector a
getIndices :: forall a. Vector Int -> Vector a -> Vector a
getIndices Vector Int
indices Vector a
xs = Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
VB.generate (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
indices) (\Int
i -> Vector a
xs Vector a -> Int -> a
forall a. Vector a -> Int -> a
VB.! (Vector Int
indices Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
i))
{-# INLINE getIndices #-}

-- | Internal helper to get indices in an unboxed vector.
getIndicesUnboxed :: (VU.Unbox a) => VU.Vector Int -> VU.Vector a -> VU.Vector a
getIndicesUnboxed :: forall a. Unbox a => Vector Int -> Vector a -> Vector a
getIndicesUnboxed Vector Int
indices Vector a
xs = Int -> (Int -> a) -> Vector a
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
indices) (\Int
i -> Vector a
xs Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.! (Vector Int
indices Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
i))
{-# INLINE getIndicesUnboxed #-}

findIndices ::
    forall a.
    (Columnable a) =>
    (a -> Bool) ->
    Column ->
    Either DataFrameException (VU.Vector Int)
findIndices :: forall a.
Columnable a =>
(a -> Bool) -> Column -> Either DataFrameException (Vector Int)
findIndices a -> Bool
predicate = \case
    BoxedColumn Maybe Bitmap
_ (Vector a
v :: VB.Vector b) -> Vector a
-> (Vector Int -> Vector Int)
-> Either DataFrameException (Vector Int)
forall b (v :: * -> *).
(Typeable b, Vector v b, Vector v Int) =>
v b
-> (v Int -> Vector Int) -> Either DataFrameException (Vector Int)
run Vector a
v Vector Int -> Vector Int
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert
    UnboxedColumn Maybe Bitmap
_ (Vector a
v :: VU.Vector b) -> Vector a
-> (Vector Int -> Vector Int)
-> Either DataFrameException (Vector Int)
forall b (v :: * -> *).
(Typeable b, Vector v b, Vector v Int) =>
v b
-> (v Int -> Vector Int) -> Either DataFrameException (Vector Int)
run Vector a
v Vector Int -> Vector Int
forall a. a -> a
id
    c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
_) -> (a -> Bool) -> Column -> Either DataFrameException (Vector Int)
forall a.
Columnable a =>
(a -> Bool) -> Column -> Either DataFrameException (Vector Int)
findIndices a -> Bool
predicate (Column -> Column
materializePacked Column
c)
    c :: Column
c@(MergedColumn Column
_ Column
_) -> (a -> Bool) -> Column -> Either DataFrameException (Vector Int)
forall a.
Columnable a =>
(a -> Bool) -> Column -> Either DataFrameException (Vector Int)
findIndices a -> Bool
predicate (Column -> Column
materializeMerged Column
c)
  where
    run ::
        forall b v.
        (Typeable b, VG.Vector v b, VG.Vector v Int) =>
        v b ->
        (v Int -> VU.Vector Int) ->
        Either DataFrameException (VU.Vector Int)
    run :: forall b (v :: * -> *).
(Typeable b, Vector v b, Vector v Int) =>
v b
-> (v Int -> Vector Int) -> Either DataFrameException (Vector Int)
run v b
column v Int -> Vector Int
finalize = 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
        Just a :~: b
Refl -> Vector Int -> Either DataFrameException (Vector Int)
forall a b. b -> Either a b
Right (Vector Int -> Either DataFrameException (Vector Int))
-> (v Int -> Vector Int)
-> v Int
-> Either DataFrameException (Vector Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v Int -> Vector Int
finalize (v Int -> Either DataFrameException (Vector Int))
-> v Int -> Either DataFrameException (Vector Int)
forall a b. (a -> b) -> a -> b
$ (a -> Bool) -> v a -> v Int
forall (v :: * -> *) a.
(Vector v a, Vector v Int) =>
(a -> Bool) -> v a -> v Int
VG.findIndices a -> Bool
predicate v a
v b
column
        Maybe (a :~: b)
Nothing ->
            DataFrameException -> Either DataFrameException (Vector Int)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Vector Int))
-> DataFrameException -> Either DataFrameException (Vector Int)
forall a b. (a -> b) -> a -> b
$
                TypeErrorContext a b -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                    MkTypeErrorContext
                        { userType :: Either [Char] (TypeRep a)
userType = TypeRep a -> Either [Char] (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 [Char] (TypeRep b)
expectedType = TypeRep b -> Either [Char] (TypeRep b)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b)
                        , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"findIndices"
                        , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                        }

-- | Fold (right) column with index.
ifoldrColumn ::
    forall a b.
    (Columnable a, Columnable b) =>
    (Int -> a -> b -> b) -> b -> Column -> Either DataFrameException b
ifoldrColumn :: forall a b.
(Columnable a, Columnable b) =>
(Int -> a -> b -> b) -> b -> Column -> Either DataFrameException b
ifoldrColumn Int -> a -> b -> b
f b
acc = \case
    BoxedColumn Maybe Bitmap
_ Vector a
column -> Vector a -> Either DataFrameException b
forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException b
foldrWorker Vector a
column
    UnboxedColumn Maybe Bitmap
_ Vector a
column -> Vector a -> Either DataFrameException b
forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException b
foldrWorker Vector a
column
    c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
_) -> (Int -> a -> b -> b) -> b -> Column -> Either DataFrameException b
forall a b.
(Columnable a, Columnable b) =>
(Int -> a -> b -> b) -> b -> Column -> Either DataFrameException b
ifoldrColumn Int -> a -> b -> b
f b
acc (Column -> Column
materializePacked Column
c)
    c :: Column
c@(MergedColumn Column
_ Column
_) -> (Int -> a -> b -> b) -> b -> Column -> Either DataFrameException b
forall a b.
(Columnable a, Columnable b) =>
(Int -> a -> b -> b) -> b -> Column -> Either DataFrameException b
ifoldrColumn Int -> a -> b -> b
f b
acc (Column -> Column
materializeMerged Column
c)
  where
    foldrWorker ::
        forall c v.
        (Typeable c, VG.Vector v c) =>
        v c ->
        Either DataFrameException b
    foldrWorker :: forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException b
foldrWorker v c
vec = case TypeRep a -> TypeRep c -> Maybe (a :~: c)
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 @c) of
        Just a :~: c
Refl -> b -> Either DataFrameException b
forall a. a -> Either DataFrameException a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> Either DataFrameException b)
-> b -> Either DataFrameException b
forall a b. (a -> b) -> a -> b
$ (Int -> a -> b -> b) -> b -> v a -> b
forall (v :: * -> *) a b.
Vector v a =>
(Int -> a -> b -> b) -> b -> v a -> b
VG.ifoldr Int -> a -> b -> b
f b
acc v a
v c
vec
        Maybe (a :~: c)
Nothing ->
            DataFrameException -> Either DataFrameException b
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException b)
-> DataFrameException -> Either DataFrameException b
forall a b. (a -> b) -> a -> b
$
                TypeErrorContext a c -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                    ( MkTypeErrorContext
                        { userType :: Either [Char] (TypeRep a)
userType = TypeRep a -> Either [Char] (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 [Char] (TypeRep c)
expectedType = TypeRep c -> Either [Char] (TypeRep c)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @c)
                        , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"ifoldrColumn"
                        , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                        }
                    )

foldlColumn ::
    forall a b.
    (Columnable a, Columnable b) =>
    (b -> a -> b) -> b -> Column -> Either DataFrameException b
foldlColumn :: forall a b.
(Columnable a, Columnable b) =>
(b -> a -> b) -> b -> Column -> Either DataFrameException b
foldlColumn b -> a -> b
f b
acc = \case
    BoxedColumn Maybe Bitmap
_ Vector a
column -> Vector a -> Either DataFrameException b
forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException b
foldlWorker Vector a
column
    UnboxedColumn Maybe Bitmap
_ Vector a
column -> Vector a -> Either DataFrameException b
forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException b
foldlWorker Vector a
column
    c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
_) -> (b -> a -> b) -> b -> Column -> Either DataFrameException b
forall a b.
(Columnable a, Columnable b) =>
(b -> a -> b) -> b -> Column -> Either DataFrameException b
foldlColumn b -> a -> b
f b
acc (Column -> Column
materializePacked Column
c)
    c :: Column
c@(MergedColumn Column
_ Column
_) -> (b -> a -> b) -> b -> Column -> Either DataFrameException b
forall a b.
(Columnable a, Columnable b) =>
(b -> a -> b) -> b -> Column -> Either DataFrameException b
foldlColumn b -> a -> b
f b
acc (Column -> Column
materializeMerged Column
c)
  where
    foldlWorker ::
        forall c v.
        (Typeable c, VG.Vector v c) =>
        v c ->
        Either DataFrameException b
    foldlWorker :: forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException b
foldlWorker v c
vec = case TypeRep a -> TypeRep c -> Maybe (a :~: c)
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 @c) of
        Just a :~: c
Refl -> b -> Either DataFrameException b
forall a. a -> Either DataFrameException a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> Either DataFrameException b)
-> b -> Either DataFrameException b
forall a b. (a -> b) -> a -> b
$ (b -> a -> b) -> b -> v a -> b
forall (v :: * -> *) b a.
Vector v b =>
(a -> b -> a) -> a -> v b -> a
VG.foldl' b -> a -> b
f b
acc v a
v c
vec
        Maybe (a :~: c)
Nothing ->
            DataFrameException -> Either DataFrameException b
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException b)
-> DataFrameException -> Either DataFrameException b
forall a b. (a -> b) -> a -> b
$
                TypeErrorContext a c -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                    ( MkTypeErrorContext
                        { userType :: Either [Char] (TypeRep a)
userType = TypeRep a -> Either [Char] (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 [Char] (TypeRep c)
expectedType = TypeRep c -> Either [Char] (TypeRep c)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @c)
                        , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"ifoldrColumn"
                        , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                        }
                    )

foldl1Column ::
    forall a.
    (Columnable a) =>
    (a -> a -> a) -> Column -> Either DataFrameException a
foldl1Column :: forall a.
Columnable a =>
(a -> a -> a) -> Column -> Either DataFrameException a
foldl1Column a -> a -> a
f = \case
    BoxedColumn Maybe Bitmap
_ Vector a
column -> Vector a -> Either DataFrameException a
forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException a
foldl1Worker Vector a
column
    UnboxedColumn Maybe Bitmap
_ Vector a
column -> Vector a -> Either DataFrameException a
forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException a
foldl1Worker Vector a
column
    c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
_) -> (a -> a -> a) -> Column -> Either DataFrameException a
forall a.
Columnable a =>
(a -> a -> a) -> Column -> Either DataFrameException a
foldl1Column a -> a -> a
f (Column -> Column
materializePacked Column
c)
    c :: Column
c@(MergedColumn Column
_ Column
_) -> (a -> a -> a) -> Column -> Either DataFrameException a
forall a.
Columnable a =>
(a -> a -> a) -> Column -> Either DataFrameException a
foldl1Column a -> a -> a
f (Column -> Column
materializeMerged Column
c)
  where
    foldl1Worker ::
        forall c v.
        (Typeable c, VG.Vector v c) =>
        v c ->
        Either DataFrameException a
    foldl1Worker :: forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException a
foldl1Worker v c
vec = case TypeRep a -> TypeRep c -> Maybe (a :~: c)
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 @c) of
        Just a :~: c
Refl -> a -> Either DataFrameException a
forall a. a -> Either DataFrameException a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Either DataFrameException a)
-> a -> Either DataFrameException a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> v a -> a
forall (v :: * -> *) a. Vector v a => (a -> a -> a) -> v a -> a
VG.foldl1' a -> a -> a
f v a
v c
vec
        Maybe (a :~: c)
Nothing ->
            DataFrameException -> Either DataFrameException a
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException a)
-> DataFrameException -> Either DataFrameException a
forall a b. (a -> b) -> a -> b
$
                TypeErrorContext a c -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                    ( MkTypeErrorContext
                        { userType :: Either [Char] (TypeRep a)
userType = TypeRep a -> Either [Char] (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 [Char] (TypeRep c)
expectedType = TypeRep c -> Either [Char] (TypeRep c)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @c)
                        , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"foldl1Column"
                        , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                        }
                    )

{- | O(n) Seedless fold over groups using the first element of each group as seed.
Like 'foldDirectGroups' but for the case where no initial accumulator is available.
-}
foldl1DirectGroups ::
    forall a.
    (Columnable a) =>
    (a -> a -> a) ->
    Column ->
    VU.Vector Int ->
    VU.Vector Int ->
    Either DataFrameException Column
foldl1DirectGroups :: forall a.
Columnable a =>
(a -> a -> a)
-> Column
-> Vector Int
-> Vector Int
-> Either DataFrameException Column
foldl1DirectGroups a -> a -> a
f Column
col Vector Int
valueIndices Vector Int
offsets
    | Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
offsets Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = 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
forall a. Vector a
VB.empty
    | Bool
otherwise = case Column
col of
        UnboxedColumn Maybe Bitmap
_ (Vector a
vec :: VU.Vector d) -> 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 -> Column)
-> Either DataFrameException (Vector a)
-> Either DataFrameException Column
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector a -> Either DataFrameException (Vector a)
forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException (v c)
foldl1Worker Vector a
vec
        BoxedColumn Maybe Bitmap
_ (Vector a
vec :: VB.Vector d) -> Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector a -> Column)
-> Either DataFrameException (Vector a)
-> Either DataFrameException Column
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector a -> Either DataFrameException (Vector a)
forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException (v c)
foldl1Worker Vector a
vec
        PackedText Maybe Bitmap
_ PackedTextData
_ -> (a -> a -> a)
-> Column
-> Vector Int
-> Vector Int
-> Either DataFrameException Column
forall a.
Columnable a =>
(a -> a -> a)
-> Column
-> Vector Int
-> Vector Int
-> Either DataFrameException Column
foldl1DirectGroups a -> a -> a
f (Column -> Column
materializePacked Column
col) Vector Int
valueIndices Vector Int
offsets
        MergedColumn Column
_ Column
_ -> (a -> a -> a)
-> Column
-> Vector Int
-> Vector Int
-> Either DataFrameException Column
forall a.
Columnable a =>
(a -> a -> a)
-> Column
-> Vector Int
-> Vector Int
-> Either DataFrameException Column
foldl1DirectGroups a -> a -> a
f (Column -> Column
materializeMerged Column
col) Vector Int
valueIndices Vector Int
offsets
  where
    foldl1Worker ::
        forall c v.
        (Typeable c, VG.Vector v c) =>
        v c ->
        Either DataFrameException (v c)
    foldl1Worker :: forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException (v c)
foldl1Worker v c
vec = case TypeRep a -> TypeRep c -> Maybe (a :~: c)
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 @c) of
        Just a :~: c
Refl ->
            v c -> Either DataFrameException (v c)
forall a b. b -> Either a b
Right (v c -> Either DataFrameException (v c))
-> v c -> Either DataFrameException (v c)
forall a b. (a -> b) -> a -> b
$
                Int -> (Int -> c) -> v c
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
VG.generate (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
offsets Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> a
Int -> c
foldGroup
          where
            foldGroup :: Int -> a
foldGroup Int
k =
                let !s :: Int
s = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offsets Int
k
                    !e :: Int
e = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offsets (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    !seed :: c
seed = v c -> Int -> c
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex v c
vec (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
valueIndices Int
s)
                 in Int -> Int -> a -> a
go (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
e a
c
seed
            go :: Int -> Int -> a -> a
go !Int
i !Int
e !a
acc
                | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
e = a
acc
                | Bool
otherwise =
                    Int -> Int -> a -> a
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
e (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$!
                        a -> a -> a
f a
acc (v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex v a
v c
vec (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
valueIndices Int
i))
        Maybe (a :~: c)
Nothing ->
            DataFrameException -> Either DataFrameException (v c)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (v c))
-> DataFrameException -> Either DataFrameException (v c)
forall a b. (a -> b) -> a -> b
$
                TypeErrorContext a c -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                    MkTypeErrorContext
                        { userType :: Either [Char] (TypeRep a)
userType = TypeRep a -> Either [Char] (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 [Char] (TypeRep c)
expectedType = TypeRep c -> Either [Char] (TypeRep c)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @c)
                        , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"foldl1DirectGroups"
                        , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                        }
{-# INLINEABLE foldl1DirectGroups #-}

{- | O(n) fold over groups by scanning the column linearly (rowToGroup[i] = group
of row i). Random writes hit the small per-group accumulator array; when @acc@ is
unboxable that array is unboxed, avoiding pointer indirection.
-}
foldLinearGroups ::
    forall b acc.
    (Columnable b, Columnable acc) =>
    (acc -> b -> acc) ->
    acc ->
    Column ->
    VU.Vector Int ->
    Int ->
    Either DataFrameException Column
foldLinearGroups :: forall b acc.
(Columnable b, Columnable acc) =>
(acc -> b -> acc)
-> acc
-> Column
-> Vector Int
-> Int
-> Either DataFrameException Column
foldLinearGroups acc -> b -> acc
f acc
seed Column
col Vector Int
rowToGroup Int
nGroups
    | Int
nGroups Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Column -> Either DataFrameException Column
forall a b. b -> Either a b
Right (forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector @acc Vector acc
forall a. Vector a
VB.empty)
    | Bool
otherwise = case Column
col of
        UnboxedColumn Maybe Bitmap
_ (Vector a
vec :: VU.Vector d) -> Vector a -> Either DataFrameException Column
forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException Column
foldLinearWorker Vector a
vec
        BoxedColumn Maybe Bitmap
_ (Vector a
vec :: VB.Vector d) -> Vector a -> Either DataFrameException Column
forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException Column
foldLinearWorker Vector a
vec
        PackedText Maybe Bitmap
_ PackedTextData
_ ->
            (acc -> b -> acc)
-> acc
-> Column
-> Vector Int
-> Int
-> Either DataFrameException Column
forall b acc.
(Columnable b, Columnable acc) =>
(acc -> b -> acc)
-> acc
-> Column
-> Vector Int
-> Int
-> Either DataFrameException Column
foldLinearGroups acc -> b -> acc
f acc
seed (Column -> Column
materializePacked Column
col) Vector Int
rowToGroup Int
nGroups
        MergedColumn Column
_ Column
_ ->
            (acc -> b -> acc)
-> acc
-> Column
-> Vector Int
-> Int
-> Either DataFrameException Column
forall b acc.
(Columnable b, Columnable acc) =>
(acc -> b -> acc)
-> acc
-> Column
-> Vector Int
-> Int
-> Either DataFrameException Column
foldLinearGroups acc -> b -> acc
f acc
seed (Column -> Column
materializeMerged Column
col) Vector Int
rowToGroup Int
nGroups
  where
    foldLinearWorker ::
        forall c v.
        (Typeable c, VG.Vector v c) =>
        v c ->
        Either DataFrameException Column
    foldLinearWorker :: forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException Column
foldLinearWorker v c
vec = case TypeRep b -> TypeRep c -> Maybe (b :~: c)
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 @c) of
        Just b :~: c
Refl ->
            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
$
                IO Column -> Column
forall a. IO a -> a
unsafePerformIO (IO Column -> Column) -> IO Column -> Column
forall a b. (a -> b) -> a -> b
$
                    ((Int -> IO acc) -> (Int -> acc -> IO ()) -> IO ()) -> IO Column
runWith
                        ( \Int -> IO acc
readAt Int -> acc -> IO ()
writeAt ->
                            v c -> (Int -> c -> IO ()) -> IO ()
forall (m :: * -> *) (v :: * -> *) a b.
(Monad m, Vector v a) =>
v a -> (Int -> a -> m b) -> m ()
VG.iforM_ v c
vec ((Int -> c -> IO ()) -> IO ()) -> (Int -> c -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
row c
x -> do
                                let !k :: Int
k = Vector Int -> Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex Vector Int
rowToGroup Int
row
                                acc
cur <- Int -> IO acc
readAt Int
k
                                Int -> acc -> IO ()
writeAt Int
k (acc -> IO ()) -> acc -> IO ()
forall a b. (a -> b) -> a -> b
$! acc -> b -> acc
f acc
cur b
c
x
                        )
        Maybe (b :~: c)
Nothing ->
            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 b c -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                    MkTypeErrorContext
                        { userType :: Either [Char] (TypeRep b)
userType = TypeRep b -> Either [Char] (TypeRep b)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b)
                        , expectedType :: Either [Char] (TypeRep c)
expectedType = TypeRep c -> Either [Char] (TypeRep c)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @c)
                        , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"foldLinearGroups"
                        , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                        }

    runWith :: ((Int -> IO acc) -> (Int -> acc -> IO ()) -> IO ()) -> IO Column
    runWith :: ((Int -> IO acc) -> (Int -> acc -> IO ()) -> IO ()) -> IO Column
runWith (Int -> IO acc) -> (Int -> acc -> IO ()) -> IO ()
body = case forall a. SBoolI (Unboxable a) => SBool (Unboxable a)
sUnbox @acc of
        SBool (Unboxable acc)
STrue -> do
            MVector RealWorld acc
accs <- Int -> acc -> IO (MVector (PrimState IO) acc)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups acc
seed
            (Int -> IO acc) -> (Int -> acc -> IO ()) -> IO ()
body (MVector (PrimState IO) acc -> Int -> IO acc
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld acc
MVector (PrimState IO) acc
accs) (MVector (PrimState IO) acc -> Int -> acc -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld acc
MVector (PrimState IO) acc
accs)
            Maybe Bitmap -> Vector acc -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector acc -> Column) -> IO (Vector acc) -> IO Column
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) acc -> IO (Vector acc)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld acc
MVector (PrimState IO) acc
accs
        SBool (Unboxable acc)
SFalse -> do
            MVector RealWorld acc
accs <- Int -> acc -> IO (MVector (PrimState IO) acc)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MVector (PrimState m) a)
VBM.replicate Int
nGroups acc
seed
            (Int -> IO acc) -> (Int -> acc -> IO ()) -> IO ()
body (MVector (PrimState IO) acc -> Int -> IO acc
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
VBM.unsafeRead MVector RealWorld acc
MVector (PrimState IO) acc
accs) (MVector (PrimState IO) acc -> Int -> acc -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
VBM.unsafeWrite MVector RealWorld acc
MVector (PrimState IO) acc
accs)
            forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector @acc (Vector acc -> Column) -> IO (Vector acc) -> IO Column
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) acc -> IO (Vector acc)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
VB.unsafeFreeze MVector RealWorld acc
MVector (PrimState IO) acc
accs
    {-# INLINE runWith #-}
{-# INLINEABLE foldLinearGroups #-}

headColumn :: forall a. (Columnable a) => Column -> Either DataFrameException a
headColumn :: forall a. Columnable a => Column -> Either DataFrameException a
headColumn = \case
    BoxedColumn Maybe Bitmap
_ Vector a
col -> Vector a -> Either DataFrameException a
forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException a
headWorker Vector a
col
    UnboxedColumn Maybe Bitmap
_ Vector a
col -> Vector a -> Either DataFrameException a
forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException a
headWorker Vector a
col
    c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
_) -> Column -> Either DataFrameException a
forall a. Columnable a => Column -> Either DataFrameException a
headColumn (Column -> Column
materializePacked Column
c)
    c :: Column
c@(MergedColumn Column
_ Column
_) -> Column -> Either DataFrameException a
forall a. Columnable a => Column -> Either DataFrameException a
headColumn (Column -> Column
mergedHead Column
c)
  where
    headWorker ::
        forall c v.
        (Typeable c, VG.Vector v c) =>
        v c ->
        Either DataFrameException a
    headWorker :: forall c (v :: * -> *).
(Typeable c, Vector v c) =>
v c -> Either DataFrameException a
headWorker v c
vec = case TypeRep a -> TypeRep c -> Maybe (a :~: c)
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 @c) of
        Just a :~: c
Refl ->
            if v c -> Bool
forall (v :: * -> *) a. Vector v a => v a -> Bool
VG.null v c
vec
                then DataFrameException -> Either DataFrameException a
forall a b. a -> Either a b
Left (Text -> DataFrameException
EmptyDataSetException Text
"headColumn")
                else a -> Either DataFrameException a
forall a. a -> Either DataFrameException a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (v a -> a
forall (v :: * -> *) a. Vector v a => v a -> a
VG.head v a
v c
vec)
        Maybe (a :~: c)
Nothing ->
            DataFrameException -> Either DataFrameException a
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException a)
-> DataFrameException -> Either DataFrameException a
forall a b. (a -> b) -> a -> b
$
                TypeErrorContext a c -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                    ( MkTypeErrorContext
                        { userType :: Either [Char] (TypeRep a)
userType = TypeRep a -> Either [Char] (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 [Char] (TypeRep c)
expectedType = TypeRep c -> Either [Char] (TypeRep c)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @c)
                        , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"headColumn"
                        , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                        }
                    )

-- | An internal, column version of zip.
zipColumns :: Column -> Column -> Column
zipColumns :: Column -> Column -> Column
zipColumns l :: Column
l@(MergedColumn Column
_ Column
_) Column
r = Column -> Column -> Column
zipColumns (Column -> Column
materializeMerged Column
l) Column
r
zipColumns Column
l r :: Column
r@(MergedColumn Column
_ Column
_) = Column -> Column -> Column
zipColumns Column
l (Column -> Column
materializeMerged Column
r)
zipColumns l :: Column
l@(PackedText Maybe Bitmap
_ PackedTextData
_) Column
r = Column -> Column -> Column
zipColumns (Column -> Column
materializePacked Column
l) Column
r
zipColumns Column
l r :: Column
r@(PackedText Maybe Bitmap
_ PackedTextData
_) = Column -> Column -> Column
zipColumns Column
l (Column -> Column
materializePacked Column
r)
zipColumns (BoxedColumn Maybe Bitmap
_ Vector a
column) (BoxedColumn Maybe Bitmap
_ Vector a
other) = Maybe Bitmap -> Vector (a, a) -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector a -> Vector a -> Vector (a, a)
forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v (a, b)) =>
v a -> v b -> v (a, b)
VG.zip Vector a
column Vector a
other)
zipColumns (BoxedColumn Maybe Bitmap
_ Vector a
column) (UnboxedColumn Maybe Bitmap
_ Vector a
other) =
    Maybe Bitmap -> Vector (a, a) -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn
        Maybe Bitmap
forall a. Maybe a
Nothing
        ( Int -> (Int -> (a, a)) -> Vector (a, a)
forall a. Int -> (Int -> a) -> Vector a
VB.generate
            (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
column) (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
other))
            (\Int
i -> (Vector a
column Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i, Vector a
other Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i))
        )
zipColumns (UnboxedColumn Maybe Bitmap
_ Vector a
column) (BoxedColumn Maybe Bitmap
_ Vector a
other) =
    Maybe Bitmap -> Vector (a, a) -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn
        Maybe Bitmap
forall a. Maybe a
Nothing
        ( Int -> (Int -> (a, a)) -> Vector (a, a)
forall a. Int -> (Int -> a) -> Vector a
VB.generate
            (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
column) (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
other))
            (\Int
i -> (Vector a
column Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i, Vector a
other Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i))
        )
zipColumns (UnboxedColumn Maybe Bitmap
_ Vector a
column) (UnboxedColumn Maybe Bitmap
_ Vector a
other) = Maybe Bitmap -> Vector (a, a) -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector a -> Vector a -> Vector (a, a)
forall (v :: * -> *) a b.
(Vector v a, Vector v b, Vector v (a, b)) =>
v a -> v b -> v (a, b)
VG.zip Vector a
column Vector a
other)
{-# INLINE zipColumns #-}

{- | Merge two columns using `These`. O(1): the sides are kept in their
native representation and 'These' values materialize on element access.
-}
mergeColumns :: Column -> Column -> Column
mergeColumns :: Column -> Column -> Column
mergeColumns = Column -> Column -> Column
MergedColumn
{-# INLINE mergeColumns #-}

-- | Decode a 'MergedColumn' into the eager @BoxedColumn (These a b)@ form.
materializeMerged :: Column -> Column
materializeMerged :: Column -> Column
materializeMerged (MergedColumn Column
colA Column
colB) =
    Column -> Column -> Column
mergeEager (Column -> Column
materializeMerged Column
colA) (Column -> Column
materializeMerged Column
colB)
materializeMerged Column
c = Column
c

mergedHead :: Column -> Column
mergedHead :: Column -> Column
mergedHead (MergedColumn Column
a Column
b) =
    Column -> Column
materializeMerged (Column -> Column -> Column
MergedColumn (Int -> Column -> Column
takeColumn Int
1 Column
a) (Int -> Column -> Column
takeColumn Int
1 Column
b))
mergedHead Column
c = Column
c

{- | The eager element-wise merge ('These' per row, boxed). Bitmaps are
honored for every representation pair: a null side yields 'This'/'That',
both-null is an error (the join kernels never produce such a row).
-}
mergeEager :: Column -> Column -> Column
mergeEager :: Column -> Column -> Column
mergeEager Column
colA Column
colB = case (Column
colA, Column
colB) of
    (MergedColumn Column
a Column
b, Column
_) -> Column -> Column -> Column
mergeEager (Column -> Column -> Column
mergeEager Column
a Column
b) Column
colB
    (Column
_, MergedColumn Column
a Column
b) -> Column -> Column -> Column
mergeEager Column
colA (Column -> Column -> Column
mergeEager Column
a Column
b)
    (PackedText Maybe Bitmap
_ PackedTextData
_, Column
_) -> Column -> Column -> Column
mergeEager (Column -> Column
materializePacked Column
colA) Column
colB
    (Column
_, PackedText Maybe Bitmap
_ PackedTextData
_) -> Column -> Column -> Column
mergeEager Column
colA (Column -> Column
materializePacked Column
colB)
    (BoxedColumn Maybe Bitmap
bmA Vector a
c1, BoxedColumn Maybe Bitmap
bmB Vector a
c2) ->
        Maybe Bitmap
-> Maybe Bitmap -> Int -> Int -> (Int -> a) -> (Int -> a) -> Column
forall a b.
(Columnable a, Columnable b) =>
Maybe Bitmap
-> Maybe Bitmap -> Int -> Int -> (Int -> a) -> (Int -> b) -> Column
merged Maybe Bitmap
bmA Maybe Bitmap
bmB (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
c1) (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
c2) (Vector a
c1 Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.!) (Vector a
c2 Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.!)
    (BoxedColumn Maybe Bitmap
bmA Vector a
c1, UnboxedColumn Maybe Bitmap
bmB Vector a
c2) ->
        Maybe Bitmap
-> Maybe Bitmap -> Int -> Int -> (Int -> a) -> (Int -> a) -> Column
forall a b.
(Columnable a, Columnable b) =>
Maybe Bitmap
-> Maybe Bitmap -> Int -> Int -> (Int -> a) -> (Int -> b) -> Column
merged Maybe Bitmap
bmA Maybe Bitmap
bmB (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
c1) (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
c2) (Vector a
c1 Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.!) (Vector a
c2 Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.!)
    (UnboxedColumn Maybe Bitmap
bmA Vector a
c1, BoxedColumn Maybe Bitmap
bmB Vector a
c2) ->
        Maybe Bitmap
-> Maybe Bitmap -> Int -> Int -> (Int -> a) -> (Int -> a) -> Column
forall a b.
(Columnable a, Columnable b) =>
Maybe Bitmap
-> Maybe Bitmap -> Int -> Int -> (Int -> a) -> (Int -> b) -> Column
merged Maybe Bitmap
bmA Maybe Bitmap
bmB (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
c1) (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
c2) (Vector a
c1 Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.!) (Vector a
c2 Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.!)
    (UnboxedColumn Maybe Bitmap
bmA Vector a
c1, UnboxedColumn Maybe Bitmap
bmB Vector a
c2) ->
        Maybe Bitmap
-> Maybe Bitmap -> Int -> Int -> (Int -> a) -> (Int -> a) -> Column
forall a b.
(Columnable a, Columnable b) =>
Maybe Bitmap
-> Maybe Bitmap -> Int -> Int -> (Int -> a) -> (Int -> b) -> Column
merged Maybe Bitmap
bmA Maybe Bitmap
bmB (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
c1) (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
c2) (Vector a
c1 Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.!) (Vector a
c2 Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.!)
  where
    merged ::
        (Columnable a, Columnable b) =>
        Maybe Bitmap ->
        Maybe Bitmap ->
        Int ->
        Int ->
        (Int -> a) ->
        (Int -> b) ->
        Column
    merged :: forall a b.
(Columnable a, Columnable b) =>
Maybe Bitmap
-> Maybe Bitmap -> Int -> Int -> (Int -> a) -> (Int -> b) -> Column
merged Maybe Bitmap
bmA Maybe Bitmap
bmB Int
lenA Int
lenB Int -> a
atA Int -> b
atB =
        Maybe Bitmap -> Vector (These a b) -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector (These a b) -> Column) -> Vector (These a b) -> Column
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> These a b) -> Vector (These a b)
forall a. Int -> (Int -> a) -> Vector a
VB.generate (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
lenA Int
lenB) ((Int -> These a b) -> Vector (These a b))
-> (Int -> These a b) -> Vector (These a b)
forall a b. (a -> b) -> a -> b
$ \Int
i ->
            case (Maybe Bitmap -> Int -> Bool
validAt Maybe Bitmap
bmA Int
i, Maybe Bitmap -> Int -> Bool
validAt Maybe Bitmap
bmB Int
i) of
                (Bool
True, Bool
True) -> a -> b -> These a b
forall a b. a -> b -> These a b
These (Int -> a
atA Int
i) (Int -> b
atB Int
i)
                (Bool
True, Bool
False) -> a -> These a b
forall a b. a -> These a b
This (Int -> a
atA Int
i)
                (Bool
False, Bool
True) -> b -> These a b
forall a b. b -> These a b
That (Int -> b
atB Int
i)
                (Bool
False, Bool
False) -> [Char] -> These a b
forall a. HasCallStack => [Char] -> a
error [Char]
"mergeColumns: both null"
    validAt :: Maybe Bitmap -> Int -> Bool
validAt Maybe Bitmap
mbm Int
i = Bool -> (Bitmap -> Bool) -> Maybe Bitmap -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (Bitmap -> Int -> Bool
`bitmapTestBit` Int
i) Maybe Bitmap
mbm
    {-# INLINE validAt #-}

-- | An internal, column version of zipWith.
zipWithColumns ::
    forall a b c.
    (Columnable a, Columnable b, Columnable c) =>
    (a -> b -> c) -> Column -> Column -> Either DataFrameException Column
zipWithColumns :: forall a b c.
(Columnable a, Columnable b, Columnable c) =>
(a -> b -> c)
-> Column -> Column -> Either DataFrameException Column
zipWithColumns a -> b -> c
f (UnboxedColumn Maybe Bitmap
bmL (Vector a
column :: VU.Vector d)) (UnboxedColumn Maybe Bitmap
bmR (Vector a
other :: VU.Vector e)) = case TypeRep a -> TypeRep a -> Maybe (a :~: 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 @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @d) of
    Just a :~: a
Refl -> case TypeRep b -> TypeRep a -> Maybe (b :~: 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 @e) of
        Just b :~: a
Refl
            | Maybe Bitmap -> Bool
forall a. Maybe a -> Bool
isNothing Maybe Bitmap
bmL
            , Maybe Bitmap -> Bool
forall a. Maybe a -> Bool
isNothing Maybe Bitmap
bmR ->
                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
$ case forall a. SBoolI (Unboxable a) => SBool (Unboxable a)
sUnbox @c of
                    SBool (Unboxable c)
STrue -> Maybe Bitmap -> Vector c -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
forall a. Maybe a
Nothing ((a -> b -> c) -> Vector a -> Vector b -> Vector c
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
VU.zipWith a -> b -> c
f Vector a
Vector a
column Vector b
Vector a
other)
                    SBool (Unboxable c)
SFalse -> Vector c -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector (Vector c -> Column) -> Vector c -> Column
forall a b. (a -> b) -> a -> b
$ (a -> b -> c) -> Vector a -> Vector b -> Vector c
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
VB.zipWith a -> b -> c
f (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
Vector a
column) (Vector b -> Vector b
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector b
Vector a
other)
        Maybe (b :~: a)
_ -> (a -> b -> c)
-> Column -> Column -> Either DataFrameException Column
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
(a -> b -> c)
-> Column -> Column -> Either DataFrameException Column
zipWithColumnsGeneral a -> b -> c
f (Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
bmL Vector a
column) (Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
bmR Vector a
other)
    Maybe (a :~: a)
Nothing -> (a -> b -> c)
-> Column -> Column -> Either DataFrameException Column
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
(a -> b -> c)
-> Column -> Column -> Either DataFrameException Column
zipWithColumnsGeneral a -> b -> c
f (Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
bmL Vector a
column) (Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
bmR Vector a
other)
-- TODO: mchavinda - reuse pattern from interpret where we augment the
-- error at the end.
zipWithColumns a -> b -> c
f Column
left Column
right = (a -> b -> c)
-> Column -> Column -> Either DataFrameException Column
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
(a -> b -> c)
-> Column -> Column -> Either DataFrameException Column
zipWithColumnsGeneral a -> b -> c
f Column
left Column
right

zipWithColumnsGeneral ::
    forall a b c.
    (Columnable a, Columnable b, Columnable c) =>
    (a -> b -> c) -> Column -> Column -> Either DataFrameException Column
zipWithColumnsGeneral :: forall a b c.
(Columnable a, Columnable b, Columnable c) =>
(a -> b -> c)
-> Column -> Column -> Either DataFrameException Column
zipWithColumnsGeneral a -> b -> c
f Column
left Column
right = case forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector @a Column
left of
    Left (TypeMismatchException TypeErrorContext a b
context) ->
        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 a b -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException (TypeErrorContext a b
context{callingFunctionName = Just "zipWithColumns"})
    Left DataFrameException
e -> DataFrameException -> Either DataFrameException Column
forall a b. a -> Either a b
Left DataFrameException
e
    Right Vector a
left' -> case forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector @b Column
right of
        Left (TypeMismatchException TypeErrorContext a b
context) ->
            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 a b -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException (TypeErrorContext a b
context{callingFunctionName = Just "zipWithColumns"})
        Left DataFrameException
e -> DataFrameException -> Either DataFrameException Column
forall a b. a -> Either a b
Left DataFrameException
e
        Right Vector b
right' -> 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
$ Vector c -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector (Vector c -> Column) -> Vector c -> Column
forall a b. (a -> b) -> a -> b
$ (a -> b -> c) -> Vector a -> Vector b -> Vector c
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
VB.zipWith a -> b -> c
f Vector a
left' Vector b
right'
{-# INLINE zipWithColumnsGeneral #-}
{-# INLINE zipWithColumns #-}

-- writeColumn and freezeColumn' (CSV-ingest helpers) moved to
-- DataFrame.IO.Internal.MutableColumn so the core column module does not
-- need to depend on DataFrame.Internal.Parsing.

{- | Freeze a mutable column into an @Either Text a@ column: every recorded
null position becomes @Left rawText@ (preserving the original input), every
other position becomes @Right v@. Used by CSV readers under 'EitherRead' mode.
-}
freezeColumnEither :: [(Int, T.Text)] -> MutableColumn -> IO Column
freezeColumnEither :: [(Int, Text)] -> MutableColumn -> IO Column
freezeColumnEither [(Int, Text)]
nulls (MBoxedColumn IOVector a
col) = do
    Vector a
frozen <- MVector (PrimState IO) a -> IO (Vector a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
VB.unsafeFreeze IOVector a
MVector (PrimState IO) a
col
    let nullMap :: [(Int, Text)]
nullMap = [(Int, Text)]
nulls
    Column -> IO Column
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Column -> IO Column) -> Column -> IO Column
forall a b. (a -> b) -> a -> b
$
        Maybe Bitmap -> Vector (Either Text a) -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector (Either Text a) -> Column)
-> Vector (Either Text a) -> Column
forall a b. (a -> b) -> a -> b
$
            (Int -> a -> Either Text a) -> Vector a -> Vector (Either Text a)
forall a b. (Int -> a -> b) -> Vector a -> Vector b
VB.imap
                ( \Int
i a
v -> case Int -> [(Int, Text)] -> Maybe Text
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Int
i [(Int, Text)]
nullMap of
                    Just Text
t -> Text -> Either Text a
forall a b. a -> Either a b
Left Text
t
                    Maybe Text
Nothing -> a -> Either Text a
forall a b. b -> Either a b
Right a
v
                )
                Vector a
frozen
freezeColumnEither [(Int, Text)]
nulls (MUnboxedColumn IOVector a
col) = do
    Vector a
c <- MVector (PrimState IO) a -> IO (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze IOVector a
MVector (PrimState IO) a
col
    let nullMap :: [(Int, Text)]
nullMap = [(Int, Text)]
nulls
    Column -> IO Column
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Column -> IO Column) -> Column -> IO Column
forall a b. (a -> b) -> a -> b
$
        Maybe Bitmap -> Vector (Either Text a) -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector (Either Text a) -> Column)
-> Vector (Either Text a) -> Column
forall a b. (a -> b) -> a -> b
$
            Int -> (Int -> Either Text a) -> Vector (Either Text a)
forall a. Int -> (Int -> a) -> Vector a
VB.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
c) ((Int -> Either Text a) -> Vector (Either Text a))
-> (Int -> Either Text a) -> Vector (Either Text a)
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                case Int -> [(Int, Text)] -> Maybe Text
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Int
i [(Int, Text)]
nullMap of
                    Just Text
t -> Text -> Either Text a
forall a b. a -> Either a b
Left Text
t
                    Maybe Text
Nothing -> a -> Either Text a
forall a b. b -> Either a b
Right (Vector a
c Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
i)
{-# INLINE freezeColumnEither #-}

{- | Promote a non-nullable column to a nullable one (add an all-valid bitmap).
No-op when already nullable.
-}
ensureOptional :: Column -> Column
ensureOptional :: Column -> Column
ensureOptional c :: Column
c@(MergedColumn Column
_ Column
_) = Column -> Column
ensureOptional (Column -> Column
materializeMerged Column
c)
ensureOptional c :: Column
c@(BoxedColumn (Just Bitmap
_) Vector a
_) = Column
c
ensureOptional (BoxedColumn Maybe Bitmap
Nothing Vector a
col) =
    Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn (Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Int -> Bitmap
allValidBitmap (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
col))) Vector a
col
ensureOptional c :: Column
c@(UnboxedColumn (Just Bitmap
_) Vector a
_) = Column
c
ensureOptional (UnboxedColumn Maybe Bitmap
Nothing Vector a
col) =
    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 (Int -> Bitmap
allValidBitmap (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
col))) Vector a
col
ensureOptional c :: Column
c@(PackedText (Just Bitmap
_) PackedTextData
_) = Column
c
ensureOptional (PackedText Maybe Bitmap
Nothing PackedTextData
p) =
    Maybe Bitmap -> PackedTextData -> Column
PackedText (Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Int -> Bitmap
allValidBitmap (PackedTextData -> Int
packedLength PackedTextData
p))) PackedTextData
p

-- | Fills the end of a column, up to n, with null rows. Does nothing if column has length >= n.
expandColumn :: Int -> Column -> Column
expandColumn :: Int -> Column -> Column
expandColumn Int
n c :: Column
c@(MergedColumn Column
a Column
b)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Column -> Int
columnLength Column
a) (Column -> Int
columnLength Column
b) = Column
c
    | Bool
otherwise = Int -> Column -> Column
expandColumn Int
n (Column -> Column
materializeMerged Column
c)
expandColumn Int
n c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
p)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= PackedTextData -> Int
packedLength PackedTextData
p = Column
c
    | Bool
otherwise = Int -> Column -> Column
expandColumn Int
n (Column -> Column
materializePacked Column
c)
expandColumn Int
n column :: Column
column@(BoxedColumn Maybe Bitmap
bm Vector a
col)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col = Column
column
    | Bool
otherwise =
        let extra :: Int
extra = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col
            newBm :: Maybe Bitmap
newBm = case Maybe Bitmap
bm of
                Maybe Bitmap
Nothing -> Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Int -> [Int] -> Bitmap
buildBitmapFromNulls Int
n [Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
                Just Bitmap
b ->
                    Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just
                        (Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col) Bitmap
b Int
extra (Int -> Word8 -> Bitmap
forall a. Unbox a => Int -> a -> Vector a
VU.replicate ((Int
extra Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3) Word8
0))
            newCol :: Vector a
newCol = Vector a
col Vector a -> Vector a -> Vector a
forall a. Semigroup a => a -> a -> a
<> Int -> a -> Vector a
forall a. Int -> a -> Vector a
VB.replicate Int
extra ([Char] -> a
forall a. [Char] -> a
errorWithoutStackTrace [Char]
"expandColumn: null slot")
         in Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
newBm Vector a
newCol
expandColumn Int
n column :: Column
column@(UnboxedColumn Maybe Bitmap
bm Vector a
col)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col = Column
column
    | Bool
otherwise =
        let extra :: Int
extra = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col
            newBm :: Maybe Bitmap
newBm = case Maybe Bitmap
bm of
                Maybe Bitmap
Nothing -> Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Int -> [Int] -> Bitmap
buildBitmapFromNulls Int
n [Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
                Just Bitmap
b ->
                    Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just
                        (Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat (Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col) Bitmap
b Int
extra (Int -> Word8 -> Bitmap
forall a. Unbox a => Int -> a -> Vector a
VU.replicate ((Int
extra Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3) Word8
0))
            newCol :: Vector a
newCol = (forall s. ST s (Vector a)) -> Vector a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector a)) -> Vector a)
-> (forall s. ST s (Vector a)) -> Vector a
forall a b. (a -> b) -> a -> b
$ do
                MVector s a
mv <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
n
                (Int -> a -> ST s ()) -> Vector a -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
(Int -> a -> m b) -> Vector a -> m ()
VU.imapM_ (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
mv) Vector a
col
                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 MVector s a
MVector (PrimState (ST s)) a
mv
         in Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
newBm Vector a
newCol

-- | Fills the beginning of a column, up to n, with null rows. Does nothing if column has length >= n.
leftExpandColumn :: Int -> Column -> Column
leftExpandColumn :: Int -> Column -> Column
leftExpandColumn Int
n c :: Column
c@(MergedColumn Column
a Column
b)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Column -> Int
columnLength Column
a) (Column -> Int
columnLength Column
b) = Column
c
    | Bool
otherwise = Int -> Column -> Column
leftExpandColumn Int
n (Column -> Column
materializeMerged Column
c)
leftExpandColumn Int
n c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
p)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= PackedTextData -> Int
packedLength PackedTextData
p = Column
c
    | Bool
otherwise = Int -> Column -> Column
leftExpandColumn Int
n (Column -> Column
materializePacked Column
c)
leftExpandColumn Int
n column :: Column
column@(BoxedColumn Maybe Bitmap
bm Vector a
col)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col = Column
column
    | Bool
otherwise =
        let extra :: Int
extra = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col
            origLen :: Int
origLen = Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col
            newBm :: Maybe Bitmap
newBm = case Maybe Bitmap
bm of
                Maybe Bitmap
Nothing -> Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Int -> [Int] -> Bitmap
buildBitmapFromNulls Int
n [Int
0 .. Int
extra Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
                Just Bitmap
b ->
                    let nullPart :: Bitmap
nullPart = Int -> Word8 -> Bitmap
forall a. Unbox a => Int -> a -> Vector a
VU.replicate ((Int
extra Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3) Word8
0
                     in Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat Int
extra Bitmap
nullPart Int
origLen Bitmap
b)
            newCol :: Vector a
newCol =
                Int -> a -> Vector a
forall a. Int -> a -> Vector a
VB.replicate Int
extra ([Char] -> a
forall a. [Char] -> a
errorWithoutStackTrace [Char]
"leftExpandColumn: null slot") Vector a -> Vector a -> Vector a
forall a. Semigroup a => a -> a -> a
<> Vector a
col
         in Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
newBm Vector a
newCol
leftExpandColumn Int
n column :: Column
column@(UnboxedColumn Maybe Bitmap
bm Vector a
col)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col = Column
column
    | Bool
otherwise =
        let extra :: Int
extra = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col
            origLen :: Int
origLen = Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector a
col
            newBm :: Maybe Bitmap
newBm = case Maybe Bitmap
bm of
                Maybe Bitmap
Nothing -> Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Int -> [Int] -> Bitmap
buildBitmapFromNulls Int
n [Int
0 .. Int
extra Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
                Just Bitmap
b ->
                    let nullPart :: Bitmap
nullPart = Int -> Word8 -> Bitmap
forall a. Unbox a => Int -> a -> Vector a
VU.replicate ((Int
extra Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3) Word8
0
                     in Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat Int
extra Bitmap
nullPart Int
origLen Bitmap
b)
            newCol :: Vector a
newCol = (forall s. ST s (Vector a)) -> Vector a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector a)) -> Vector a)
-> (forall s. ST s (Vector a)) -> Vector a
forall a b. (a -> b) -> a -> b
$ do
                MVector s a
mv <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
n
                (Int -> a -> ST s ()) -> Vector a -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
(Int -> a -> m b) -> Vector a -> m ()
VU.imapM_ (\Int
i a
x -> 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
mv (Int
extra Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) a
x) Vector a
col
                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 MVector s a
MVector (PrimState (ST s)) a
mv
         in Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
newBm Vector a
newCol

{- | Concatenates two columns.
Returns Nothing if the columns are of different types.
-}
concatColumns :: Column -> Column -> Either DataFrameException Column
concatColumns :: Column -> Column -> Either DataFrameException Column
concatColumns Column
left Column
right = case (Column
left, Column
right) of
    (MergedColumn Column
_ Column
_, Column
_) -> Column -> Column -> Either DataFrameException Column
concatColumns (Column -> Column
materializeMerged Column
left) Column
right
    (Column
_, MergedColumn Column
_ Column
_) -> Column -> Column -> Either DataFrameException Column
concatColumns Column
left (Column -> Column
materializeMerged Column
right)
    (PackedText Maybe Bitmap
_ PackedTextData
_, Column
_) -> Column -> Column -> Either DataFrameException Column
concatColumns (Column -> Column
materializePacked Column
left) Column
right
    (Column
_, PackedText Maybe Bitmap
_ PackedTextData
_) -> Column -> Column -> Either DataFrameException Column
concatColumns Column
left (Column -> Column
materializePacked Column
right)
    (BoxedColumn Maybe Bitmap
bmL Vector a
l, BoxedColumn Maybe Bitmap
bmR Vector a
r) -> case TypeRep (Vector a)
-> TypeRep (Vector a) -> Maybe (Vector a :~: Vector 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 (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
l) (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
r) of
        Just Vector a :~: Vector a
Refl ->
            let newBm :: Maybe Bitmap
newBm = case (Maybe Bitmap
bmL, Maybe Bitmap
bmR) of
                    (Maybe Bitmap
Nothing, Maybe Bitmap
Nothing) -> Maybe Bitmap
forall a. Maybe a
Nothing
                    (Just Bitmap
bl, Maybe Bitmap
Nothing) ->
                        Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just
                            (Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
l) Bitmap
bl (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
r) (Int -> Bitmap
allValidBitmap (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
r)))
                    (Maybe Bitmap
Nothing, Just Bitmap
br) ->
                        Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just
                            (Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
l) (Int -> Bitmap
allValidBitmap (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
l)) (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
r) Bitmap
br)
                    (Just Bitmap
bl, Just Bitmap
br) -> Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
l) Bitmap
bl (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
r) Bitmap
br)
             in Column -> Either DataFrameException Column
forall a. a -> Either DataFrameException a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
newBm (Vector a
l Vector a -> Vector a -> Vector a
forall a. Semigroup a => a -> a -> a
<> Vector a
Vector a
r))
        Maybe (Vector a :~: Vector a)
Nothing -> DataFrameException -> Either DataFrameException Column
forall a b. a -> Either a b
Left (TypeRep (Vector a) -> TypeRep (Vector a) -> DataFrameException
forall x y. TypeRep x -> TypeRep y -> DataFrameException
mismatchErr (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
r) (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
l))
    (UnboxedColumn Maybe Bitmap
bmL Vector a
l, UnboxedColumn Maybe Bitmap
bmR Vector a
r) -> case TypeRep (Vector a)
-> TypeRep (Vector a) -> Maybe (Vector a :~: Vector 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 (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
l) (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
r) of
        Just Vector a :~: Vector a
Refl ->
            let newBm :: Maybe Bitmap
newBm = case (Maybe Bitmap
bmL, Maybe Bitmap
bmR) of
                    (Maybe Bitmap
Nothing, Maybe Bitmap
Nothing) -> Maybe Bitmap
forall a. Maybe a
Nothing
                    (Just Bitmap
bl, Maybe Bitmap
Nothing) ->
                        Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just
                            (Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
l) Bitmap
bl (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
r) (Int -> Bitmap
allValidBitmap (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
r)))
                    (Maybe Bitmap
Nothing, Just Bitmap
br) ->
                        Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just
                            (Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
l) (Int -> Bitmap
allValidBitmap (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
l)) (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
r) Bitmap
br)
                    (Just Bitmap
bl, Just Bitmap
br) -> Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
l) Bitmap
bl (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
r) Bitmap
br)
             in Column -> Either DataFrameException Column
forall a. a -> Either DataFrameException a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
newBm (Vector a
l Vector a -> Vector a -> Vector a
forall a. Semigroup a => a -> a -> a
<> Vector a
Vector a
r))
        Maybe (Vector a :~: Vector a)
Nothing -> DataFrameException -> Either DataFrameException Column
forall a b. a -> Either a b
Left (TypeRep (Vector a) -> TypeRep (Vector a) -> DataFrameException
forall x y. TypeRep x -> TypeRep y -> DataFrameException
mismatchErr (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
r) (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
l))
    (Column, Column)
_ -> DataFrameException -> Either DataFrameException Column
forall a b. a -> Either a b
Left (TypeRep Column -> TypeRep Column -> DataFrameException
forall x y. TypeRep x -> TypeRep y -> DataFrameException
mismatchErr (Column -> TypeRep Column
forall a. Typeable a => a -> TypeRep a
typeOf Column
right) (Column -> TypeRep Column
forall a. Typeable a => a -> TypeRep a
typeOf Column
left))
  where
    mismatchErr ::
        forall (x :: Type) (y :: Type). TypeRep x -> TypeRep y -> DataFrameException
    mismatchErr :: forall x y. TypeRep x -> TypeRep y -> DataFrameException
mismatchErr TypeRep x
ta TypeRep y
tb =
        TypeRep x
-> (Typeable x => DataFrameException) -> DataFrameException
forall k (a :: k) r. TypeRep a -> (Typeable a => r) -> r
withTypeable TypeRep x
ta ((Typeable x => DataFrameException) -> DataFrameException)
-> (Typeable x => DataFrameException) -> DataFrameException
forall a b. (a -> b) -> a -> b
$
            TypeRep y
-> (Typeable y => DataFrameException) -> DataFrameException
forall k (a :: k) r. TypeRep a -> (Typeable a => r) -> r
withTypeable TypeRep y
tb ((Typeable y => DataFrameException) -> DataFrameException)
-> (Typeable y => DataFrameException) -> DataFrameException
forall a b. (a -> b) -> a -> b
$
                TypeErrorContext x y -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                    ( MkTypeErrorContext
                        { userType :: Either [Char] (TypeRep x)
userType = TypeRep x -> Either [Char] (TypeRep x)
forall a b. b -> Either a b
Right TypeRep x
ta
                        , expectedType :: Either [Char] (TypeRep y)
expectedType = TypeRep y -> Either [Char] (TypeRep y)
forall a b. b -> Either a b
Right TypeRep y
tb
                        , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"concatColumns"
                        , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                        }
                    )

{- | Like 'concatColumns' but also combines columns of different types by wrapping
values in 'Either' (e.g. @[1,2]@ and @["a","b"]@ become
@[Left 1, Left 2, Right "a", Right "b"]@).
-}

{- | O(n) Concatenate a list of same-type columns in a single allocation.
All columns must have the same constructor and element type (as they will
within a single Parquet column). Calls 'error' on mismatch.
-}
concatManyColumns :: [Column] -> Column
concatManyColumns :: [Column] -> Column
concatManyColumns [] = [Maybe Int] -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList ([] :: [Maybe Int])
concatManyColumns [Column
c] = Column
c
concatManyColumns [Column]
all'
    | (Column -> Bool) -> [Column] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Column -> Bool
isMergedColumn [Column]
all' =
        [Column] -> Column
concatManyColumns ((Column -> Column) -> [Column] -> [Column]
forall a b. (a -> b) -> [a] -> [b]
map Column -> Column
materializeMerged [Column]
all')
    | (Column -> Bool) -> [Column] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Column -> Bool
isPackedText [Column]
all' =
        [Column] -> Column
concatManyColumns ((Column -> Column) -> [Column] -> [Column]
forall a b. (a -> b) -> [a] -> [b]
map Column -> Column
materializePacked [Column]
all')
concatManyColumns (Column
c0 : [Column]
cs) = case Column
c0 of
    BoxedColumn Maybe Bitmap
bm0 Vector a
v0 ->
        let getCol :: Column -> (Maybe Bitmap, Vector a)
getCol (BoxedColumn Maybe Bitmap
bm Vector a
v) = case TypeRep (Vector a)
-> TypeRep (Vector a) -> Maybe (Vector a :~: Vector 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 (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
v0) (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
v) of
                Just Vector a :~: Vector a
Refl -> (Maybe Bitmap
bm, Vector a
v)
                Maybe (Vector a :~: Vector a)
Nothing -> [Char] -> (Maybe Bitmap, Vector a)
forall a. HasCallStack => [Char] -> a
error [Char]
"concatManyColumns: BoxedColumn type mismatch"
            getCol Column
_ = [Char] -> (Maybe Bitmap, Vector a)
forall a. HasCallStack => [Char] -> a
error [Char]
"concatManyColumns: column constructor mismatch"
            rest :: [(Maybe Bitmap, Vector a)]
rest = (Column -> (Maybe Bitmap, Vector a))
-> [Column] -> [(Maybe Bitmap, Vector a)]
forall a b. (a -> b) -> [a] -> [b]
map Column -> (Maybe Bitmap, Vector a)
getCol [Column]
cs
            allVecs :: [Vector a]
allVecs = Vector a
v0 Vector a -> [Vector a] -> [Vector a]
forall a. a -> [a] -> [a]
: ((Maybe Bitmap, Vector a) -> Vector a)
-> [(Maybe Bitmap, Vector a)] -> [Vector a]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe Bitmap, Vector a) -> Vector a
forall a b. (a, b) -> b
snd [(Maybe Bitmap, Vector a)]
rest
            allBms :: [Maybe Bitmap]
allBms = Maybe Bitmap
bm0 Maybe Bitmap -> [Maybe Bitmap] -> [Maybe Bitmap]
forall a. a -> [a] -> [a]
: ((Maybe Bitmap, Vector a) -> Maybe Bitmap)
-> [(Maybe Bitmap, Vector a)] -> [Maybe Bitmap]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe Bitmap, Vector a) -> Maybe Bitmap
forall a b. (a, b) -> a
fst [(Maybe Bitmap, Vector a)]
rest
            newBm :: Maybe Bitmap
newBm
                | (Maybe Bitmap -> Bool) -> [Maybe Bitmap] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Maybe Bitmap -> Bool
forall a. Maybe a -> Bool
isNothing [Maybe Bitmap]
allBms = Maybe Bitmap
forall a. Maybe a
Nothing
                | Bool
otherwise =
                    let pairs :: [(Vector a, Maybe Bitmap)]
pairs = [Vector a] -> [Maybe Bitmap] -> [(Vector a, Maybe Bitmap)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Vector a]
allVecs [Maybe Bitmap]
allBms
                        expandedBms :: [Bitmap]
expandedBms = ((Vector a, Maybe Bitmap) -> Bitmap)
-> [(Vector a, Maybe Bitmap)] -> [Bitmap]
forall a b. (a -> b) -> [a] -> [b]
map (\(Vector a
v, Maybe Bitmap
mb) -> Bitmap -> Maybe Bitmap -> Bitmap
forall a. a -> Maybe a -> a
fromMaybe (Int -> Bitmap
allValidBitmap (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
v)) Maybe Bitmap
mb) [(Vector a, Maybe Bitmap)]
pairs
                        go :: Bitmap -> Int -> Bitmap -> Int -> Bitmap
go Bitmap
b1 Int
n1 Bitmap
b2 Int
n2 = Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat Int
n1 Bitmap
b1 Int
n2 Bitmap
b2
                        concatBms :: [(Bitmap, Vector a)] -> Bitmap
concatBms [] = Bitmap
forall a. Unbox a => Vector a
VU.empty
                        concatBms [(Bitmap
b, Vector a
_v)] = Bitmap
b
                        concatBms ((Bitmap
b1, Vector a
v1) : (Bitmap
b2, Vector a
v2) : [(Bitmap, Vector a)]
rest') =
                            let merged :: Bitmap
merged = Bitmap -> Int -> Bitmap -> Int -> Bitmap
go Bitmap
b1 (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
v1) Bitmap
b2 (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
v2)
                             in [(Bitmap, Vector a)] -> Bitmap
concatBms ((Bitmap
merged, Vector a
v1 Vector a -> Vector a -> Vector a
forall a. Semigroup a => a -> a -> a
<> Vector a
v2) (Bitmap, Vector a) -> [(Bitmap, Vector a)] -> [(Bitmap, Vector a)]
forall a. a -> [a] -> [a]
: [(Bitmap, Vector a)]
rest')
                     in Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Bitmap -> Maybe Bitmap) -> Bitmap -> Maybe Bitmap
forall a b. (a -> b) -> a -> b
$ [(Bitmap, Vector a)] -> Bitmap
forall {a}. [(Bitmap, Vector a)] -> Bitmap
concatBms ([Bitmap] -> [Vector a] -> [(Bitmap, Vector a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Bitmap]
expandedBms [Vector a]
allVecs)
         in Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
newBm ([Vector a] -> Vector a
forall a. [Vector a] -> Vector a
VB.concat [Vector a]
allVecs)
    UnboxedColumn Maybe Bitmap
bm0 Vector a
v0 ->
        let getCol :: Column -> (Maybe Bitmap, Vector a)
getCol (UnboxedColumn Maybe Bitmap
bm Vector a
v) = case TypeRep (Vector a)
-> TypeRep (Vector a) -> Maybe (Vector a :~: Vector 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 (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
v0) (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
v) of
                Just Vector a :~: Vector a
Refl -> (Maybe Bitmap
bm, Vector a
v)
                Maybe (Vector a :~: Vector a)
Nothing -> [Char] -> (Maybe Bitmap, Vector a)
forall a. HasCallStack => [Char] -> a
error [Char]
"concatManyColumns: UnboxedColumn type mismatch"
            getCol Column
_ = [Char] -> (Maybe Bitmap, Vector a)
forall a. HasCallStack => [Char] -> a
error [Char]
"concatManyColumns: column constructor mismatch"
            rest :: [(Maybe Bitmap, Vector a)]
rest = (Column -> (Maybe Bitmap, Vector a))
-> [Column] -> [(Maybe Bitmap, Vector a)]
forall a b. (a -> b) -> [a] -> [b]
map Column -> (Maybe Bitmap, Vector a)
getCol [Column]
cs
            allVecs :: [Vector a]
allVecs = Vector a
v0 Vector a -> [Vector a] -> [Vector a]
forall a. a -> [a] -> [a]
: ((Maybe Bitmap, Vector a) -> Vector a)
-> [(Maybe Bitmap, Vector a)] -> [Vector a]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe Bitmap, Vector a) -> Vector a
forall a b. (a, b) -> b
snd [(Maybe Bitmap, Vector a)]
rest
            allBms :: [Maybe Bitmap]
allBms = Maybe Bitmap
bm0 Maybe Bitmap -> [Maybe Bitmap] -> [Maybe Bitmap]
forall a. a -> [a] -> [a]
: ((Maybe Bitmap, Vector a) -> Maybe Bitmap)
-> [(Maybe Bitmap, Vector a)] -> [Maybe Bitmap]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe Bitmap, Vector a) -> Maybe Bitmap
forall a b. (a, b) -> a
fst [(Maybe Bitmap, Vector a)]
rest
            newBm :: Maybe Bitmap
newBm
                | (Maybe Bitmap -> Bool) -> [Maybe Bitmap] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Maybe Bitmap -> Bool
forall a. Maybe a -> Bool
isNothing [Maybe Bitmap]
allBms = Maybe Bitmap
forall a. Maybe a
Nothing
                | Bool
otherwise =
                    let pairs :: [(Vector a, Maybe Bitmap)]
pairs = [Vector a] -> [Maybe Bitmap] -> [(Vector a, Maybe Bitmap)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Vector a]
allVecs [Maybe Bitmap]
allBms
                        expandedBms :: [Bitmap]
expandedBms = ((Vector a, Maybe Bitmap) -> Bitmap)
-> [(Vector a, Maybe Bitmap)] -> [Bitmap]
forall a b. (a -> b) -> [a] -> [b]
map (\(Vector a
v, Maybe Bitmap
mb) -> Bitmap -> Maybe Bitmap -> Bitmap
forall a. a -> Maybe a -> a
fromMaybe (Int -> Bitmap
allValidBitmap (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
v)) Maybe Bitmap
mb) [(Vector a, Maybe Bitmap)]
pairs
                        go :: Bitmap -> Int -> Bitmap -> Int -> Bitmap
go Bitmap
b1 Int
n1 Bitmap
b2 Int
n2 = Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat Int
n1 Bitmap
b1 Int
n2 Bitmap
b2
                        concatBms :: [(Bitmap, Vector a)] -> Bitmap
concatBms [] = Bitmap
forall a. Unbox a => Vector a
VU.empty
                        concatBms [(Bitmap
b, Vector a
_)] = Bitmap
b
                        concatBms ((Bitmap
b1, Vector a
v1) : (Bitmap
b2, Vector a
v2) : [(Bitmap, Vector a)]
rest') =
                            let merged :: Bitmap
merged = Bitmap -> Int -> Bitmap -> Int -> Bitmap
go Bitmap
b1 (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
v1) Bitmap
b2 (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
v2)
                             in [(Bitmap, Vector a)] -> Bitmap
concatBms ((Bitmap
merged, Vector a
v1 Vector a -> Vector a -> Vector a
forall a. Semigroup a => a -> a -> a
<> Vector a
v2) (Bitmap, Vector a) -> [(Bitmap, Vector a)] -> [(Bitmap, Vector a)]
forall a. a -> [a] -> [a]
: [(Bitmap, Vector a)]
rest')
                     in Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Bitmap -> Maybe Bitmap) -> Bitmap -> Maybe Bitmap
forall a b. (a -> b) -> a -> b
$ [(Bitmap, Vector a)] -> Bitmap
forall {a}. Unbox a => [(Bitmap, Vector a)] -> Bitmap
concatBms ([Bitmap] -> [Vector a] -> [(Bitmap, Vector a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Bitmap]
expandedBms [Vector a]
allVecs)
         in Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
newBm ([Vector a] -> Vector a
forall a. Unbox a => [Vector a] -> Vector a
VU.concat [Vector a]
allVecs)
    PackedText Maybe Bitmap
_ PackedTextData
_ -> [Column] -> Column
concatManyColumns ((Column -> Column) -> [Column] -> [Column]
forall a b. (a -> b) -> [a] -> [b]
map Column -> Column
materializePacked (Column
c0 Column -> [Column] -> [Column]
forall a. a -> [a] -> [a]
: [Column]
cs))
    MergedColumn Column
_ Column
_ -> [Column] -> Column
concatManyColumns ((Column -> Column) -> [Column] -> [Column]
forall a b. (a -> b) -> [a] -> [b]
map Column -> Column
materializeMerged (Column
c0 Column -> [Column] -> [Column]
forall a. a -> [a] -> [a]
: [Column]
cs))

concatColumnsEither :: Column -> Column -> Column
concatColumnsEither :: Column -> Column -> Column
concatColumnsEither l :: Column
l@(MergedColumn Column
_ Column
_) Column
r =
    Column -> Column -> Column
concatColumnsEither (Column -> Column
materializeMerged Column
l) Column
r
concatColumnsEither Column
l r :: Column
r@(MergedColumn Column
_ Column
_) =
    Column -> Column -> Column
concatColumnsEither Column
l (Column -> Column
materializeMerged Column
r)
concatColumnsEither l :: Column
l@(PackedText Maybe Bitmap
_ PackedTextData
_) Column
r = Column -> Column -> Column
concatColumnsEither (Column -> Column
materializePacked Column
l) Column
r
concatColumnsEither Column
l r :: Column
r@(PackedText Maybe Bitmap
_ PackedTextData
_) = Column -> Column -> Column
concatColumnsEither Column
l (Column -> Column
materializePacked Column
r)
concatColumnsEither (BoxedColumn Maybe Bitmap
bmL Vector a
left) (BoxedColumn Maybe Bitmap
bmR Vector a
right) = case TypeRep (Vector a)
-> TypeRep (Vector a) -> Maybe (Vector a :~: Vector 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 (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
left) (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
right) of
    Maybe (Vector a :~: Vector a)
Nothing ->
        Maybe Bitmap -> Vector (Either a a) -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector (Either a a) -> Column) -> Vector (Either a a) -> Column
forall a b. (a -> b) -> a -> b
$ (a -> Either a a) -> Vector a -> Vector (Either a a)
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Either a a
forall a b. a -> Either a b
Left Vector a
left Vector (Either a a) -> Vector (Either a a) -> Vector (Either a a)
forall a. Semigroup a => a -> a -> a
<> (a -> Either a a) -> Vector a -> Vector (Either a a)
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Either a a
forall a b. b -> Either a b
Right Vector a
right
    Just Vector a :~: Vector a
Refl ->
        let newBm :: Maybe Bitmap
newBm = case (Maybe Bitmap
bmL, Maybe Bitmap
bmR) of
                (Maybe Bitmap
Nothing, Maybe Bitmap
Nothing) -> Maybe Bitmap
forall a. Maybe a
Nothing
                (Just Bitmap
bl, Maybe Bitmap
Nothing) ->
                    Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just
                        ( Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat
                            (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
left)
                            Bitmap
bl
                            (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
right)
                            (Int -> Bitmap
allValidBitmap (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
right))
                        )
                (Maybe Bitmap
Nothing, Just Bitmap
br) ->
                    Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just
                        ( Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat
                            (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
left)
                            (Int -> Bitmap
allValidBitmap (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
left))
                            (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
right)
                            Bitmap
br
                        )
                (Just Bitmap
bl, Just Bitmap
br) -> Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
left) Bitmap
bl (Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
right) Bitmap
br)
         in Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
newBm (Vector a -> Column) -> Vector a -> Column
forall a b. (a -> b) -> a -> b
$ Vector a
left Vector a -> Vector a -> Vector a
forall a. Semigroup a => a -> a -> a
<> Vector a
Vector a
right
concatColumnsEither (UnboxedColumn Maybe Bitmap
bmL Vector a
left) (UnboxedColumn Maybe Bitmap
bmR Vector a
right) = case TypeRep (Vector a)
-> TypeRep (Vector a) -> Maybe (Vector a :~: Vector 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 (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
left) (Vector a -> TypeRep (Vector a)
forall a. Typeable a => a -> TypeRep a
typeOf Vector a
right) of
    Maybe (Vector a :~: Vector a)
Nothing ->
        Maybe Bitmap -> Vector (Either a a) -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector (Either a a) -> Column) -> Vector (Either a a) -> Column
forall a b. (a -> b) -> a -> b
$
            (a -> Either a a) -> Vector a -> Vector (Either a a)
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Either a a
forall a b. a -> Either a b
Left (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
left) Vector (Either a a) -> Vector (Either a a) -> Vector (Either a a)
forall a. Semigroup a => a -> a -> a
<> (a -> Either a a) -> Vector a -> Vector (Either a a)
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Either a a
forall a b. b -> Either a b
Right (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
right)
    Just Vector a :~: Vector a
Refl ->
        let newBm :: Maybe Bitmap
newBm = case (Maybe Bitmap
bmL, Maybe Bitmap
bmR) of
                (Maybe Bitmap
Nothing, Maybe Bitmap
Nothing) -> Maybe Bitmap
forall a. Maybe a
Nothing
                (Just Bitmap
bl, Maybe Bitmap
Nothing) ->
                    Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just
                        ( Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat
                            (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
left)
                            Bitmap
bl
                            (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
right)
                            (Int -> Bitmap
allValidBitmap (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
right))
                        )
                (Maybe Bitmap
Nothing, Just Bitmap
br) ->
                    Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just
                        ( Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat
                            (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
left)
                            (Int -> Bitmap
allValidBitmap (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
left))
                            (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
right)
                            Bitmap
br
                        )
                (Just Bitmap
bl, Just Bitmap
br) -> Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Int -> Bitmap -> Int -> Bitmap -> Bitmap
bitmapConcat (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
left) Bitmap
bl (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
right) Bitmap
br)
         in Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
newBm (Vector a -> Column) -> Vector a -> Column
forall a b. (a -> b) -> a -> b
$ Vector a
left Vector a -> Vector a -> Vector a
forall a. Semigroup a => a -> a -> a
<> Vector a
Vector a
right
concatColumnsEither (BoxedColumn Maybe Bitmap
_ Vector a
left) (UnboxedColumn Maybe Bitmap
_ Vector a
right) =
    Maybe Bitmap -> Vector (Either a a) -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector (Either a a) -> Column) -> Vector (Either a a) -> Column
forall a b. (a -> b) -> a -> b
$ (a -> Either a a) -> Vector a -> Vector (Either a a)
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Either a a
forall a b. a -> Either a b
Left Vector a
left Vector (Either a a) -> Vector (Either a a) -> Vector (Either a a)
forall a. Semigroup a => a -> a -> a
<> (a -> Either a a) -> Vector a -> Vector (Either a a)
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Either a a
forall a b. b -> Either a b
Right (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
right)
concatColumnsEither (UnboxedColumn Maybe Bitmap
_ Vector a
left) (BoxedColumn Maybe Bitmap
_ Vector a
right) =
    Maybe Bitmap -> Vector (Either a a) -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector (Either a a) -> Column) -> Vector (Either a a) -> Column
forall a b. (a -> b) -> a -> b
$ (a -> Either a a) -> Vector a -> Vector (Either a a)
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Either a a
forall a b. a -> Either a b
Left (Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
left) Vector (Either a a) -> Vector (Either a a) -> Vector (Either a a)
forall a. Semigroup a => a -> a -> a
<> (a -> Either a a) -> Vector a -> Vector (Either a a)
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Either a a
forall a b. b -> Either a b
Right Vector a
right

-- | Allocate a mutable column of size @n@ matching the constructor/type of the given column.
newMutableColumn :: Int -> Column -> IO MutableColumn
newMutableColumn :: Int -> Column -> IO MutableColumn
newMutableColumn Int
n (BoxedColumn Maybe Bitmap
_ (Vector a
_ :: VB.Vector a)) =
    IOVector a -> MutableColumn
forall a. Columnable a => IOVector a -> MutableColumn
MBoxedColumn (IOVector a -> MutableColumn)
-> IO (IOVector a) -> IO MutableColumn
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> IO (MVector (PrimState IO) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
VBM.new Int
n :: IO (VBM.IOVector a))
newMutableColumn Int
n (UnboxedColumn Maybe Bitmap
_ (Vector a
_ :: VU.Vector a)) =
    IOVector a -> MutableColumn
forall a. (Columnable a, Unbox a) => IOVector a -> MutableColumn
MUnboxedColumn (IOVector a -> MutableColumn)
-> IO (IOVector a) -> IO MutableColumn
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> IO (MVector (PrimState IO) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
n :: IO (VUM.IOVector a))
newMutableColumn Int
n c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
_) = Int -> Column -> IO MutableColumn
newMutableColumn Int
n (Column -> Column
materializePacked Column
c)
newMutableColumn Int
n c :: Column
c@(MergedColumn Column
_ Column
_) = Int -> Column -> IO MutableColumn
newMutableColumn Int
n (Column -> Column
materializeMerged Column
c)

-- | Copy a column chunk into a mutable column starting at offset @off@.
copyIntoMutableColumn :: MutableColumn -> Int -> Column -> IO ()
copyIntoMutableColumn :: MutableColumn -> Int -> Column -> IO ()
copyIntoMutableColumn MutableColumn
mv Int
off c :: Column
c@(MergedColumn Column
_ Column
_) =
    MutableColumn -> Int -> Column -> IO ()
copyIntoMutableColumn MutableColumn
mv Int
off (Column -> Column
materializeMerged Column
c)
copyIntoMutableColumn (MBoxedColumn (IOVector a
mv :: VBM.IOVector b)) Int
off (BoxedColumn Maybe Bitmap
_ (Vector a
v :: VB.Vector a)) =
    case TypeRep a -> TypeRep a -> Maybe (a :~: 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 @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) of
        Just a :~: a
Refl -> (Int -> a -> IO ()) -> Vector a -> IO ()
forall (m :: * -> *) (v :: * -> *) a b.
(Monad m, Vector v a) =>
(Int -> a -> m b) -> v a -> m ()
VG.imapM_ (\Int
i a
x -> MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
VBM.unsafeWrite IOVector a
MVector (PrimState IO) a
mv (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) a
x) Vector a
Vector a
v
        Maybe (a :~: a)
Nothing -> [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"copyIntoMutableColumn: Boxed type mismatch"
copyIntoMutableColumn (MUnboxedColumn (IOVector a
mv :: VUM.IOVector b)) Int
off (UnboxedColumn Maybe Bitmap
_ (Vector a
v :: VU.Vector a)) =
    case TypeRep a -> TypeRep a -> Maybe (a :~: 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 @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) of
        Just a :~: a
Refl -> (Int -> a -> IO ()) -> Vector a -> IO ()
forall (m :: * -> *) (v :: * -> *) a b.
(Monad m, Vector v a) =>
(Int -> a -> m b) -> v a -> m ()
VG.imapM_ (\Int
i a
x -> MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite IOVector a
MVector (PrimState IO) a
mv (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) a
x) Vector a
Vector a
v
        Maybe (a :~: a)
Nothing -> [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"copyIntoMutableColumn: Unboxed type mismatch"
copyIntoMutableColumn MutableColumn
mc Int
off c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
_) =
    MutableColumn -> Int -> Column -> IO ()
copyIntoMutableColumn MutableColumn
mc Int
off (Column -> Column
materializePacked Column
c)
copyIntoMutableColumn MutableColumn
_ Int
_ Column
_ =
    [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"copyIntoMutableColumn: constructor mismatch"

-- | Freeze a mutable column into an immutable column.
freezeMutableColumn :: MutableColumn -> IO Column
freezeMutableColumn :: MutableColumn -> IO Column
freezeMutableColumn (MBoxedColumn IOVector a
mv) = Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
forall a. Maybe a
Nothing (Vector a -> Column) -> IO (Vector a) -> IO Column
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) a -> IO (Vector a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
VB.unsafeFreeze IOVector a
MVector (PrimState IO) a
mv
freezeMutableColumn (MUnboxedColumn IOVector a
mv) = 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 -> Column) -> IO (Vector a) -> IO Column
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) a -> IO (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze IOVector a
MVector (PrimState IO) a
mv

{- | O(n) Converts a column to a list. Throws an exception if the wrong type is specified.

__Examples:__

@
> column = fromList [(1 :: Int), 2, 3, 4]
> toList @Int column
[1,2,3,4]
> toList @Double column
exception: ...
@
-}
toList :: forall a. (Columnable a) => Column -> [a]
toList :: forall a. Columnable a => Column -> [a]
toList Column
xs = case forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector @a Column
xs of
    Left DataFrameException
err -> DataFrameException -> [a]
forall a e. Exception e => e -> a
throw DataFrameException
err
    Right Vector a
val -> Vector a -> [a]
forall a. Vector a -> [a]
VB.toList Vector a
val

{- | Type-safe conversion of a column to a vector of element type @a@ (specify via
type application); 'Left' 'TypeMismatchException' when the column's type differs.

>>> toVector @Int @VU.Vector column
Right (unboxed vector of Ints)

>>> toVector @Text @VB.Vector column
Right (boxed vector of Text)
-}
toVector ::
    forall a v.
    (VG.Vector v a, Columnable a) => Column -> Either DataFrameException (v a)
toVector :: forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector Column
col = case Column
col of
    PackedText Maybe Bitmap
_ PackedTextData
_ -> Column -> Either DataFrameException (v a)
forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector (Column -> Column
materializePacked Column
col)
    MergedColumn Column
_ Column
_ -> Column -> Either DataFrameException (v a)
forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector (Column -> Column
materializeMerged Column
col)
    BoxedColumn Maybe Bitmap
bm (Vector a
inner :: VB.Vector c) ->
        -- Check if user wants Maybe c (nullable) or c directly
        case TypeRep a -> TypeRep a -> Maybe (a :~: 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 @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @c) of
            Just a :~: a
Refl -> v a -> Either DataFrameException (v a)
forall a b. b -> Either a b
Right (v a -> Either DataFrameException (v a))
-> v a -> Either DataFrameException (v a)
forall a b. (a -> b) -> a -> b
$ Vector a -> v a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
Vector a
inner
            Maybe (a :~: a)
Nothing ->
                -- Try: a = Maybe c
                case TypeRep a -> TypeRep (Maybe a) -> Maybe (a :~: 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 @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @(Maybe c)) of
                    Just a :~: Maybe a
Refl ->
                        -- Use VB.generate to avoid fusion forcing null slots
                        let !n :: Int
n = Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
inner
                            maybeVec :: Vector (Maybe a)
maybeVec = case Maybe Bitmap
bm of
                                Maybe Bitmap
Nothing -> Int -> (Int -> Maybe a) -> Vector (Maybe a)
forall a. Int -> (Int -> a) -> Vector a
VB.generate Int
n (a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> (Int -> a) -> Int -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Int -> a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector a
inner)
                                Just Bitmap
bitmap -> Int -> (Int -> Maybe a) -> Vector (Maybe a)
forall a. Int -> (Int -> a) -> Vector a
VB.generate Int
n ((Int -> Maybe a) -> Vector (Maybe a))
-> (Int -> Maybe a) -> Vector (Maybe a)
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                                    if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bitmap Int
i then a -> Maybe a
forall a. a -> Maybe a
Just (Vector a -> Int -> a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector a
inner Int
i) else Maybe a
forall a. Maybe a
Nothing
                         in v a -> Either DataFrameException (v a)
forall a b. b -> Either a b
Right (v a -> Either DataFrameException (v a))
-> v a -> Either DataFrameException (v a)
forall a b. (a -> b) -> a -> b
$ Vector a -> v a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
Vector (Maybe a)
maybeVec
                    Maybe (a :~: Maybe a)
Nothing ->
                        DataFrameException -> Either DataFrameException (v a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (v a))
-> DataFrameException -> Either DataFrameException (v a)
forall a b. (a -> b) -> a -> b
$
                            TypeErrorContext a a -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                                ( MkTypeErrorContext
                                    { userType :: Either [Char] (TypeRep a)
userType = TypeRep a -> Either [Char] (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 [Char] (TypeRep a)
expectedType = TypeRep a -> Either [Char] (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 @c)
                                    , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"toVector"
                                    , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                                    }
                                )
    UnboxedColumn Maybe Bitmap
bm (Vector a
inner :: VU.Vector c) ->
        case TypeRep a -> TypeRep a -> Maybe (a :~: 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 @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @c) of
            Just a :~: a
Refl -> v a -> Either DataFrameException (v a)
forall a b. b -> Either a b
Right (v a -> Either DataFrameException (v a))
-> v a -> Either DataFrameException (v a)
forall a b. (a -> b) -> a -> b
$ Vector a -> v a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
Vector a
inner
            Maybe (a :~: a)
Nothing ->
                case TypeRep a -> TypeRep (Maybe a) -> Maybe (a :~: 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 @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @(Maybe c)) of
                    Just a :~: Maybe a
Refl ->
                        let maybeVec :: Vector (Maybe a)
maybeVec = case Maybe Bitmap
bm of
                                Maybe Bitmap
Nothing -> Int -> (Int -> Maybe a) -> Vector (Maybe a)
forall a. Int -> (Int -> a) -> Vector a
VB.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
inner) (a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> (Int -> a) -> Int -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
inner)
                                Just Bitmap
bitmap -> Int -> (Int -> Maybe a) -> Vector (Maybe a)
forall a. Int -> (Int -> a) -> Vector a
VB.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
inner) ((Int -> Maybe a) -> Vector (Maybe a))
-> (Int -> Maybe a) -> Vector (Maybe a)
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                                    if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bitmap Int
i then a -> Maybe a
forall a. a -> Maybe a
Just (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
inner Int
i) else Maybe a
forall a. Maybe a
Nothing
                         in v a -> Either DataFrameException (v a)
forall a b. b -> Either a b
Right (v a -> Either DataFrameException (v a))
-> v a -> Either DataFrameException (v a)
forall a b. (a -> b) -> a -> b
$ Vector a -> v a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VG.convert Vector a
Vector (Maybe a)
maybeVec
                    Maybe (a :~: Maybe a)
Nothing ->
                        DataFrameException -> Either DataFrameException (v a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (v a))
-> DataFrameException -> Either DataFrameException (v a)
forall a b. (a -> b) -> a -> b
$
                            TypeErrorContext a a -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                                ( MkTypeErrorContext
                                    { userType :: Either [Char] (TypeRep a)
userType = TypeRep a -> Either [Char] (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 [Char] (TypeRep a)
expectedType = TypeRep a -> Either [Char] (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 @c)
                                    , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"toVector"
                                    , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                                    }
                                )

-- Some common types we will use for numerical computing.

{- | Convert a column to an unboxed 'Double' vector, coercing numeric types
('realToFrac' for floats, 'fromIntegral' for integrals; nulls become @NaN@).
'Left' 'TypeMismatchException' when the column is not numeric.
-}
toDoubleVector :: Column -> Either DataFrameException (VU.Vector Double)
toDoubleVector :: Column -> Either DataFrameException (Vector Double)
toDoubleVector Column
column =
    case Column
column of
        PackedText Maybe Bitmap
_ PackedTextData
_ -> Column -> Either DataFrameException (Vector Double)
toDoubleVector (Column -> Column
materializePacked Column
column)
        MergedColumn Column
_ Column
_ -> Column -> Either DataFrameException (Vector Double)
toDoubleVector (Column -> Column
materializeMerged Column
column)
        UnboxedColumn Maybe Bitmap
bm (Vector a
f :: VU.Vector a) -> 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 -> case Maybe Bitmap
bm of
                Maybe Bitmap
Nothing -> Vector Double -> Either DataFrameException (Vector Double)
forall a b. b -> Either a b
Right Vector a
Vector Double
f
                Just Bitmap
bitmap -> Vector Double -> Either DataFrameException (Vector Double)
forall a b. b -> Either a b
Right (Vector Double -> Either DataFrameException (Vector Double))
-> Vector Double -> Either DataFrameException (Vector Double)
forall a b. (a -> b) -> a -> b
$ (Int -> Double -> Double) -> Vector Double -> Vector Double
forall a b.
(Unbox a, Unbox b) =>
(Int -> a -> b) -> Vector a -> Vector b
VU.imap (\Int
i Double
x -> if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bitmap Int
i then Double
x else [Char] -> Double
forall a. Read a => [Char] -> a
read [Char]
"NaN") Vector a
Vector Double
f
            Maybe (a :~: Double)
Nothing -> case forall a. SBoolI (FloatingTypes a) => SBool (FloatingTypes a)
sFloating @a of
                SBool (FloatingTypes a)
STrue ->
                    Vector Double -> Either DataFrameException (Vector Double)
forall a b. b -> Either a b
Right
                        ( (Int -> a -> Double) -> Vector a -> Vector Double
forall a b.
(Unbox a, Unbox b) =>
(Int -> a -> b) -> Vector a -> Vector b
VU.imap
                            ( \Int
i a
x -> case Maybe Bitmap
bm of
                                Just Bitmap
bitmap | Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bitmap Int
i) -> [Char] -> Double
forall a. Read a => [Char] -> a
read [Char]
"NaN"
                                Maybe Bitmap
_ -> a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac a
x
                            )
                            Vector a
f
                        )
                SBool (FloatingTypes a)
SFalse -> case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @a of
                    SBool (IntegralTypes a)
STrue ->
                        Vector Double -> Either DataFrameException (Vector Double)
forall a b. b -> Either a b
Right
                            ( (Int -> a -> Double) -> Vector a -> Vector Double
forall a b.
(Unbox a, Unbox b) =>
(Int -> a -> b) -> Vector a -> Vector b
VU.imap
                                ( \Int
i a
x -> case Maybe Bitmap
bm of
                                    Just Bitmap
bitmap | Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bitmap Int
i) -> [Char] -> Double
forall a. Read a => [Char] -> a
read [Char]
"NaN"
                                    Maybe Bitmap
_ -> a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x
                                )
                                Vector a
f
                            )
                    SBool (IntegralTypes a)
SFalse ->
                        DataFrameException -> Either DataFrameException (Vector Double)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Vector Double))
-> DataFrameException -> Either DataFrameException (Vector Double)
forall a b. (a -> b) -> a -> b
$
                            TypeErrorContext Double a -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                                ( MkTypeErrorContext
                                    { userType :: Either [Char] (TypeRep Double)
userType = TypeRep Double -> Either [Char] (TypeRep Double)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Double)
                                    , expectedType :: Either [Char] (TypeRep a)
expectedType = TypeRep a -> Either [Char] (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)
                                    , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"toDoubleVector"
                                    , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                                    }
                                )
        BoxedColumn Maybe Bitmap
bm (Vector a
f :: VB.Vector a) -> case TypeRep a -> TypeRep Integer -> Maybe (a :~: Integer)
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 @Integer) of
            Just a :~: Integer
Refl ->
                Vector Double -> Either DataFrameException (Vector Double)
forall a b. b -> Either a b
Right
                    ( Vector Double -> Vector Double
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VB.convert (Vector Double -> Vector Double) -> Vector Double -> Vector Double
forall a b. (a -> b) -> a -> b
$
                        (Int -> a -> Double) -> Vector a -> Vector Double
forall a b. (Int -> a -> b) -> Vector a -> Vector b
VB.imap
                            ( \Int
i a
x -> case Maybe Bitmap
bm of
                                Just Bitmap
bitmap | Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bitmap Int
i) -> [Char] -> Double
forall a. Read a => [Char] -> a
read [Char]
"NaN"
                                Maybe Bitmap
_ -> a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x
                            )
                            Vector a
f
                    )
            Maybe (a :~: Integer)
Nothing ->
                DataFrameException -> Either DataFrameException (Vector Double)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Vector Double))
-> DataFrameException -> Either DataFrameException (Vector Double)
forall a b. (a -> b) -> a -> b
$
                    TypeErrorContext Double () -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                        ( MkTypeErrorContext
                            { userType :: Either [Char] (TypeRep Double)
userType = TypeRep Double -> Either [Char] (TypeRep Double)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Double)
                            , expectedType :: Either [Char] (TypeRep ())
expectedType = [Char] -> Either [Char] (TypeRep ())
forall a b. a -> Either a b
Left (Column -> [Char]
columnTypeString Column
column) :: Either String (TypeRep ())
                            , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"toDoubleVector"
                            , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                            }
                        )

{- | Convert a column to an unboxed 'Float' vector, coercing numeric types (nulls
become @NaN@); 'Left' 'TypeMismatchException' when not numeric. Converting from
'Double' may lose precision.
-}
toFloatVector :: Column -> Either DataFrameException (VU.Vector Float)
toFloatVector :: Column -> Either DataFrameException (Vector Float)
toFloatVector Column
column =
    case Column
column of
        PackedText Maybe Bitmap
_ PackedTextData
_ -> Column -> Either DataFrameException (Vector Float)
toFloatVector (Column -> Column
materializePacked Column
column)
        MergedColumn Column
_ Column
_ -> Column -> Either DataFrameException (Vector Float)
toFloatVector (Column -> Column
materializeMerged Column
column)
        UnboxedColumn Maybe Bitmap
bm (Vector a
f :: VU.Vector a) -> 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 -> case Maybe Bitmap
bm of
                Maybe Bitmap
Nothing -> Vector Float -> Either DataFrameException (Vector Float)
forall a b. b -> Either a b
Right Vector a
Vector Float
f
                Just Bitmap
bitmap -> Vector Float -> Either DataFrameException (Vector Float)
forall a b. b -> Either a b
Right (Vector Float -> Either DataFrameException (Vector Float))
-> Vector Float -> Either DataFrameException (Vector Float)
forall a b. (a -> b) -> a -> b
$ (Int -> Float -> Float) -> Vector Float -> Vector Float
forall a b.
(Unbox a, Unbox b) =>
(Int -> a -> b) -> Vector a -> Vector b
VU.imap (\Int
i Float
x -> if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bitmap Int
i then Float
x else [Char] -> Float
forall a. Read a => [Char] -> a
read [Char]
"NaN") Vector a
Vector Float
f
            Maybe (a :~: Float)
Nothing -> case forall a. SBoolI (FloatingTypes a) => SBool (FloatingTypes a)
sFloating @a of
                SBool (FloatingTypes a)
STrue ->
                    Vector Float -> Either DataFrameException (Vector Float)
forall a b. b -> Either a b
Right
                        ( (Int -> a -> Float) -> Vector a -> Vector Float
forall a b.
(Unbox a, Unbox b) =>
(Int -> a -> b) -> Vector a -> Vector b
VU.imap
                            ( \Int
i a
x -> case Maybe Bitmap
bm of
                                Just Bitmap
bitmap | Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bitmap Int
i) -> [Char] -> Float
forall a. Read a => [Char] -> a
read [Char]
"NaN"
                                Maybe Bitmap
_ -> a -> Float
forall a b. (Real a, Fractional b) => a -> b
realToFrac a
x
                            )
                            Vector a
f
                        )
                SBool (FloatingTypes a)
SFalse -> case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @a of
                    SBool (IntegralTypes a)
STrue ->
                        Vector Float -> Either DataFrameException (Vector Float)
forall a b. b -> Either a b
Right
                            ( (Int -> a -> Float) -> Vector a -> Vector Float
forall a b.
(Unbox a, Unbox b) =>
(Int -> a -> b) -> Vector a -> Vector b
VU.imap
                                ( \Int
i a
x -> case Maybe Bitmap
bm of
                                    Just Bitmap
bitmap | Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bitmap Int
i) -> [Char] -> Float
forall a. Read a => [Char] -> a
read [Char]
"NaN"
                                    Maybe Bitmap
_ -> a -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x
                                )
                                Vector a
f
                            )
                    SBool (IntegralTypes a)
SFalse ->
                        DataFrameException -> Either DataFrameException (Vector Float)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Vector Float))
-> DataFrameException -> Either DataFrameException (Vector Float)
forall a b. (a -> b) -> a -> b
$
                            TypeErrorContext Float a -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                                ( MkTypeErrorContext
                                    { userType :: Either [Char] (TypeRep Float)
userType = TypeRep Float -> Either [Char] (TypeRep Float)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Float)
                                    , expectedType :: Either [Char] (TypeRep a)
expectedType = TypeRep a -> Either [Char] (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)
                                    , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"toFloatVector"
                                    , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                                    }
                                )
        BoxedColumn Maybe Bitmap
bm (Vector a
f :: VB.Vector a) -> case TypeRep a -> TypeRep Integer -> Maybe (a :~: Integer)
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 @Integer) of
            Just a :~: Integer
Refl ->
                Vector Float -> Either DataFrameException (Vector Float)
forall a b. b -> Either a b
Right
                    ( Vector Float -> Vector Float
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VB.convert (Vector Float -> Vector Float) -> Vector Float -> Vector Float
forall a b. (a -> b) -> a -> b
$
                        (Int -> a -> Float) -> Vector a -> Vector Float
forall a b. (Int -> a -> b) -> Vector a -> Vector b
VB.imap
                            ( \Int
i a
x -> case Maybe Bitmap
bm of
                                Just Bitmap
bitmap | Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bitmap Int
i) -> [Char] -> Float
forall a. Read a => [Char] -> a
read [Char]
"NaN"
                                Maybe Bitmap
_ -> a -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x
                            )
                            Vector a
f
                    )
            Maybe (a :~: Integer)
Nothing ->
                DataFrameException -> Either DataFrameException (Vector Float)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Vector Float))
-> DataFrameException -> Either DataFrameException (Vector Float)
forall a b. (a -> b) -> a -> b
$
                    TypeErrorContext Float () -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                        ( MkTypeErrorContext
                            { userType :: Either [Char] (TypeRep Float)
userType = TypeRep Float -> Either [Char] (TypeRep Float)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Float)
                            , expectedType :: Either [Char] (TypeRep ())
expectedType = [Char] -> Either [Char] (TypeRep ())
forall a b. a -> Either a b
Left (Column -> [Char]
columnTypeString Column
column) :: Either String (TypeRep ())
                            , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"toFloatVector"
                            , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                            }
                        )

{- | Convert a column to an unboxed 'Int' vector, coercing numeric types
(floats are 'round'ed via banker's rounding); 'Left' 'TypeMismatchException'
when the column is not numeric. Does not support nullable columns.
-}
toIntVector :: Column -> Either DataFrameException (VU.Vector Int)
toIntVector :: Column -> Either DataFrameException (Vector Int)
toIntVector Column
column =
    case Column
column of
        PackedText Maybe Bitmap
_ PackedTextData
_ -> Column -> Either DataFrameException (Vector Int)
toIntVector (Column -> Column
materializePacked Column
column)
        MergedColumn Column
_ Column
_ -> Column -> Either DataFrameException (Vector Int)
toIntVector (Column -> Column
materializeMerged Column
column)
        UnboxedColumn Maybe Bitmap
_ (Vector a
f :: VU.Vector a) -> 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 -> Vector Int -> Either DataFrameException (Vector Int)
forall a b. b -> Either a b
Right Vector a
Vector Int
f
            Maybe (a :~: Int)
Nothing -> case forall a. SBoolI (FloatingTypes a) => SBool (FloatingTypes a)
sFloating @a of
                SBool (FloatingTypes a)
STrue -> Vector Int -> Either DataFrameException (Vector Int)
forall a b. b -> Either a b
Right ((a -> Int) -> Vector a -> Vector Int
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (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 :: a -> Double)) Vector a
f)
                SBool (FloatingTypes a)
SFalse -> case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @a of
                    SBool (IntegralTypes a)
STrue -> Vector Int -> Either DataFrameException (Vector Int)
forall a b. b -> Either a b
Right ((a -> Int) -> Vector a -> Vector Int
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Vector a
f)
                    SBool (IntegralTypes a)
SFalse ->
                        DataFrameException -> Either DataFrameException (Vector Int)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Vector Int))
-> DataFrameException -> Either DataFrameException (Vector Int)
forall a b. (a -> b) -> a -> b
$
                            TypeErrorContext Int a -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                                ( MkTypeErrorContext
                                    { userType :: Either [Char] (TypeRep Int)
userType = TypeRep Int -> Either [Char] (TypeRep Int)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Int)
                                    , expectedType :: Either [Char] (TypeRep a)
expectedType = TypeRep a -> Either [Char] (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)
                                    , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"toIntVector"
                                    , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                                    }
                                )
        BoxedColumn Maybe Bitmap
_ (Vector a
f :: VB.Vector a) -> case TypeRep a -> TypeRep Integer -> Maybe (a :~: Integer)
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 @Integer) of
            Just a :~: Integer
Refl -> Vector Int -> Either DataFrameException (Vector Int)
forall a b. b -> Either a b
Right (Vector Int -> Vector Int
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VB.convert (Vector Int -> Vector Int) -> Vector Int -> Vector Int
forall a b. (a -> b) -> a -> b
$ (a -> Int) -> Vector a -> Vector Int
forall a b. (a -> b) -> Vector a -> Vector b
VB.map a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Vector a
f)
            Maybe (a :~: Integer)
Nothing ->
                DataFrameException -> Either DataFrameException (Vector Int)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Vector Int))
-> DataFrameException -> Either DataFrameException (Vector Int)
forall a b. (a -> b) -> a -> b
$
                    TypeErrorContext Int () -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                        ( MkTypeErrorContext
                            { userType :: Either [Char] (TypeRep Int)
userType = TypeRep Int -> Either [Char] (TypeRep Int)
forall a b. b -> Either a b
Right (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Int)
                            , expectedType :: Either [Char] (TypeRep ())
expectedType = [Char] -> Either [Char] (TypeRep ())
forall a b. a -> Either a b
Left (Column -> [Char]
columnTypeString Column
column) :: Either String (TypeRep ())
                            , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"toIntVector"
                            , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                            }
                        )

toUnboxedVector ::
    forall a.
    (Columnable a, VU.Unbox a) => Column -> Either DataFrameException (VU.Vector a)
toUnboxedVector :: forall a.
(Columnable a, Unbox a) =>
Column -> Either DataFrameException (Vector a)
toUnboxedVector Column
column =
    case Column
column of
        UnboxedColumn Maybe Bitmap
_ (Vector a
f :: VU.Vector b) -> case TypeRep a -> TypeRep a -> Maybe (a :~: 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 @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) of
            Just a :~: a
Refl -> Vector a -> Either DataFrameException (Vector a)
forall a b. b -> Either a b
Right Vector a
Vector a
f
            Maybe (a :~: a)
Nothing ->
                DataFrameException -> Either DataFrameException (Vector a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Vector a))
-> DataFrameException -> Either DataFrameException (Vector a)
forall a b. (a -> b) -> a -> b
$
                    TypeErrorContext a a -> DataFrameException
forall a b.
(Typeable a, Typeable b) =>
TypeErrorContext a b -> DataFrameException
TypeMismatchException
                        ( MkTypeErrorContext
                            { userType :: Either [Char] (TypeRep a)
userType = TypeRep a -> Either [Char] (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 [Char] (TypeRep a)
expectedType = TypeRep a -> Either [Char] (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 @b)
                            , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"toUnboxedVector"
                            , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                            }
                        )
        Column
_ ->
            DataFrameException -> Either DataFrameException (Vector a)
forall a b. a -> Either a b
Left (DataFrameException -> Either DataFrameException (Vector a))
-> DataFrameException -> Either DataFrameException (Vector 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 [Char] (TypeRep a)
userType = TypeRep a -> Either [Char] (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 [Char] (TypeRep ())
expectedType = [Char] -> Either [Char] (TypeRep ())
forall a b. a -> Either a b
Left (Column -> [Char]
columnTypeString Column
column) :: Either String (TypeRep ())
                        , callingFunctionName :: Maybe [Char]
callingFunctionName = [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
"toUnboxedVector"
                        , errorColumnName :: Maybe [Char]
errorColumnName = Maybe [Char]
forall a. Maybe a
Nothing
                        }
                    )
{-# INLINE toUnboxedVector #-}

-- Shared finaliser for the two parseUnboxedColumn* helpers.  Freezes
-- the mutable data vector, and only materialises the bitmap when the
-- column actually had nulls.
{-# INLINE finalizeParseResult #-}
finalizeParseResult ::
    (VU.Unbox a) =>
    VUM.STVector s a ->
    VUM.STVector s Word8 ->
    Bool ->
    ST s (Maybe (Maybe Bitmap, VU.Vector a))
finalizeParseResult :: forall a s.
Unbox a =>
STVector s a
-> STVector s Word8
-> Bool
-> ST s (Maybe (Maybe Bitmap, Vector a))
finalizeParseResult STVector s a
values STVector s Word8
vmask Bool
anyNull
    | Bool
anyNull = do
        Vector a
vs <- 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 STVector s a
MVector (PrimState (ST s)) a
values
        Bitmap
vm <- MVector (PrimState (ST s)) Word8 -> ST s Bitmap
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze STVector s Word8
MVector (PrimState (ST s)) Word8
vmask
        Maybe (Maybe Bitmap, Vector a)
-> ST s (Maybe (Maybe Bitmap, Vector a))
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Maybe Bitmap, Vector a) -> Maybe (Maybe Bitmap, Vector a)
forall a. a -> Maybe a
Just (Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Bitmap -> Bitmap
buildBitmapFromValid Bitmap
vm), Vector a
vs))
    | Bool
otherwise = do
        Vector a
vs <- 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 STVector s a
MVector (PrimState (ST s)) a
values
        Maybe (Maybe Bitmap, Vector a)
-> ST s (Maybe (Maybe Bitmap, Vector a))
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Maybe Bitmap, Vector a) -> Maybe (Maybe Bitmap, Vector a)
forall a. a -> Maybe a
Just (Maybe Bitmap
forall a. Maybe a
Nothing, Vector a
vs))