{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

{- | Driver for the default CSV reader: a single strict scan feeding
per-column sinks, with cassava-parity semantics. Ragged rows pad
trailing columns with null, drop extra fields, and force the result.
-}
module DataFrame.IO.CSV.Internal.Read (decodeCsvStrict) where

import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BSU
import qualified Data.Map.Strict as M
import qualified Data.Proxy as P
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU

import Control.Monad (when, zipWithM)
import Control.Monad.ST (stToIO)
import Data.Char (ord)
import Data.Maybe (fromMaybe)
import Data.Type.Equality (TestEquality (testEquality))
import Data.Word (Word8)
import DataFrame.IO.CSV.Internal.Infer
import DataFrame.IO.CSV.Internal.Options
import DataFrame.IO.CSV.Internal.Scanner
import DataFrame.IO.CSV.Internal.Sink
import DataFrame.Internal.Column (Column, ensureOptional)
import DataFrame.Internal.ColumnBuilder
import DataFrame.Internal.DataFrame (DataFrame (..), forceDataFrame)
import DataFrame.Operations.Typing (
    SafeReadMode (..),
    effectiveSafeRead,
    parseWithTypes,
 )
import DataFrame.Schema (SchemaType (..), schemaType)
import Foreign.Ptr (castPtr)
import Type.Reflection (typeRep)

-- | Decode a strict CSV buffer into a fully forced DataFrame.
decodeCsvStrict :: ReadOptions -> BS.ByteString -> IO DataFrame
decodeCsvStrict :: ReadOptions -> ByteString -> IO DataFrame
decodeCsvStrict ReadOptions
opts ByteString
bs = do
    ReadOptions -> IO ()
validateReadOptions ReadOptions
opts
    ByteString -> (CStringLen -> IO DataFrame) -> IO DataFrame
forall a. ByteString -> (CStringLen -> IO a) -> IO a
BSU.unsafeUseAsCStringLen ByteString
bs ((CStringLen -> IO DataFrame) -> IO DataFrame)
-> (CStringLen -> IO DataFrame) -> IO DataFrame
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
cstr, Int
_) -> do
        let !len :: Int
len = ByteString -> Int
BS.length ByteString
bs
            !sep :: Word8
sep = Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Char -> Int
ord (ReadOptions -> Char
columnSeparator ReadOptions
opts)) :: Word8

            findRecord :: Int -> Int
findRecord !Int
pos
                | Int
pos Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len = Int
len
                | Bool
otherwise = ByteString
-> Int
-> Word8
-> Int
-> (Int -> Int -> Bool -> Int -> Int -> Int)
-> Int
forall r.
ByteString
-> Int
-> Word8
-> Int
-> (Int -> Int -> Bool -> Int -> Int -> r)
-> r
withField ByteString
bs Int
len Word8
sep Int
pos ((Int -> Int -> Bool -> Int -> Int -> Int) -> Int)
-> (Int -> Int -> Bool -> Int -> Int -> Int) -> Int
forall a b. (a -> b) -> a -> b
$ \Int
cs Int
ce Bool
_ Int
term Int
next ->
                    if Int
cs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
ce Bool -> Bool -> Bool
&& Int
term Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
termSep
                        then if Int
term Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
termEof then Int
len else Int -> Int
findRecord Int
next
                        else Int
pos

            headerFields :: Int -> ([Text], Int)
headerFields !Int
pos = Int -> ([Text] -> [Text]) -> ([Text], Int)
go Int
pos [Text] -> [Text]
forall a. a -> a
id
              where
                go :: Int -> ([Text] -> [Text]) -> ([Text], Int)
go !Int
p [Text] -> [Text]
acc = ByteString
-> Int
-> Word8
-> Int
-> (Int -> Int -> Bool -> Int -> Int -> ([Text], Int))
-> ([Text], Int)
forall r.
ByteString
-> Int
-> Word8
-> Int
-> (Int -> Int -> Bool -> Int -> Int -> r)
-> r
withField ByteString
bs Int
len Word8
sep Int
p ((Int -> Int -> Bool -> Int -> Int -> ([Text], Int))
 -> ([Text], Int))
-> (Int -> Int -> Bool -> Int -> Int -> ([Text], Int))
-> ([Text], Int)
forall a b. (a -> b) -> a -> b
$ \Int
cs Int
ce Bool
unesc Int
term Int
next ->
                    let t :: Text
