CSV.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module DataFrame.IO.CSV where

import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as C
import qualified Data.ByteString.Lazy as BL
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.Text.IO as TIO
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 Data.Csv.Streaming (Records (..))
import qualified Data.Csv.Streaming as CsvStream

import Control.Exception (SomeException, catch)
import Control.Monad
import Control.Monad.ST (runST)
import Data.Char
import qualified Data.Csv as Csv
import Data.Either
import Data.Functor
import Data.IORef
import Data.Maybe
import Data.Type.Equality (TestEquality (testEquality))
import Data.Word (Word8)
import DataFrame.Internal.Column
import DataFrame.Internal.DataFrame (
    DataFrame (..),
    forceDataFrame,
    toSeparated,
 )
import DataFrame.Internal.Parsing
import DataFrame.Internal.Schema
import DataFrame.Operations.Typing
import System.IO
import Type.Reflection
import Prelude hiding (concat, takeWhile)

chunkSize :: Int
chunkSize = 16_384

data PagedVector a = PagedVector
    { pvChunks :: !(IORef [V.Vector a])
    -- ^ Finished chunks (reverse order)
    , pvActive :: !(IORef (VM.IOVector a))
    -- ^ Current mutable chunk
    , pvCount :: !(IORef Int)
    -- ^ Items written in current chunk
    }

data PagedUnboxedVector a = PagedUnboxedVector
    { puvChunks :: !(IORef [VU.Vector a])
    , puvActive :: !(IORef (VUM.IOVector a))
    , puvCount :: !(IORef Int)
    }

data BuilderColumn
    = BuilderInt !(PagedUnboxedVector Int) !(PagedUnboxedVector Word8)
    | BuilderDouble !(PagedUnboxedVector Double) !(PagedUnboxedVector Word8)
    | BuilderText !(PagedVector T.Text) !(PagedUnboxedVector Word8)
    | BuilderBS !(PagedVector BS.ByteString) !(PagedUnboxedVector Word8)

newPagedVector :: IO (PagedVector a)
newPagedVector = do
    active <- VM.unsafeNew chunkSize
    PagedVector <$> newIORef [] <*> newIORef active <*> newIORef 0

newPagedUnboxedVector :: (VUM.Unbox a) => IO (PagedUnboxedVector a)
newPagedUnboxedVector = do
    active <- VUM.unsafeNew chunkSize
    PagedUnboxedVector <$> newIORef [] <*> newIORef active <*> newIORef 0

