{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}

{- | Column inference over raw field bytes for the default CSV reader:
classify a bounded sample with the shared lattice
("DataFrame.Operations.Inference"), then build the column in one fused
pass with in-place Int -> Double promotion (audit T2). Cells arrive as
'Cells' — unboxed offsets into the input buffer plus a builder validity
vector — so the passes parse slices without materializing cell values.
-}
module DataFrame.IO.CSV.Internal.Infer (
    inferColumnFromBS,
) where

import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BSU
import qualified Data.IntMap.Strict as IM
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM

import Control.Monad.ST (runST, stToIO)
import Data.Time (Day)
import Data.Word (Word8)
import DataFrame.IO.CSV.Internal.Options
import DataFrame.IO.CSV.Internal.Sink (Cells (..), cellAt, withCell)
import DataFrame.Internal.Column
import DataFrame.Internal.ColumnBuilder
import DataFrame.Internal.Parsing
import DataFrame.Internal.Parsing.Fast (
    parseBoolFieldSlice,
    parseDateFieldSlice,
    parseDoubleFieldSlice,
    parseIntFieldSlice,
 )
import DataFrame.Operations.Typing
import Foreign.Ptr (Ptr, castPtr, plusPtr)

inferColumnFromBS :: SafeReadMode -> ReadOptions -> Cells -> IO Column
inferColumnFromBS :: SafeReadMode -> ReadOptions -> Cells -> IO Column
inferColumnFromBS SafeReadMode
mode ReadOptions
opts Cells
cells =
    let sampleN :: Int
sampleN = let n :: Int
n = TypeSpec -> Int
typeInferenceSampleSize (ReadOptions -> TypeSpec
typeSpec ReadOptions
opts) in if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Int
100 else Int
n
        dfmt :: String
dfmt = ReadOptions -> String
dateFormat ReadOptions
opts
        -- The sample IS still Maybe-wrapped; it's bounded at 100 rows
        -- by default, so the allocation is ignorable.
        samples :: Vector (Maybe ByteString)
samples = Int -> (Int -> Maybe ByteString) -> Vector (Maybe ByteString)
forall a. Int -> (Int -> a) -> Vector a
V.generate (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
sampleN (Cells -> Int
cLen Cells
cells)) ((Int -> Maybe ByteString) -> Vector (Maybe ByteString))
-> (Int -> Maybe ByteString) -> Vector (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ \Int
i ->
            if Cells -> Int -> Bool
validAt Cells
cells Int
i then ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (Cells -> Int -> ByteString
cellAt Cells
cells Int
i) else Maybe ByteString
forall a. Maybe a
Nothing
        assumption :: ParsingAssumption
assumption = String -> Vector (Maybe ByteString) -> ParsingAssumption
makeParsingAssumptionBytes String
dfmt Vector (Maybe ByteString)
samples
        orText :: Maybe Column -> IO Column
orText = IO Column -> (Column -> IO Column) -> Maybe Column -> IO Column
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Cells -> IO Column
textColumn Cells
cells) Column -> IO Column
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
     in case SafeReadMode
mode of
            SafeReadMode
EitherRead -> Column -> IO Column
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> ParsingAssumption -> Cells -> Column
handleBSEither String
dfmt ParsingAssumption
assumption Cells
cells)
            SafeReadMode
_ -> case ParsingAssumption
assumption of
                ParsingAssumption
IntAssumption -> Maybe Column -> IO Column
orText (Cells -> Maybe Column
handleBSInt Cells
cells)
                ParsingAssumption
DoubleAssumption -> Maybe Column -> IO Column
orText (Cells -> Maybe Column
handleBSDouble Cells
cells)
                ParsingAssumption
BoolAssumption -> Maybe Column -> IO Column
orText (Cells -> Maybe Column
handleBSBool Cells
cells)
                ParsingAssumption
DateAssumption -> Maybe Column -> IO Column
orText (String -> Cells -> Maybe Column
handleBSDate String
dfmt Cells
cells)
                ParsingAssumption
TextAssumption -> Cells -> IO Column
textColumn Cells
cells
                ParsingAssumption