t =
                            ByteString -> Text
TE.decodeUtf8Lenient
                                (if Bool
unesc then ByteString -> Int -> Int -> ByteString
unescapeQuotes ByteString
bs Int
cs Int
ce else ByteString -> Int -> Int -> ByteString
sliceBS ByteString
bs Int
cs Int
ce)
                     in if Int
term Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
termSep
                            then Int -> ([Text] -> [Text]) -> ([Text], Int)
go Int
next ([Text] -> [Text]
acc ([Text] -> [Text]) -> ([Text] -> [Text]) -> [Text] -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text
t Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
:))
                            else ([Text] -> [Text]
acc [Text
t], Int
next)

        let hdrPos :: Int
hdrPos = Int -> Int
findRecord Int
0
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
hdrPos Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len) ([Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty CSV file")
        let ([Text]
hdrFields, Int
afterHdr) = Int -> ([Text], Int)
headerFields Int
hdrPos
            positional :: Int -> [Text]
positional Int
n = (Int -> Text) -> [Int] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map ([Char] -> Text
T.pack ([Char] -> Text) -> (Int -> [Char]) -> Int -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Char]
forall a. Show a => a -> [Char]
show) [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 :: Int]
            ([Text]
names, Int
dataPos0) = case ReadOptions -> HeaderSpec
headerSpec ReadOptions
opts of
                HeaderSpec
UseFirstRow -> ((Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Text
T.strip [Text]
hdrFields, Int
afterHdr)
                HeaderSpec
NoHeader -> (Int -> [Text]
positional ([Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Text]
hdrFields), Int
hdrPos)
                ProvideNames [Text]
ns ->
                    ([Text]
ns [Text] -> [Text] -> [Text]
forall a. [a] -> [a] -> [a]
++ Int -> [Text] -> [Text]
forall a. Int -> [a] -> [a]
drop ([Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Text]
ns) (Int -> [Text]
positional ([Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Text]
hdrFields)), Int
hdrPos)
            !nfields :: Int
nfields = [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Text]
names
            dataPos :: Int
dataPos = Int -> Int
findRecord Int
dataPos0
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
dataPos Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len) ([Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty CSV file")

        let selection :: [(Text, Int)]
selection = ReadOptions -> [Text] -> [(Text, Int)]
resolveSelection ReadOptions
opts [Text]
names
            !keptNames :: [Text]
keptNames = ((Text, Int) -> Text) -> [(Text, Int)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Int) -> Text
forall a b. (a, b) -> a
fst [(Text, Int)]
selection
            !ncols :: Int
ncols = [(Text, Int)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Text, Int)]
selection
            !sinkOf :: Vector Int
sinkOf =
                Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
nfields (-Int
1)
                    Vector Int -> [(Int, Int)] -> Vector Int
forall a. Unbox a => Vector a -> [(Int, a)] -> Vector a
VU.// [(Int
fieldIx, Int
slot) | (Int
slot, (Text
_, Int
fieldIx)) <- [Int] -> [(Text, Int)] -> [(Int, (Text, Int))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [(Text, Int)]
selection]

        let resolveMode :: Text -> SafeReadMode
resolveMode = SafeReadMode -> [(Text, SafeReadMode)] -> Text -> SafeReadMode
effectiveSafeRead (ReadOptions -> SafeReadMode
safeRead ReadOptions
opts) (ReadOptions -> [(Text, SafeReadMode)]
safeReadOverrides ReadOptions
opts)
            anyEither :: Bool
anyEither = (Text -> Bool) -> [Text] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Text
n -> Text -> SafeReadMode
resolveMode Text
n SafeReadMode -> SafeReadMode -> Bool
forall a. Eq a => a -> a -> Bool
== SafeReadMode
EitherRead) [Text]
keptNames
            missing :: [Text]
missing = if Bool
anyEither then [] else ReadOptions -> [Text]
missingIndicators ReadOptions
opts
            missMode :: MissingMode
missMode
                | [Text] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Text]
missing = MissingMode
MissNone
                | [Text]
missing [Text] -> [Text] -> Bool
forall a. Eq a => a -> a -> Bool
== ReadOptions -> [Text]
missingIndicators ReadOptions
defaultReadOptions = MissingMode
MissCanonical
                | Bool