appendPagedVector :: PagedVector a -> a -> IO ()
appendPagedVector (PagedVector chunksRef activeRef countRef) !val = do
    count <- readIORef countRef
    active <- readIORef activeRef

    if count < chunkSize
        then do
            VM.unsafeWrite active count val
            writeIORef countRef $! count + 1
        else do
            frozen <- V.unsafeFreeze active
            modifyIORef' chunksRef (frozen :)

            newActive <- VM.unsafeNew chunkSize
            VM.unsafeWrite newActive 0 val

            writeIORef activeRef newActive
            writeIORef countRef 1
{-# INLINE appendPagedVector #-}

appendPagedUnboxedVector :: (VUM.Unbox a) => PagedUnboxedVector a -> a -> IO ()
appendPagedUnboxedVector (PagedUnboxedVector chunksRef activeRef countRef) !val = do
    count <- readIORef countRef
    active <- readIORef activeRef

    if count < chunkSize
        then do
            VUM.unsafeWrite active count val
            writeIORef countRef $! count + 1
        else do
            frozen <- VU.unsafeFreeze active
            modifyIORef' chunksRef (frozen :)

            newActive <- VUM.unsafeNew chunkSize
            VUM.unsafeWrite newActive 0 val

            writeIORef activeRef newActive
            writeIORef countRef 1
{-# INLINE appendPagedUnboxedVector #-}

freezePagedVector :: PagedVector a -> IO (V.Vector a)
freezePagedVector (PagedVector chunksRef activeRef countRef) = do
    count <- readIORef countRef
    active <- readIORef activeRef
    chunks <- readIORef chunksRef

    writeIORef chunksRef [] -- release chunk references
    let frozenChunks = reverse chunks
        totalLen = count + sum (map V.length frozenChunks)

    mv <- VM.unsafeNew totalLen

    let copyChunk !offset chunk = do
            V.copy (VM.slice offset (V.length chunk) mv) chunk
            pure (offset + V.length chunk)

    offset <- foldM copyChunk 0 frozenChunks
    VM.copy (VM.slice offset count mv) (VM.slice 0 count active)

    V.unsafeFreeze mv

freezePagedUnboxedVector ::
    (VUM.Unbox a) => PagedUnboxedVector a -> IO (VU.Vector a)
freezePagedUnboxedVector (PagedUnboxedVector chunksRef activeRef countRef) = do
    count <- readIORef countRef
    active <- readIORef activeRef
    chunks <- readIORef chunksRef

    writeIORef chunksRef [] -- release chunk references
    let frozenChunks = reverse chunks
        totalLen = count + sum (map VU.length frozenChunks)

    mv <- VUM.unsafeNew totalLen

    let copyChunk !offset chunk = do
            VU.copy (VUM.slice offset (VU.length chunk) mv) chunk
            pure (offset + VU.length chunk)

    offset <- foldM copyChunk 0 frozenChunks
    VUM.copy (VUM.slice offset count mv) (VUM.slice 0 count active)

    VU.unsafeFreeze mv

-- | STANDARD CONFIG TYPES
data HeaderSpec = NoHeader | UseFirstRow | ProvideNames [T.Text]
    deriving (Eq, Show)

data TypeSpec
    = InferFromSample Int
    | SpecifyTypes [(T.Text, SchemaType)] TypeSpec
    | NoInference

{- | How the fast reader should treat a row whose field count does not
match the header row.  Only consulted by @dataframe-fastcsv@; the pure
Haskell reader has its own semantics.
-}
data RaggedRowPolicy
    = {- | Fill missing cells with nulls; silently drop extras.
      Matches pandas / polars lenient defaults.
      -}
      PadWithNull
    | {- | Fill missing cells with nulls; silently drop extras.  Alias
      kept for ergonomic naming when the caller only cares about the
      "don't raise, just forget the extras" half of 'PadWithNull'.
      -}
      Truncate
    | {- | Raise a 'DataFrame.IO.CSV.Fast.CsvParseError' on any row whose
      field count differs from the header.  Matches polars' strict
      schema-bound mode.
      -}
      RaiseOnRagged
    deriving (Eq, Show)

-- | How the fast reader should treat an unclosed quoted field at EOF.
data UnclosedQuotePolicy
    = -- | Raise a 'DataFrame.IO.CSV.Fast.CsvParseError'.  Default.
      RaiseOnUnclosedQuote
    | {- | Return whatever rows were parsed before the stray quote; the
      remainder is silently dropped.
      -}
      BestEffort
    deriving (Eq, Show)

-- | CSV read parameters.
data ReadOptions = ReadOptions
    { headerSpec :: HeaderSpec
    -- ^ Where to get the headers from. (default: UseFirstRow)
    , typeSpec :: TypeSpec
    -- ^ Whether/how to infer types. (default: InferFromSample 100)
    , safeRead :: SafeReadMode
    {- ^ Default 'SafeReadMode' for columns without an entry in
    'safeReadOverrides'. (default: 'NoSafeRead')
    -}
    , safeReadOverrides :: [(T.Text, SafeReadMode)]
    -- ^ Per-column 'SafeReadMode' overrides; takes precedence over 'safeRead'.
    , dateFormat :: String
    {- ^ Format of date fields as recognized by the Data.Time.Format module.

    __Examples:__

    @
    > parseTimeM True defaultTimeLocale "%Y/%-m/%-d" "2010/3/04" :: Maybe Day
    Just 2010-03-04
    > parseTimeM True defaultTimeLocale "%d/%-m/%-Y" "04/3/2010" :: Maybe Day
    Just 2010-03-04
    @
    -}
    , columnSeparator :: Char
    -- ^ Character that separates column values.
    , numColumns :: Maybe Int
    -- ^ Number of columns to read.
    , missingIndicators :: [T.Text]
    -- ^ Values that should be read as `Nothing`.
    , fastCsvOnRaggedRow :: RaggedRowPolicy
    {- ^ @dataframe-fastcsv@: how to treat rows with a non-header field count.
    (default: 'PadWithNull')
    -}
    , fastCsvOnUnclosedQuote :: UnclosedQuotePolicy
    {- ^ @dataframe-fastcsv@: how to treat an unclosed quoted field at EOF.
    (default: 'RaiseOnUnclosedQuote')
    -}
    , fastCsvTrimUnquoted :: Bool
    {- ^ @dataframe-fastcsv@: if 'True', leading/trailing whitespace is
    stripped from unquoted fields after decoding.  RFC 4180 preserves
    this whitespace, and that is the default ('False').
    -}
    }

shouldInferFromSample :: TypeSpec -> Bool
shouldInferFromSample (InferFromSample _) = True
shouldInferFromSample (SpecifyTypes _ fallback) = shouldInferFromSample fallback
shouldInferFromSample _ = False

schemaTypeMap :: TypeSpec -> M.Map T.Text SchemaType
schemaTypeMap (SpecifyTypes xs _) = M.fromList xs
schemaTypeMap _ = M.empty

typeInferenceSampleSize :: TypeSpec -> Int
typeInferenceSampleSize (InferFromSample n) = n
typeInferenceSampleSize (SpecifyTypes _ fallback) = typeInferenceSampleSize fallback
typeInferenceSampleSize _ = 0

defaultReadOptions :: ReadOptions
defaultReadOptions =
    ReadOptions
        { headerSpec = UseFirstRow
        , typeSpec = InferFromSample 100
        , safeRead = NoSafeRead
        , safeReadOverrides = []
        , dateFormat = "%Y-%m-%d"
        , columnSeparator = ','
        , numColumns = Nothing
        , missingIndicators =
            ["Nothing", "NULL", "", " ", "nan", "null", "N/A", "NaN", "NAN", "NA"]
        , fastCsvOnRaggedRow = PadWithNull
        , fastCsvOnUnclosedQuote = RaiseOnUnclosedQuote
        , fastCsvTrimUnquoted = False
        }

{- | Read CSV file from path and load it into a dataframe.

==== __Example__
@
ghci> D.readCsv ".\/data\/taxi.csv"

@
-}
readCsv :: FilePath -> IO DataFrame
readCsv = readSeparated defaultReadOptions

type CsvReader = Schema -> FilePath -> IO DataFrame

{- | Schema-driven attoparsec CSV reader.  Coerces each column to the
type declared in 'Schema'; columns absent from the schema fall back to
the default inference path.  Defined in terms of 'readSeparated' with
the 'TypeSpec' filled in.

@
import qualified DataFrame as D
df <- D.readCsvWithSchema schema "input.csv"
@
-}
readCsvWithSchema :: CsvReader
readCsvWithSchema schema =
    readSeparated
        defaultReadOptions
            { typeSpec =
                SpecifyTypes
                    (M.toList (elements schema))
                    (typeSpec defaultReadOptions)
            }

{- | Read CSV file from path and load it into a dataframe.

==== __Example__
@
ghci> D.readCsvWithOpts ".\/data\/taxi.csv" (D.defaultReadOptions { dateFormat = "%d/%-m/%-Y" })

@
-}
readCsvWithOpts :: ReadOptions -> FilePath -> IO DataFrame
readCsvWithOpts = readSeparated

{- | Read TSV (tab separated) file from path and load it into a dataframe.

==== __Example__
@
ghci> D.readTsv ".\/data\/taxi.tsv"

@
-}
readTsv :: FilePath -> IO DataFrame
readTsv = readSeparated (defaultReadOptions{columnSeparator = '\t'})

{- | Read text file with specified delimiter into a dataframe.

==== __Example__
@
ghci> D.readSeparated (D.defaultReadOptions { columnSeparator = ';' }) ".\/data\/taxi.txt"

@
-}
readSeparated :: ReadOptions -> FilePath -> IO DataFrame
readSeparated opts !path = do
    let stripUtf8Bom bs = fromMaybe bs (BL.stripPrefix "\xEF\xBB\xBF" bs)
    csvData <- stripUtf8Bom <$> BL.readFile path
    fmap forceDataFrame (decodeSeparated opts csvData)

decodeSeparated :: ReadOptions -> BL.ByteString -> IO DataFrame
decodeSeparated !opts csvData = do
    let sep = columnSeparator opts
    let decodeOpts = Csv.defaultDecodeOptions{Csv.decDelimiter = fromIntegral (ord sep)}
    let stream = CsvStream.decodeWith decodeOpts Csv.NoHeader csvData

    let peekStream (Cons (Right row) rest) = return (row, rest)
        peekStream (Cons (Left err) _) = error $ "Error parsing CSV header: " ++ err
        peekStream (Nil Nothing _) = error "Empty CSV file"
        peekStream (Nil (Just err) _) = error err

    (firstRowRaw, dataStream) <- peekStream stream

    let (columnNames, rowsToProcess) = case headerSpec opts of
            NoHeader ->
                ( map (T.pack . show) [0 .. V.length firstRowRaw - 1]
                , Cons (Right firstRowRaw) dataStream
                )
            UseFirstRow ->
                ( map (T.strip . TE.decodeUtf8Lenient . BL.toStrict) (V.toList firstRowRaw)
                , dataStream
                )
            ProvideNames ns ->
                ( ns ++ drop (length ns) (map (T.pack . show) [0 .. V.length firstRowRaw - 1])
                , Cons (Right firstRowRaw) dataStream
                )

    (sampleRow, _) <- peekStream rowsToProcess
    builderCols <- initializeColumns columnNames (V.toList sampleRow) opts
    let !builderColsV = V.fromList builderCols
    let colNamesV = V.fromList columnNames
        resolveMode =
            effectiveSafeRead
                (safeRead opts)
                (safeReadOverrides opts)
        -- If ANY column is EitherRead we keep every raw cell (including
        -- "N/A" etc.) verbatim; otherwise the missing-indicator list applies.
        anyEither =
            any (\n -> resolveMode n == EitherRead) columnNames
        missing = if anyEither then [] else missingIndicators opts
    processStream missing rowsToProcess builderColsV (numColumns opts)

    frozenCols <-
        V.zipWithM
            (\name bc -> finalizeBuilderColumn (resolveMode name) opts bc)
            colNamesV
            builderColsV
    let numRows = maybe 0 columnLength (frozenCols V.!? 0)

    let df =
            DataFrame
                frozenCols
                (M.fromList (zip columnNames [0 ..]))
                (numRows, V.length frozenCols)
                M.empty -- TODO give typed column references
    pure $ parseWithTypes resolveMode (schemaTypeMap (typeSpec opts)) df

initializeColumns ::
    [T.Text] -> [BL.ByteString] -> ReadOptions -> IO [BuilderColumn]
initializeColumns names _row opts = zipWithM initColumn names (map lookupType names)
  where
    typeMap = schemaTypeMap (typeSpec opts)
    -- Return Nothing for columns that should be inferred from BS
    shouldInfer = case typeSpec opts of
        InferFromSample _ -> True
        SpecifyTypes _ fallback -> shouldInferFromSample fallback
        NoInference -> False
    lookupType name = M.lookup name typeMap
    resolveMode =
        effectiveSafeRead (safeRead opts) (safeReadOverrides opts)
    initColumn :: T.Text -> Maybe SchemaType -> IO BuilderColumn
    initColumn name _ | resolveMode name == EitherRead = do
        validityRef <- newPagedUnboxedVector
        BuilderBS <$> newPagedVector <*> pure validityRef
    initColumn _ Nothing | shouldInfer = do
        validityRef <- newPagedUnboxedVector
        BuilderBS <$> newPagedVector <*> pure validityRef
    initColumn _ mtype = do
        validityRef <- newPagedUnboxedVector
        let t = fromMaybe (schemaType @T.Text) mtype
        case t of
            SType (_ :: P.Proxy a) -> case testEquality (typeRep @a) (typeRep @Int) of
                Just Refl -> BuilderInt <$> newPagedUnboxedVector <*> pure validityRef
                Nothing -> case testEquality (typeRep @a) (typeRep @Double) of
                    Just Refl -> BuilderDouble <$> newPagedUnboxedVector <*> pure validityRef
                    Nothing -> BuilderText <$> newPagedVector <*> pure validityRef

processStream ::
    [T.Text] ->
    CsvStream.Records (V.Vector BL.ByteString) ->
    V.Vector BuilderColumn ->
    Maybe Int ->
    IO ()
processStream _ _ _ (Just 0) = return ()
processStream missing (Cons (Right row) rest) cols n =
    processRow missing row cols
        >> processStream missing rest cols (fmap (flip (-) 1) n)
processStream _missing (Cons (Left err) _) _ _ = error ("CSV Parse Error: " ++ err)
processStream _missing (Nil _ _) _ _ = return ()

processRow ::
    [T.Text] -> V.Vector BL.ByteString -> V.Vector BuilderColumn -> IO ()
processRow missing !vals !cols = V.zipWithM_ processValue vals cols
  where
    processValue !bs !col = do
        let !bs' = BL.toStrict bs
        case col of
            BuilderInt gv valid -> case readByteStringInt bs' of
                Just !i -> appendPagedUnboxedVector gv i >> appendPagedUnboxedVector valid 1
                Nothing -> appendPagedUnboxedVector gv 0 >> appendPagedUnboxedVector valid 0
            BuilderDouble gv valid -> case readByteStringDouble bs' of
                Just !d -> appendPagedUnboxedVector gv d >> appendPagedUnboxedVector valid 1
                Nothing -> appendPagedUnboxedVector gv 0.0 >> appendPagedUnboxedVector valid 0
            BuilderText gv valid -> do
                let !val = T.strip (TE.decodeUtf8Lenient bs')
                appendPagedVector gv val
                appendPagedUnboxedVector valid (if val `elem` missing then 0 else 1)
            BuilderBS gv valid -> do
                let !bs'' = C.strip bs'
                appendPagedVector gv bs''
                appendPagedUnboxedVector
                    valid
                    (if TE.decodeUtf8Lenient bs'' `elem` missing then 0 else 1)

freezeBuilderColumn :: BuilderColumn -> IO Column
freezeBuilderColumn (BuilderInt gv validRef) = do
    vec <- freezePagedUnboxedVector gv
    valid <- freezePagedUnboxedVector validRef
    if VU.all (== 1) valid
        then return $! UnboxedColumn Nothing vec
        else constructOptional vec valid
freezeBuilderColumn (BuilderDouble gv validRef) = do
    vec <- freezePagedUnboxedVector gv
    valid <- freezePagedUnboxedVector validRef
    if VU.all (== 1) valid
        then return $! UnboxedColumn Nothing vec
        else constructOptional vec valid
freezeBuilderColumn (BuilderText gv validRef) = do
    vec <- freezePagedVector gv
    valid <- freezePagedUnboxedVector validRef
    if VU.all (== 1) valid
        then return $! BoxedColumn Nothing vec
        else constructOptionalBoxed vec valid
freezeBuilderColumn (BuilderBS _ _) =
    error
        "freezeBuilderColumn: BuilderBS must be finalized via finalizeBuilderColumn"

finalizeBuilderColumn ::
    SafeReadMode -> ReadOptions -> BuilderColumn -> IO Column
finalizeBuilderColumn mode opts bc = do
    col <- case bc of
        BuilderBS gv validRef -> do
            vec <- freezePagedVector gv
            valid <- freezePagedUnboxedVector validRef
            return $! inferColumnFromBS mode opts vec valid
        _ -> freezeBuilderColumn bc
    return $! case mode of
        NoSafeRead -> col
        MaybeRead -> ensureOptional col
        EitherRead -> col

inferColumnFromBS ::
    SafeReadMode ->
    ReadOptions ->
    V.Vector BS.ByteString ->
    VU.Vector Word8 ->
    Column
inferColumnFromBS mode opts vec valid =
    let sampleN = let n = typeInferenceSampleSize (typeSpec opts) in if n == 0 then 100 else n
        dfmt = dateFormat opts
        -- The sample IS still Maybe-wrapped; it's bounded at 100 rows
        -- by default, so the allocation is ignorable.  The previous
        -- full-column `asMaybeFull = V.generate ...` allocation is
        -- gone — handlers walk (vec, valid) directly.
        samples = V.generate (min sampleN (V.length vec)) $ \i ->
            if valid VU.! i == 1 then Just (vec V.! i) else Nothing
        assumption = makeParsingAssumptionBS dfmt samples
     in case mode of
            EitherRead -> handleBSEither dfmt assumption vec valid
            _ -> case assumption of
                IntAssumption -> handleBSInt dfmt vec valid
                DoubleAssumption -> handleBSDouble vec valid
                BoolAssumption -> handleBSBool vec valid
                DateAssumption -> handleBSDate dfmt vec valid
                TextAssumption -> handleBSText vec valid
                NoAssumption -> handleBSNo dfmt vec valid

{- | '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 ->
    V.Vector BS.ByteString ->
    VU.Vector Word8 ->
    Column
handleBSEither dfmt assumption vec valid = case assumption of
    BoolAssumption -> wrap readByteStringBool
    IntAssumption -> wrap readByteStringInt
    DoubleAssumption -> wrap readByteStringDouble
    DateAssumption -> wrap (readByteStringDate dfmt)
    -- Text / No assumption: column is Either Text Text; empty cells become
    -- Left "" to keep the "Left means missing/failure" convention.
    TextAssumption -> fromVector (V.imap textEither vec)
    NoAssumption -> fromVector (V.imap textEither vec)
  where
    wrap ::
        forall a. (Columnable a) => (BS.ByteString -> Maybe a) -> Column
    wrap p = fromVector (V.imap (toEither p) vec)

    toEither ::
        forall a.
        (BS.ByteString -> Maybe a) ->
        Int ->
        BS.ByteString ->
        Either T.Text a
    toEither p i bs
        | valid VU.! i == 0 = Left (TE.decodeUtf8Lenient bs)
        | otherwise = case p bs of
            Just v -> Right v
            Nothing -> Left (TE.decodeUtf8Lenient bs)

    textEither :: Int -> BS.ByteString -> Either T.Text T.Text
    textEither i bs =
        let t = TE.decodeUtf8Lenient bs
         in if valid VU.! i == 0 || T.null t then Left t else Right t

makeParsingAssumptionBS ::
    String -> V.Vector (Maybe BS.ByteString) -> ParsingAssumption
makeParsingAssumptionBS dfmt asMaybe
    | V.all (== Nothing) asMaybe = NoAssumption
    | vecSameConstructor asMaybe asMaybeBool = BoolAssumption
    | vecSameConstructor asMaybe asMaybeInt
        && vecSameConstructor asMaybe asMaybeDouble =
        IntAssumption
    | vecSameConstructor asMaybe asMaybeDouble = DoubleAssumption
    | vecSameConstructor asMaybe asMaybeDate = DateAssumption
    | otherwise = TextAssumption
  where
    asMaybeBool = V.map (>>= readByteStringBool) asMaybe
    asMaybeInt = V.map (>>= readByteStringInt) asMaybe
    asMaybeDouble = V.map (>>= readByteStringDouble) asMaybe
    asMaybeDate = V.map (>>= readByteStringDate dfmt) asMaybe

-- All @handleBS*@ helpers now take the raw @V.Vector BS.ByteString@
-- plus the builder's @VU.Vector Word8@ validity vector, fusing the
-- parse + validity check into a single pass via
-- 'parseUnboxedColumnWithValid'.  The previous @V.Vector (Maybe
-- BS.ByteString)@ intermediate — allocated upstream in
-- 'inferColumnFromBS' — is gone, along with the paired Int/Double
-- parses and the two 'V.zipWith' Bool vectors from
-- 'vecSameConstructor'.

handleBSBool ::
    V.Vector BS.ByteString -> VU.Vector Word8 -> Column
handleBSBool vec valid =
    case parseUnboxedColumnWithValid False readByteStringBool vec valid of
        Just (mbm, out) -> UnboxedColumn mbm out
        Nothing -> handleBSText vec valid

handleBSInt ::
    String -> V.Vector BS.ByteString -> VU.Vector Word8 -> Column
handleBSInt _dfmt vec valid =
    case parseUnboxedColumnWithValid 0 readByteStringInt vec valid of
        Just (mbm, out) -> UnboxedColumn mbm out
        Nothing -> case parseUnboxedColumnWithValid 0 readByteStringDouble vec valid of
            Just (mbm, out) -> UnboxedColumn mbm out
            Nothing -> handleBSText vec valid

handleBSDouble ::
    V.Vector BS.ByteString -> VU.Vector Word8 -> Column
handleBSDouble vec valid =
    case parseUnboxedColumnWithValid 0 readByteStringDouble vec valid of
        Just (mbm, out) -> UnboxedColumn mbm out
        Nothing -> handleBSText vec valid

-- 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 -> V.Vector BS.ByteString -> VU.Vector Word8 -> Column
handleBSDate dfmt vec valid =
    case parseBoxedMaybeBSColumn valid (readByteStringDate dfmt) vec of
        Just (anyNull, out)
            | anyNull -> fromVector out
            | otherwise -> fromVector (V.mapMaybe id out)
        Nothing -> handleBSText vec valid

-- Fused Text handler: decode each cell's UTF-8 bytes once, mark nulls
-- directly from the validity vector.  Replaces the two-pass
-- `V.map (fmap decodeUtf8Lenient) ... sequenceA ...` pattern.
handleBSText ::
    V.Vector BS.ByteString -> VU.Vector Word8 -> Column
handleBSText vec valid
    | VU.any (== 0) valid =
        fromVector
            ( V.imap
                ( \i bs ->
                    if valid VU.! i == 0
                        then Nothing
                        else Just (TE.decodeUtf8Lenient bs)
                )
                vec
            )
    | otherwise = fromVector (V.map TE.decodeUtf8Lenient vec)

handleBSNo ::
    String -> V.Vector BS.ByteString -> VU.Vector Word8 -> Column
handleBSNo dfmt vec valid
    | VU.all (== 0) valid =
        fromVector (V.map (const (Nothing :: Maybe T.Text)) vec)
    | Just (mbm, out) <-
        parseUnboxedColumnWithValid False readByteStringBool vec valid =
        UnboxedColumn mbm out
    | Just (mbm, out) <- parseUnboxedColumnWithValid 0 readByteStringInt vec valid =
        UnboxedColumn mbm out
    | Just (mbm, out) <- parseUnboxedColumnWithValid 0 readByteStringDouble vec valid =
        UnboxedColumn mbm out
    | otherwise = case parseBoxedMaybeBSColumn valid (readByteStringDate dfmt) vec of
        Just (anyNull, out)
            | anyNull -> fromVector out
            | otherwise -> fromVector (V.mapMaybe id out)
        Nothing -> handleBSText vec valid

-- Boxed counterpart to 'parseUnboxedColumnWithValid' for types that
-- aren't 'VU.Unbox' (e.g. 'Day').  Same one-pass + early-bail shape.
parseBoxedMaybeBSColumn ::
    VU.Vector Word8 ->
    (BS.ByteString -> Maybe a) ->
    V.Vector BS.ByteString ->
    Maybe (Bool, V.Vector (Maybe a))
parseBoxedMaybeBSColumn valid parser vec = runST $ do
    let n = V.length vec
    out <- VM.new n
    let loop !i !anyNull
            | i >= n = do
                frozen <- V.unsafeFreeze out
                return (Just (anyNull, frozen))
            | VU.unsafeIndex valid i == 0 = do
                VM.unsafeWrite out i Nothing
                loop (i + 1) True
            | otherwise = case parser (V.unsafeIndex vec i) of
                Just v -> do
                    VM.unsafeWrite out i (Just v)
                    loop (i + 1) anyNull
                Nothing -> return Nothing
    loop 0 False

{- | One-pass fused parse into a typed unboxed column.  Avoids the
@V.Vector (Maybe a)@ intermediate that the "parse then 'sequenceA'"
idiom requires, and avoids the upstream @V.Vector (Maybe src)@
classification by reading nullability from a precomputed validity
vector (produced by the CSV builder alongside the raw cells).

Returns @Just (mbm, vec)@ only when every non-null cell parses.  The
first unparseable cell short-circuits to @Nothing@ so the caller can
fall back to the next assumption (Int → Double → Text).  @mbm@ is
@Nothing@ when no nulls exist (the column is non-nullable) and
@Just bm@ otherwise.

Null slots are filled with @nullValue@; downstream consumers only see
them through the bitmap, so the sentinel never escapes.

Memory shape for a length-@n@ input:

  * 1 × 'VUM.STVector' of @a@        (final data, @sizeOf a × n@ bytes)
  * 1 × 'VUM.STVector' of 'Word8'    (per-element validity, @n@ bytes)
  * 1 × 'Bitmap' (bit-packed, @⌈n\/8⌉@ bytes) — only when nulls exist.
-}
parseUnboxedColumnWithValid ::
    forall src a.
    (VU.Unbox a) =>
    a ->
    (src -> Maybe a) ->
    V.Vector src ->
    VU.Vector Word8 ->
    Maybe (Maybe Bitmap, VU.Vector a)
parseUnboxedColumnWithValid nullValue parser vec valid = runST $ do
    let n = V.length vec
    values <- VUM.unsafeNew n
    vmask <- VUM.unsafeNew n
    let go !i !anyNull
            | i >= n = finalizeParseResult values vmask anyNull
            | VU.unsafeIndex valid i == 0 = do
                VUM.unsafeWrite vmask i 0
                VUM.unsafeWrite values i nullValue
                go (i + 1) True
            | otherwise = case parser (V.unsafeIndex vec i) of
                Just v -> do
                    VUM.unsafeWrite vmask i 1
                    VUM.unsafeWrite values i v
                    go (i + 1) anyNull
                Nothing -> return Nothing
    go 0 False
{-# INLINE parseUnboxedColumnWithValid #-}

constructOptional ::
    (VU.Unbox a, Columnable a) => VU.Vector a -> VU.Vector Word8 -> IO Column
constructOptional vec valid = do
    let bm = buildBitmapFromValid valid
    pure $ UnboxedColumn (Just bm) vec

constructOptionalBoxed :: V.Vector T.Text -> VU.Vector Word8 -> IO Column
constructOptionalBoxed vec valid = do
    let bm = buildBitmapFromValid valid
    pure $ BoxedColumn (Just bm) vec

writeCsv :: FilePath -> DataFrame -> IO ()
writeCsv = writeSeparated ','

writeTsv :: FilePath -> DataFrame -> IO ()
writeTsv = writeSeparated '\t'

writeSeparated ::
    -- | Separator
    Char ->
    -- | Path to write to
    FilePath ->
    DataFrame ->
    IO ()
writeSeparated c filepath df = TIO.writeFile filepath (toSeparated c df)

-- | Parse a CSV string into a DataFrame using default options.
fromCsv :: String -> IO (Either String DataFrame)
fromCsv s = do
    let bs = BL.fromStrict (TE.encodeUtf8 (T.pack s))
    (Right <$> decodeSeparated defaultReadOptions bs)
        `catch` (\(e :: SomeException) -> pure (Left (show e)))

-- | Parse a lazy 'ByteString' containing CSV data into a DataFrame using default options.
fromCsvBytes :: BL.ByteString -> IO DataFrame
fromCsvBytes = decodeSeparated defaultReadOptions

stripQuotes :: T.Text -> T.Text
stripQuotes txt =
    case T.uncons txt of
        Just ('"', rest) ->
            case T.unsnoc rest of
                Just (middle, '"') -> middle
                _ -> txt
        _ -> txt