NoAssumption -> String -> Cells -> IO Column
handleBSNo String
dfmt Cells
cells

validAt :: Cells -> Int -> Bool
validAt :: Cells -> Int -> Bool
validAt Cells
cells Int
i = Vector Word8 -> Int -> Word8
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex (Cells -> Vector Word8
cValid Cells
cells) Int
i Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0
{-# INLINE validAt #-}

{- | @%Y-%m-%d@ takes the WS-B slice fast path; custom formats keep the
reference parser over a materialized cell.
-}
dateSliceParser :: String -> Cells -> Int -> Maybe Day
dateSliceParser :: String -> Cells -> Int -> Maybe Day
dateSliceParser String
"%Y-%m-%d" Cells
cells Int
i = Cells
-> Int -> (ByteString -> Int -> Int -> Maybe Day) -> Maybe Day
forall r. Cells -> Int -> (ByteString -> Int -> Int -> r) -> r
withCell Cells
cells Int
i ByteString -> Int -> Int -> Maybe Day
parseDateFieldSlice
dateSliceParser String
fmt Cells
cells Int
i = String -> ByteString -> Maybe Day
readByteStringDate String
fmt (Cells -> Int -> ByteString
cellAt Cells
cells Int
i)
{-# INLINE dateSliceParser #-}

{- | 'EitherRead' wrap for the ByteString inference path: produce an
@Either Text a@ column. Raw input bytes are preserved verbatim; rows that were
marked invalid by the builder (e.g. empty cells before EitherRead disabled
missing-indicator detection) become @Left \"\"@.
-}
handleBSEither :: String -> ParsingAssumption -> Cells -> Column
handleBSEither :: String -> ParsingAssumption -> Cells -> Column
handleBSEither String
dfmt ParsingAssumption
assumption Cells
cells = case ParsingAssumption
assumption of
    ParsingAssumption
BoolAssumption -> (ByteString -> Maybe Bool) -> Column
forall a. Columnable a => (ByteString -> Maybe a) -> Column
wrap ByteString -> Maybe Bool
readByteStringBool
    ParsingAssumption
IntAssumption -> (ByteString -> Maybe Int) -> Column
forall a. Columnable a => (ByteString -> Maybe a) -> Column
wrap HasCallStack => ByteString -> Maybe Int
ByteString -> Maybe Int
readByteStringInt
    ParsingAssumption
DoubleAssumption -> (ByteString -> Maybe Double) -> Column
forall a. Columnable a => (ByteString -> Maybe a) -> Column
wrap HasCallStack => ByteString -> Maybe Double
ByteString -> Maybe Double
readByteStringDouble
    ParsingAssumption
DateAssumption -> (ByteString -> Maybe Day) -> Column
forall a. Columnable a => (ByteString -> Maybe a) -> Column
wrap (String -> ByteString -> Maybe Day
readByteStringDate String
dfmt)
    -- Text / No assumption: column is Either Text Text; empty cells become
    -- Left "" to keep the "Left means missing/failure" convention.
    ParsingAssumption
TextAssumption -> Vector (Either Text Text) -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector (Int -> (Int -> Either Text Text) -> Vector (Either Text Text)
forall a. Int -> (Int -> a) -> Vector a
V.generate (Cells -> Int
cLen Cells
cells) Int -> Either Text Text
textEither)
    ParsingAssumption
NoAssumption -> Vector (Either Text Text) -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector (Int -> (Int -> Either Text Text) -> Vector (Either Text Text)
forall a. Int -> (Int -> a) -> Vector a
V.generate (Cells -> Int
cLen Cells
cells) Int -> Either Text Text
textEither)
  where
    wrap ::
        forall a. (Columnable a) => (BS.ByteString -> Maybe a) -> Column
    wrap :: forall a. Columnable a => (ByteString -> Maybe a) -> Column
wrap ByteString -> Maybe a
p = Vector (Either Text a) -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector (Int -> (Int -> Either Text a) -> Vector (Either Text a)
forall a. Int -> (Int -> a) -> Vector a
V.generate (Cells -> Int
cLen Cells
cells) ((ByteString -> Maybe a) -> Int -> Either Text a
forall a. (ByteString -> Maybe a) -> Int -> Either Text a
toEither ByteString -> Maybe a
p))

    toEither ::
        forall a. (BS.ByteString -> Maybe a) -> Int -> Either T.Text a
    toEither :: forall a. (ByteString -> Maybe a) -> Int -> Either Text a
toEither ByteString -> Maybe a
p Int
i
        | Bool -> Bool
not (Cells -> Int -> Bool
validAt Cells
cells Int
i) = Text -> Either Text a
forall a b. a -> Either a b
Left (ByteString -> Text
TE.decodeUtf8Lenient ByteString
bs)
        | Bool
otherwise = case ByteString -> Maybe a
p ByteString
bs of
            Just a
v -> a -> Either Text a
forall a b. b -> Either a b
Right a
v
            Maybe a
Nothing -> Text -> Either Text a
forall a b. a -> Either a b
Left (ByteString -> Text
TE.decodeUtf8Lenient ByteString
bs)
      where
        bs :: ByteString
bs = Cells -> Int -> ByteString
cellAt Cells
cells Int
i

    textEither :: Int -> Either T.Text T.Text
    textEither :: Int -> Either Text Text
textEither Int
i =
        let t :: Text
t = ByteString -> Text
TE.decodeUtf8Lenient (Cells -> Int -> ByteString
cellAt Cells
cells Int
i)
         in if Bool -> Bool
not (Cells -> Int -> Bool
validAt Cells
cells Int
i) Bool -> Bool -> Bool
|| Text -> Bool
T.null Text
t then Text -> Either Text Text
forall a b. a -> Either a b
Left Text
t else Text -> Either Text Text
forall a b. b -> Either a b
Right Text
t

handleBSBool :: Cells -> Maybe Column
handleBSBool :: Cells -> Maybe Column
handleBSBool Cells
cells =
    (Maybe (Vector Word8) -> Vector Bool -> Column)
-> (Maybe (Vector Word8), Vector Bool) -> Column
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Maybe (Vector Word8) -> Vector Bool -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe (Vector Word8) -> Vector a -> Column
UnboxedColumn
        ((Maybe (Vector Word8), Vector Bool) -> Column)
-> Maybe (Maybe (Vector Word8), Vector Bool) -> Maybe Column
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool
-> (Int -> Maybe Bool)
-> Cells
-> Maybe (Maybe (Vector Word8), Vector Bool)
forall a.
Unbox a =>
a
-> (Int -> Maybe a)
-> Cells
-> Maybe (Maybe (Vector Word8), Vector a)
parseUnboxedCells Bool
False (\Int
i -> Cells
-> Int -> (ByteString -> Int -> Int -> Maybe Bool) -> Maybe Bool
forall r. Cells -> Int -> (ByteString -> Int -> Int -> r) -> r
withCell Cells
cells Int
i ByteString -> Int -> Int -> Maybe Bool
parseBoolFieldSlice) Cells
cells

{- | Int columns: one fused pass with in-place Int -> Double promotion
('promoteIntColumnIndexed', audit T2) instead of the full-column Double
re-parse; both parsers are the WS-B slice parsers (bit-exact with the
old @readByteString*@ pair).
-}
handleBSInt :: Cells -> Maybe Column
handleBSInt :: Cells -> Maybe Column
handleBSInt Cells
cells =
    Int
-> (Int -> Bool)
-> (Int -> Maybe Int)
-> (Int -> Maybe Double)
-> Maybe Column
promoteIntColumnIndexed
        (Cells -> Int
cLen Cells
cells)
        (Bool -> Bool
not (Bool -> Bool) -> (Int -> Bool) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Cells -> Int -> Bool
validAt Cells
cells)
        (\Int
i -> Cells
-> Int -> (ByteString -> Int -> Int -> Maybe Int) -> Maybe Int
forall r. Cells -> Int -> (ByteString -> Int -> Int -> r) -> r
withCell Cells
cells Int
i ByteString -> Int -> Int -> Maybe Int
parseIntFieldSlice)
        (\Int
i -> Cells
-> Int
-> (ByteString -> Int -> Int -> Maybe Double)
-> Maybe Double
forall r. Cells -> Int -> (ByteString -> Int -> Int -> r) -> r
withCell Cells
cells Int
i ByteString -> Int -> Int -> Maybe Double
parseDoubleFieldSlice)

handleBSDouble :: Cells -> Maybe Column
handleBSDouble :: Cells -> Maybe Column
handleBSDouble Cells
cells =
    (Maybe (Vector Word8) -> Vector Double -> Column)
-> (Maybe (Vector Word8), Vector Double) -> Column
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Maybe (Vector Word8) -> Vector Double -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe (Vector Word8) -> Vector a -> Column
UnboxedColumn
        ((Maybe (Vector Word8), Vector Double) -> Column)
-> Maybe (Maybe (Vector Word8), Vector Double) -> Maybe Column
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Double
-> (Int -> Maybe Double)
-> Cells
-> Maybe (Maybe (Vector Word8), Vector Double)
forall a.
Unbox a =>
a
-> (Int -> Maybe a)
-> Cells
-> Maybe (Maybe (Vector Word8), Vector a)
parseUnboxedCells Double
0 (\Int
i -> Cells
-> Int
-> (ByteString -> Int -> Int -> Maybe Double)
-> Maybe Double
forall r. Cells -> Int -> (ByteString -> Int -> Int -> r) -> r
withCell Cells
cells Int
i ByteString -> Int -> Int -> Maybe Double
parseDoubleFieldSlice) Cells
cells

-- Dates are boxed ('Day' isn't 'VU.Unbox'), so fuse into a V.Vector
-- (Maybe Day) in one pass.  Bails on the first non-null cell that
-- fails to parse.
handleBSDate :: String -> Cells -> Maybe Column
handleBSDate :: String -> Cells -> Maybe Column
handleBSDate String
dfmt Cells
cells =
    case (Int -> Maybe Day) -> Cells -> Maybe (Bool, Vector (Maybe Day))
forall a.
(Int -> Maybe a) -> Cells -> Maybe (Bool, Vector (Maybe a))
parseBoxedMaybeCells (String -> Cells -> Int -> Maybe Day
dateSliceParser String
dfmt Cells
cells) Cells
cells of
        Just (Bool
anyNull, Vector (Maybe Day)
out)
            | Bool
anyNull -> Column -> Maybe Column
forall a. a -> Maybe a
Just (Vector (Maybe Day) -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector Vector (Maybe Day)
out)
            | Bool
otherwise -> Column -> Maybe Column
forall a. a -> Maybe a
Just (Vector Day -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector ((Maybe Day -> Maybe Day) -> Vector (Maybe Day) -> Vector Day
forall a b. (a -> Maybe b) -> Vector a -> Vector b
V.mapMaybe Maybe Day -> Maybe Day
forall a. a -> a
id Vector (Maybe Day)
out))
        Maybe (Bool, Vector (Maybe Day))
Nothing -> Maybe Column
forall a. Maybe a
Nothing

{- | Text columns build through the WS-C 'TextBuilder': one shared byte
array + offsets, 'Data.Text.Text' values sliced off it at freeze (UTF-8
validated once, lenient per-field decode only when invalid).
-}
textColumn :: Cells -> IO Column
textColumn :: Cells -> IO Column
textColumn Cells
cells = do
    let n :: Int
n = Cells -> Int
cLen Cells
cells
    TextBuilder RealWorld
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
n (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
8))
    ByteString -> (CStringLen -> IO ()) -> IO ()
forall a. ByteString -> (CStringLen -> IO a) -> IO a
BSU.unsafeUseAsCStringLen (Cells -> ByteString
cBS Cells
cells) ((CStringLen -> IO ()) -> IO ()) -> (CStringLen -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
p, Int
_) -> do
        let base :: Ptr Word8
base = Ptr CChar -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr CChar
p :: Ptr Word8
            go :: Int -> IO ()
go !Int
i
                | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                | Bool
otherwise = do
                    let s :: Int
s = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex (Cells -> Vector Int
cOffs Cells
cells) (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i)
                        e :: Int
e = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex (Cells -> Vector Int
cOffs Cells
cells) (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    if Bool -> Bool
not (Cells -> Int -> Bool
validAt Cells
cells Int
i)
                        then ST RealWorld () -> IO ()
forall a. ST RealWorld a -> IO a
stToIO (TextBuilder RealWorld -> ST RealWorld ()
forall s. TextBuilder s -> ST s ()
forall (b :: * -> *) s. ColumnBuilder b => b s -> ST s ()
appendNull TextBuilder RealWorld
b)
                        else
                            if Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
                                then ST RealWorld () -> IO ()
forall a. ST RealWorld a -> IO a
stToIO (TextBuilder RealWorld -> Ptr Word8 -> Int -> ST RealWorld ()
forall s. TextBuilder s -> Ptr Word8 -> Int -> ST s ()
appendTextSliceFromPtr TextBuilder RealWorld
b (Ptr Word8
base Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
s) (Int
e Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
s))
                                else
                                    let cell :: ByteString
cell = Cells -> IntMap ByteString
cEsc Cells
cells IntMap ByteString -> Int -> ByteString
forall a. IntMap a -> Int -> a
IM.! Int
e
                                     in ByteString -> (CStringLen -> IO ()) -> IO ()
forall a. ByteString -> (CStringLen -> IO a) -> IO a
BSU.unsafeUseAsCStringLen ByteString
cell ((CStringLen -> IO ()) -> IO ()) -> (CStringLen -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
q, Int
l) ->
                                            ST RealWorld () -> IO ()
forall a. ST RealWorld a -> IO a
stToIO
                                                (TextBuilder RealWorld -> Ptr Word8 -> Int -> ST RealWorld ()
forall s. TextBuilder s -> Ptr Word8 -> Int -> ST s ()
appendTextSliceFromPtr TextBuilder RealWorld
b (Ptr CChar -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr CChar
q) Int
l)
                    Int -> IO ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        Int -> IO ()
go Int
0
    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)

handleBSNo :: String -> Cells -> IO Column
handleBSNo :: String -> Cells -> IO Column
handleBSNo String
dfmt Cells
cells
    | (Word8 -> Bool) -> Vector Word8 -> Bool
forall a. Unbox a => (a -> Bool) -> Vector a -> Bool
VU.all (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0) (Cells -> Vector Word8
cValid Cells
cells) =
        Column -> IO Column
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector (Maybe Text) -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
fromVector (Int -> Maybe Text -> Vector (Maybe Text)
forall a. Int -> a -> Vector a
V.replicate (Cells -> Int
cLen Cells
cells) (Maybe Text
forall a. Maybe a
Nothing :: Maybe T.Text)))
    | Just (Maybe (Vector Word8)
mbm, Vector Bool
out) <-
        Bool
-> (Int -> Maybe Bool)
-> Cells
-> Maybe (Maybe (Vector Word8), Vector Bool)
forall a.
Unbox a =>
a
-> (Int -> Maybe a)
-> Cells
-> Maybe (Maybe (Vector Word8), Vector a)
parseUnboxedCells Bool
False (\Int
i -> Cells
-> Int -> (ByteString -> Int -> Int -> Maybe Bool) -> Maybe Bool
forall r. Cells -> Int -> (ByteString -> Int -> Int -> r) -> r
withCell Cells
cells Int
i ByteString -> Int -> Int -> Maybe Bool
parseBoolFieldSlice) Cells
cells =
        Column -> IO Column
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Vector Word8) -> Vector Bool -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe (Vector Word8) -> Vector a -> Column
UnboxedColumn Maybe (Vector Word8)
mbm Vector Bool
out)
    | Just (Maybe (Vector Word8)
mbm, Vector Int
out) <-
        Int