otherwise = MissingMode
MissCustom
            env :: Env
env = ByteString -> Ptr Word8 -> MissingMode -> [Text] -> Env
Env ByteString
bs (Ptr CChar -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr CChar
cstr) MissingMode
missMode [Text]
missing
            rowHint :: Int
rowHint = Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int
nfields Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
8) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
16

        [Sink]
sinks <- (Text -> IO Sink) -> [Text] -> IO [Sink]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (ReadOptions -> (Text -> SafeReadMode) -> Int -> Text -> IO Sink
newSink ReadOptions
opts Text -> SafeReadMode
resolveMode Int
rowHint) [Text]
keptNames
        let !sinksV :: Vector Sink
sinksV = [Sink] -> Vector Sink
forall a. [a] -> Vector a
V.fromList [Sink]
sinks
            sinkAt :: Int -> Int
sinkAt !Int
col
                | Int
col Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nfields = -Int
1
                | Bool
otherwise = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
sinkOf Int
col

        let pad :: Int -> Int -> IO ()
pad !Int
row !Int
col
                | Int
col Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nfields = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                | Bool
otherwise = do
                    let !slot :: Int
slot = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
sinkOf Int
col
                    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slot Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Sink -> Int -> IO ()
nullSink (Vector Sink -> Int -> Sink
forall a. Vector a -> Int -> a
V.unsafeIndex Vector Sink
sinksV Int
slot) Int
row
                    Int -> Int -> IO ()
