dataframe-fastcsv-1.4.1.2: src/DataFrame/IO/CSV/Fast/Core.hs

{-# LANGUAGE OverloadedStrings #-}

-- | The shared read pipeline behind every fast CSV entry point.
module DataFrame.IO.CSV.Fast.Core (
    readSeparatedCore,
    checkSeparatorAgrees,
) where

import qualified Data.Map as M
import qualified Data.Set as S
import qualified Data.Text as Text
import qualified Data.Vector as Vector
import qualified Data.Vector.Storable as VS

import Control.Exception (evaluate, throwIO)
import Control.Monad (unless, when)
import Data.Char (ord)
import qualified Data.Text as T
import Data.Word (Word8)
import Debug.Trace (traceMarkerIO)

import DataFrame.IO.CSV (
    HeaderSpec (..),
    InvalidReadOptions (..),
    RaggedRowPolicy (..),
    ReadOptions (..),
    UnclosedQuotePolicy (..),
    resolveSelection,
    schemaTypeMap,
    validateReadOptions,
 )
import DataFrame.IO.CSV.Fast.Index (
    CsvParseError (..),
    byteStringView,
    checkNoRaggedRows,
    classifyRowEnds,
    collectDataRows,
    fieldsInRow,
    getDelimiterIndicesPolicy,
 )
import DataFrame.IO.CSV.Fast.IndexPar (
    ParIndex (..),
    getDelimiterIndicesPar,
 )
import DataFrame.IO.CSV.Fast.Parallel (autoChunkCount, buildAllColumns)
import DataFrame.IO.CSV.Fast.Passes (ColumnEnv (..))
import DataFrame.IO.CSV.Fast.Slice (
    FieldCtx (..),
    FieldLayout (..),
    compactRows,
    extractField,
    stripBom,
 )
import DataFrame.Internal.DataFrame (DataFrame (..), forceDataFrame)
import DataFrame.Operations.Typing (effectiveSafeRead, parseWithTypes)

-- | Run the read pipeline; @chunkSpec@ pins the chunk count if given.
readSeparatedCore ::
    Maybe Int ->
    Word8 ->
    ReadOptions ->
    VS.Vector Word8 ->
    IO DataFrame
readSeparatedCore chunkSpec separator opts fileRaw = do
    validateReadOptions opts
    checkSeparatorAgrees separator opts
    let (file, _bomLen) = stripBom fileRaw
        contentLen = VS.length file
    nChunks <- maybe (autoChunkCount contentLen) pure chunkSpec
    -- Parallel reads scan in parallel too.
    mPar <-
        if nChunks > 1
            then getDelimiterIndicesPar nChunks separator contentLen file
            else pure Nothing
    (indices, rowEnds) <- case mPar of
        Just par -> do
            when
                ( piEndsInQuote par
                    && fastCsvOnUnclosedQuote opts == RaiseOnUnclosedQuote
                )
                (throwIO CsvUnclosedQuote)
            pure (piDelims par, piRowEnds par)
        Nothing -> do
            is <-
                getDelimiterIndicesPolicy
                    (fastCsvOnUnclosedQuote opts)
                    separator
                    contentLen
                    file
            pure (is, classifyRowEnds file contentLen is)
    -- Eventlog phase markers (visible under +RTS -l; no-ops otherwise).
    traceMarkerIO "fastcsv:scan-done"
    let totalRows = VS.length rowEnds
    if totalRows == 0
        then return emptyDataFrame
        else do
            let ctx =
                    FieldCtx
                        { fcFile = file
                        , fcBS = byteStringView contentLen file
                        , fcLayout = LayoutFlat indices rowEnds
                        , fcContentLen = contentLen
                        , fcTrim = fastCsvTrimUnquoted opts
                        }
                headerNumCol = fieldsInRow rowEnds 0
                (headerNames, dataStartRow, numCol) = case headerSpec opts of
                    NoHeader ->
                        ( map (Text.pack . show) [0 .. headerNumCol - 1]
                        , 0
                        , headerNumCol
                        )
                    UseFirstRow ->
                        ( map (extractField ctx 0) [0 .. headerNumCol - 1]
                        , 1
                        , headerNumCol
                        )
                    ProvideNames ns ->
                        (ns, 0, length ns)
                wanted = resolveSelection opts headerNames
                columnNames = Vector.fromList (map fst wanted)
                numKept = length wanted
            if numCol == 0
                then return emptyDataFrame
                else do
                    let dataRows =
                            collectDataRows
                                file
                                indices
                                rowEnds
                                contentLen
                                dataStartRow
                                totalRows
                        numRow = VS.length dataRows
                    when (fastCsvOnRaggedRow opts == RaiseOnRagged) $
                        checkNoRaggedRows rowEnds dataRows numCol
                    traceMarkerIO "fastcsv:rows-done"
                    mapM_ (\(nm, ix) -> evaluate nm >> evaluate ix) wanted
                    ctx' <-
                        evaluate $
                            case compactRows contentLen indices rowEnds numCol of
                                Just lay -> ctx{fcLayout = lay}
                                Nothing -> ctx
                    traceMarkerIO "fastcsv:index-compacted"
                    VS.unsafeWith file $ \filePtr -> do
                        let env =
                                ColumnEnv
                                    { ceCtx = ctx'
                                    , ceRows = dataRows
                                    , ceNumRow = numRow
                                    , cePtr = filePtr
                                    , ceTextHint =
                                        max 64 (contentLen `div` numCol)
                                    , ceOpts = opts
                                    }
                        results <- buildAllColumns nChunks env wanted
                        traceMarkerIO "fastcsv:build-done"
                        let cols = Vector.fromListN numKept (map fst results)
                            legacySchema =
                                S.fromList
                                    [ columnNames Vector.! c
                                    | (c, (_, True)) <- zip [0 ..] results
                                    ]
                            colIndices =
                                M.fromList $
                                    zip (Vector.toList columnNames) [0 ..]
                            df =
                                DataFrame
                                    cols
                                    colIndices
                                    (numRow, numKept)
                                    M.empty
                            resolveMode =
                                effectiveSafeRead
                                    (safeRead opts)
                                    (safeReadOverrides opts)
                            df' =
                                if S.null legacySchema
                                    then df
                                    else
                                        parseWithTypes
                                            resolveMode
                                            ( M.restrictKeys
                                                (schemaTypeMap (typeSpec opts))
                                                legacySchema
                                            )
                                            df
                        -- Parallel merge already forced the payloads.
                        if nChunks > 1 && S.null legacySchema
                            then do
                                Vector.foldl' (\() c -> c `seq` ()) () cols `seq`
                                    return df
                            else return $! forceDataFrame df'

-- | Reject a 'columnSeparator' that disagrees with the separator byte.
checkSeparatorAgrees :: Word8 -> ReadOptions -> IO ()
checkSeparatorAgrees separator opts =
    unless (separator == fromIntegral (ord asked)) $
        throwIO . InvalidReadOptions $
            [ "columnSeparator is "
                <> T.pack (show asked)
                <> " but this reader was given "
                <> T.pack (show (toEnum (fromIntegral separator) :: Char))
            ]
  where
    asked = columnSeparator opts

-- | An empty 'DataFrame', for input with no delimiters.
emptyDataFrame :: DataFrame
emptyDataFrame = DataFrame Vector.empty M.empty (0, 0) M.empty