-> (Int -> Maybe Int)
-> Cells
-> Maybe (Maybe (Vector Word8), Vector Int)
forall a.
Unbox a =>
a
-> (Int -> Maybe a)
-> Cells
-> Maybe (Maybe (Vector Word8), Vector a)
parseUnboxedCells Int
0 (\Int
i -> Cells
-> Int -> (ByteString -> Int -> Int -> Maybe Int) -> Maybe Int
forall r. Cells -> Int -> (ByteString -> Int -> Int -> r) -> r
withCell Cells
cells Int
i ByteString -> Int -> Int -> Maybe Int
parseIntFieldSlice) Cells
cells =
        Column -> IO Column
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Vector Word8) -> Vector Int -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe (Vector Word8) -> Vector a -> Column
UnboxedColumn Maybe (Vector Word8)
mbm Vector Int
out)
    | Just (Maybe (Vector Word8)
mbm, Vector Double
out) <-
        Double
-> (Int -> Maybe Double)
-> Cells
-> Maybe (Maybe (Vector Word8), Vector Double)
forall a.
Unbox a =>
a
-> (Int -> Maybe a)
-> Cells
-> Maybe (Maybe (Vector Word8), Vector a)
parseUnboxedCells Double
0 (\Int
i -> Cells
-> Int
-> (ByteString -> Int -> Int -> Maybe Double)
-> Maybe Double
forall r. Cells -> Int -> (ByteString -> Int -> Int -> r) -> r
withCell Cells
cells Int
i ByteString -> Int -> Int -> Maybe Double
parseDoubleFieldSlice) Cells
cells =
        Column -> IO Column
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Vector Word8) -> Vector Double -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe (Vector Word8) -> Vector a -> Column
UnboxedColumn Maybe (Vector Word8)
mbm Vector Double
out)
    | Bool