pad Int
row (Int
col Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            feedAt :: Int -> Int -> Int -> Int -> Bool -> IO ()
feedAt !Int
slot !Int
row !Int
cs !Int
ce !Bool
unesc =
                Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slot Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                    Env -> Sink -> Int -> Int -> Int -> Bool -> IO ()
feedSink Env
env (Vector Sink -> Int -> Sink
forall a. Vector a -> Int -> a
V.unsafeIndex Vector Sink
sinksV Int
slot) Int
row Int
cs Int
ce Bool
unesc
            dec :: a -> a
dec a
b = if a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0 then a
b a -> a -> a
forall a. Num a => a -> a -> a
- a
1 else a
b
            rowLoop :: Int -> Int -> Int -> IO Int
rowLoop !Int
pos !Int
row !Int
budget
                | Int
budget Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| Int
pos Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len = Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
row
                | Bool
otherwise = ByteString
-> Int
-> Word8
-> Int
-> (Int -> Int -> Bool -> Int -> Int -> IO Int)
-> IO Int
forall r.
ByteString
-> Int
-> Word8
-> Int
-> (Int -> Int -> Bool -> Int -> Int -> r)
-> r
withField ByteString
bs Int
len Word8
sep Int
pos ((Int -> Int -> Bool -> Int -> Int -> IO Int) -> IO Int)
-> (Int -> Int -> Bool -> Int -> Int -> IO Int) -> IO Int
forall a b. (a -> b) -> a -> b
$ \Int
cs Int
ce Bool
unesc Int
term Int
next ->
                    if Int
cs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
ce Bool -> Bool -> Bool
&& Int
term Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
termSep
                        then if Int
term Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
termEof then Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
row else Int -> Int -> Int -> IO Int
rowLoop Int
next Int
row Int
budget
                        else do
                            Int -> Int -> Int -> Int -> Bool -> IO ()
feedAt (Int -> Int
sinkAt Int
0) Int
row Int
cs Int
ce Bool
unesc
                            if Int
term Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
termSep
                                then Int -> Int -> Int -> Int -> IO Int
fieldLoop Int
next Int
row Int
1 Int
budget
                                else Int -> Int -> IO ()
pad Int
row Int
1 IO () -> IO Int -> IO Int
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> Int -> IO Int
rowLoop Int
next (Int
row Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> Int
forall {a}. (Ord a, Num a) => a -> a
dec Int
budget)
            fieldLoop :: Int -> Int -> Int -> Int -> IO Int
fieldLoop !Int
pos !Int
row !Int
col !Int
budget =
                ByteString
-> Int
-> Word8
-> Int
-> (Int -> Int -> Bool -> Int -> Int -> IO Int)
-> IO Int
forall r.
ByteString
-> Int
-> Word8
-> Int
-> (Int -> Int -> Bool -> Int -> Int -> r)
-> r
withField ByteString
bs Int
len Word8
sep Int
pos ((Int -> Int -> Bool -> Int -> Int -> IO Int) -> IO Int)
-> (Int -> Int -> Bool -> Int -> Int -> IO Int) -> IO Int
forall a b. (a -> b) -> a -> b
$ \Int
cs Int
ce Bool
unesc Int
term Int
next -> do
                    Int -> Int -> Int -> Int -> Bool -> IO ()
feedAt (Int -> Int
sinkAt Int
col) Int
row Int
cs Int
ce Bool
unesc
                    if Int
term Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
termSep
                        then Int -> Int -> Int -> Int -> IO Int
fieldLoop Int
next Int
row (Int
col Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
budget
                        else Int -> Int -> IO ()
pad Int
row (Int
col Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) IO () -> IO Int -> IO Int
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> Int -> IO Int
rowLoop Int
next (Int
row Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> Int
forall {a}. (Ord a, Num a) => a -> a
dec Int
budget)

        Int
nrows <- Int -> Int -> Int -> IO Int
rowLoop Int
dataPos Int
0 (Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe (-Int
1) (ReadOptions -> Maybe Int
numRowsToRead ReadOptions
opts))

        [Column]
cols <- (Text -> Sink -> IO Column) -> [Text] -> [Sink] -> IO [Column]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (ByteString
-> ReadOptions
-> (Text -> SafeReadMode)
-> Int
-> Text
-> Sink
-> IO Column
finalizeSink ByteString
bs ReadOptions
opts Text -> SafeReadMode
resolveMode Int
nrows) [Text]
keptNames [Sink]
sinks
        let df :: DataFrame
df =
                Vector Column
-> Map Text Int -> (Int, Int) -> Map Text UExpr -> DataFrame
DataFrame
                    ([Column] -> Vector Column
forall a. [a] -> Vector a
V.fromList [Column]
cols)
                    ([(Text, Int)] -> Map Text Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([Text] -> [Int] -> [(Text, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Text]
keptNames [Int
0 ..]))
                    (Int
nrows, Int
ncols)
                    Map Text UExpr
forall k a. Map k a
M.empty
        DataFrame -> IO DataFrame
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DataFrame -> IO DataFrame) -> DataFrame -> IO DataFrame
forall a b. (a -> b) -> a -> b
$!
            DataFrame -> DataFrame
forceDataFrame
                ((Text -> SafeReadMode)
-> Map Text SchemaType -> DataFrame -> DataFrame
parseWithTypes Text -> SafeReadMode
resolveMode (TypeSpec -> Map Text SchemaType
schemaTypeMap (ReadOptions -> TypeSpec
typeSpec ReadOptions
opts)) DataFrame
df)

{- | Resolve a column's sink: EitherRead and to-be-inferred columns
collect raw slices; schema'd Int/Double parse straight into builders;
every other schema type builds Text ('parseWithTypes' converts after).
-}
newSink ::
    ReadOptions -> (T.Text -> SafeReadMode) -> Int -> T.Text -> IO Sink
newSink :: ReadOptions -> (Text -> SafeReadMode) -> Int -> Text -> IO Sink
newSink ReadOptions
opts Text -> SafeReadMode
resolveMode Int
rowHint Text
name
    | Text -> SafeReadMode
resolveMode Text
name SafeReadMode -> SafeReadMode -> Bool
forall a. Eq a => a -> a -> Bool
== SafeReadMode
EitherRead = SliceCol -> Sink
SinkBS (SliceCol -> Sink) -> IO SliceCol -> IO Sink
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO SliceCol
newSliceCol Int
rowHint
    | Maybe SchemaType
Nothing <- Maybe SchemaType
entry
    , TypeSpec -> Bool
shouldInferFromSample (ReadOptions -> TypeSpec
typeSpec ReadOptions
opts) =
        SliceCol -> Sink
SinkBS (SliceCol -> Sink) -> IO SliceCol -> IO Sink
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO SliceCol
newSliceCol Int
rowHint
    | Bool
otherwise = case SchemaType -> Maybe SchemaType -> SchemaType
forall a. a -> Maybe a -> a
fromMaybe (forall a. (Columnable a, Read a) => SchemaType
schemaType @T.Text) Maybe SchemaType
entry of
        SType (Proxy a
_ :: P.Proxy 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
_ -> IntBuilder RealWorld -> Sink
SinkInt (IntBuilder RealWorld -> Sink)
-> IO (IntBuilder RealWorld) -> IO Sink
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ST RealWorld (IntBuilder RealWorld) -> IO (IntBuilder RealWorld)
forall a. ST RealWorld a -> IO a
stToIO (Int -> ST RealWorld (IntBuilder RealWorld)
forall s. Int -> ST s (IntBuilder s)
newIntBuilder Int
rowHint)
            Maybe (a :~: Int)
Nothing -> 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
_ -> DoubleBuilder RealWorld -> Sink
SinkDouble (DoubleBuilder RealWorld -> Sink)
-> IO (DoubleBuilder RealWorld) -> IO Sink
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ST RealWorld (DoubleBuilder RealWorld)
-> IO (DoubleBuilder RealWorld)
forall a. ST RealWorld a -> IO a
stToIO (Int -> ST RealWorld (DoubleBuilder RealWorld)
forall s. Int -> ST s (DoubleBuilder s)
newDoubleBuilder Int
rowHint)
                Maybe (a :~: Double)
Nothing ->
                    TextBuilder RealWorld -> Sink
SinkText (TextBuilder RealWorld -> Sink)
-> IO (TextBuilder RealWorld) -> IO Sink
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ST RealWorld (TextBuilder RealWorld) -> IO (TextBuilder RealWorld)
forall a. ST RealWorld a -> IO a
stToIO (Int -> Int -> ST RealWorld (TextBuilder RealWorld)
forall s. Int -> Int -> ST s (TextBuilder s)
newTextBuilder Int
rowHint (Int
rowHint Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
8))
  where
    entry :: Maybe SchemaType
entry = Text -> Map Text SchemaType -> Maybe SchemaType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Text
name (TypeSpec -> Map Text SchemaType
schemaTypeMap (ReadOptions -> TypeSpec
typeSpec ReadOptions
opts))

finalizeSink ::
    BS.ByteString ->
    ReadOptions ->
    (T.Text -> SafeReadMode) ->
    Int ->
    T.Text ->
    Sink ->
    IO Column
finalizeSink :: ByteString
-> ReadOptions
-> (Text -> SafeReadMode)
-> Int
-> Text
-> Sink
-> IO Column
finalizeSink ByteString
bs ReadOptions
opts Text -> SafeReadMode
resolveMode Int
nrows Text
name Sink
sink = do
    Column
col <- case Sink
sink of
        SinkInt IntBuilder RealWorld
b -> ST RealWorld Column -> IO Column
forall a. ST RealWorld a -> IO a
stToIO (IntBuilder RealWorld -> ST RealWorld Column
forall s. NumBuilder Int s -> ST s Column
forall (b :: * -> *) s. ColumnBuilder b => b s -> ST s Column
freezeBuilder IntBuilder RealWorld
b)
        SinkDouble DoubleBuilder RealWorld
b -> ST RealWorld Column -> IO Column
forall a. ST RealWorld a -> IO a
stToIO (DoubleBuilder RealWorld -> ST RealWorld Column
forall s. NumBuilder Double s -> ST s Column
forall (b :: * -> *) s. ColumnBuilder b => b s -> ST s Column
freezeBuilder DoubleBuilder RealWorld
b)
        SinkText TextBuilder RealWorld
b -> ST RealWorld Column -> IO Column
forall a. ST RealWorld a -> IO a
stToIO (TextBuilder RealWorld -> ST RealWorld Column
forall s. TextBuilder s -> ST s Column
forall (b :: * -> *) s. ColumnBuilder b => b s -> ST s Column
freezeBuilder TextBuilder RealWorld
b)
        SinkBS SliceCol
sc -> do
            Cells
cells <- ByteString -> SliceCol -> Int -> IO Cells
freezeCells ByteString
bs SliceCol
sc Int
nrows
            SafeReadMode -> ReadOptions -> Cells -> IO Column
inferColumnFromBS (Text -> SafeReadMode
resolveMode Text
name) ReadOptions
opts Cells
cells
    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
$! case Text -> SafeReadMode
resolveMode Text
name of
        SafeReadMode
MaybeRead -> Column -> Column
ensureOptional Column
col
        SafeReadMode
_ -> Column
col