otherwise = IO Column -> (Column -> IO Column) -> Maybe Column -> IO Column
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Cells -> IO Column
textColumn Cells
cells) Column -> IO Column
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Cells -> Maybe Column
handleBSDate String
dfmt Cells
cells)

-- Boxed counterpart to 'parseUnboxedCells' for types that aren't
-- 'VU.Unbox' (e.g. 'Day').  Same one-pass + early-bail shape.
parseBoxedMaybeCells ::
    (Int -> Maybe a) -> Cells -> Maybe (Bool, V.Vector (Maybe a))
parseBoxedMaybeCells :: forall a.
(Int -> Maybe a) -> Cells -> Maybe (Bool, Vector (Maybe a))
parseBoxedMaybeCells Int -> Maybe a
parser Cells
cells = (forall s. ST s (Maybe (Bool, Vector (Maybe a))))
-> Maybe (Bool, Vector (Maybe a))
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Maybe (Bool, Vector (Maybe a))))
 -> Maybe (Bool, Vector (Maybe a)))
-> (forall s. ST s (Maybe (Bool, Vector (Maybe a))))
-> Maybe (Bool, Vector (Maybe a))
forall a b. (a -> b) -> a -> b
$ do
    let n :: Int
n = Cells -> Int
cLen Cells
cells
    MVector s (Maybe a)
out <- Int -> ST s (MVector (PrimState (ST s)) (Maybe a))
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
VM.new Int
n
    let loop :: Int -> Bool -> ST s (Maybe (Bool, Vector (Maybe a)))
loop !Int
i !Bool
anyNull
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = do
                Vector (Maybe a)
frozen <- MVector (PrimState (ST s)) (Maybe a) -> ST s (Vector (Maybe a))
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.unsafeFreeze MVector s (Maybe a)
MVector (PrimState (ST s)) (Maybe a)
out
                Maybe (Bool, Vector (Maybe a))
-> ST s (Maybe (Bool, Vector (Maybe a)))
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Bool, Vector (Maybe a)) -> Maybe (Bool, Vector (Maybe a))
forall a. a -> Maybe a
Just (Bool
anyNull, Vector (Maybe a)
frozen))
            | Bool -> Bool
not (Cells -> Int -> Bool
validAt Cells
cells Int
i) = do
                MVector (PrimState (ST s)) (Maybe a) -> Int -> Maybe a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
VM.unsafeWrite MVector s (Maybe a)
MVector (PrimState (ST s)) (Maybe a)
out Int
i Maybe a
forall a. Maybe a
Nothing
                Int -> Bool -> ST s (Maybe (Bool, Vector (Maybe a)))
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bool
True
            | Bool
otherwise = case Int -> Maybe a
parser Int
i of
                Just a
v -> do
                    MVector (PrimState (ST s)) (Maybe a) -> Int -> Maybe a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
VM.unsafeWrite MVector s (Maybe a)
MVector (PrimState (ST s)) (Maybe a)
out Int
i (a -> Maybe a
forall a. a -> Maybe a
Just a
v)
                    Int -> Bool -> ST s (Maybe (Bool, Vector (Maybe a)))
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bool
anyNull
                Maybe a
Nothing -> Maybe (Bool, Vector (Maybe a))
-> ST s (Maybe (Bool, Vector (Maybe a)))
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Bool, Vector (Maybe a))
forall a. Maybe a
Nothing
    Int -> Bool -> ST s (Maybe (Bool, Vector (Maybe a)))
loop Int
0 Bool
False

{- | One-pass fused parse into a typed unboxed column, reading
nullability from the builder's validity vector. Returns @Just@ only
when every non-null cell parses; the first unparseable cell
short-circuits so the caller can fall back to the next assumption.
-}
parseUnboxedCells ::
    forall a.
    (VU.Unbox a) =>
    a ->
    (Int -> Maybe a) ->
    Cells ->
    Maybe (Maybe Bitmap, VU.Vector a)
parseUnboxedCells :: forall a.
Unbox a =>
a
-> (Int -> Maybe a)
-> Cells
-> Maybe (Maybe (Vector Word8), Vector a)
parseUnboxedCells a
nullValue Int -> Maybe a
parser Cells
cells = (forall s. ST s (Maybe (Maybe (Vector Word8), Vector a)))
-> Maybe (Maybe (Vector Word8), Vector a)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Maybe (Maybe (Vector Word8), Vector a)))
 -> Maybe (Maybe (Vector Word8), Vector a))
-> (forall s. ST s (Maybe (Maybe (Vector Word8), Vector a)))
-> Maybe (Maybe (Vector Word8), Vector a)
forall a b. (a -> b) -> a -> b
$ do
    let n :: Int
n = Cells -> Int
cLen Cells
cells
    STVector s a
values <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
n
    STVector s Word8
vmask <- Int -> ST s (MVector (PrimState (ST s)) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
n
    let go :: Int -> Bool -> ST s (Maybe (Maybe (Vector Word8), Vector a))
go !Int
i !Bool
anyNull
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = STVector s a
-> STVector s Word8
-> Bool
-> ST s (Maybe (Maybe (Vector Word8), Vector a))
forall a s.
Unbox a =>
STVector s a
-> STVector s Word8
-> Bool
-> ST s (Maybe (Maybe (Vector Word8), Vector a))
finalizeParseResult STVector s a
values STVector s Word8
vmask Bool
anyNull
            | Bool -> Bool
not (Cells -> Int -> Bool
validAt Cells
cells Int
i) = do
                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 STVector s Word8
MVector (PrimState (ST s)) Word8
vmask Int
i Word8
0
                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 STVector s a
MVector (PrimState (ST s)) a
values Int
i a
nullValue
                Int -> Bool -> ST s (Maybe (Maybe (Vector Word8), Vector a))
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bool
True
            | Bool
otherwise = case Int -> Maybe a
parser Int
i of
                Just a
v -> do
                    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 STVector s Word8
MVector (PrimState (ST s)) Word8
vmask Int
i Word8
1
                    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 STVector s a
MVector (PrimState (ST s)) a
values Int
i a
v
                    Int -> Bool -> ST s (Maybe (Maybe (Vector Word8), Vector a))
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bool
anyNull
                Maybe a
Nothing -> Maybe (Maybe (Vector Word8), Vector a)
-> ST s (Maybe (Maybe (Vector Word8), Vector a))
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Maybe (Vector Word8), Vector a)
forall a. Maybe a
Nothing
    Int -> Bool -> ST s (Maybe (Maybe (Vector Word8), Vector a))
go Int
0 Bool
False
{-# INLINE parseUnboxedCells #-}