dataframe-csv 1.0.2.0 → 2.3.1.1
raw patch · 10 files changed
Files
- dataframe-csv.cabal +13/−11
- src-internal/DataFrame/IO/Internal/MutableColumn.hs +84/−0
- src/DataFrame/IO/CSV.hs +37/−57
- src/DataFrame/IO/CSV/Internal/Infer.hs +34/−42
- src/DataFrame/IO/CSV/Internal/Options.hs +99/−56
- src/DataFrame/IO/CSV/Internal/Read.hs +95/−90
- src/DataFrame/IO/CSV/Internal/Scanner.hs +7/−28
- src/DataFrame/IO/CSV/Internal/Sink.hs +7/−27
- src/DataFrame/IO/Internal/MutableColumn.hs +0/−89
- src/DataFrame/Typed/IO/CSV.hs +78/−0
dataframe-csv.cabal view
@@ -1,6 +1,6 @@-cabal-version: 2.4+cabal-version: 3.4 name: dataframe-csv-version: 1.0.2.0+version: 2.3.1.1 synopsis: CSV reader and writer for the dataframe ecosystem. description: @DataFrame.IO.CSV@ — strict single-pass CSV read/write (pure@@ -28,20 +28,22 @@ import: warnings exposed-modules: DataFrame.IO.CSV+ DataFrame.Typed.IO.CSV+ DataFrame.IO.Internal.MutableColumn+ other-modules: DataFrame.IO.CSV.Internal.Infer DataFrame.IO.CSV.Internal.Options DataFrame.IO.CSV.Internal.Read DataFrame.IO.CSV.Internal.Scanner DataFrame.IO.CSV.Internal.Sink- DataFrame.IO.Internal.MutableColumn build-depends: base >= 4 && < 5,- bytestring >= 0.11 && < 0.13,- containers >= 0.6.7 && < 0.9,- dataframe-core ^>= 1.1,- dataframe-operations ^>= 1.1.1,- dataframe-parsing ^>= 1.0.2,- text >= 2.0 && < 3,+ bytestring >= 0.11 && < 0.14,+ containers >= 0.6.7 && < 0.10,+ dataframe-core >= 2.6 && < 2.7,+ dataframe-operations >= 2.6 && < 2.7,+ dataframe-parsing >= 2.2 && < 2.3,+ text >= 2.1 && < 3, time >= 1.12 && < 2,- vector ^>= 0.13- hs-source-dirs: src+ vector >= 0.13 && < 0.15+ hs-source-dirs: src, src-internal default-language: Haskell2010
+ src-internal/DataFrame/IO/Internal/MutableColumn.hs view
@@ -0,0 +1,84 @@+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | Mutable-column ingest helpers used by the CSV readers.+module DataFrame.IO.Internal.MutableColumn (+ writeColumn,+ freezeColumn',+) where++import qualified Data.Text as T+import qualified Data.Vector as VB+import qualified Data.Vector.Mutable as VBM+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM++import Data.Maybe (fromMaybe)+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))+import DataFrame.Internal.Column (+ Column (..),+ MutableColumn (..),+ )+import DataFrame.Internal.Column.Bitmap (buildBitmapFromNulls)+import DataFrame.Internal.Parsing (isNullish, readDouble, readInt)+import Type.Reflection (typeRep)++writeColumn :: Int -> T.Text -> MutableColumn -> IO (Either T.Text Bool)+writeColumn i value (MBoxedColumn (col :: VBM.IOVector a)) =+ case testEquality (typeRep @a) (typeRep @T.Text) of+ Just Refl ->+ if isNullish value+ then VBM.unsafeWrite col i "" >> return (Left $! value)+ else VBM.unsafeWrite col i value >> return (Right True)+ Nothing -> return (Left value)+writeColumn i value (MUnboxedColumn (col :: VUM.IOVector a)) =+ case testEquality (typeRep @a) (typeRep @Int) of+ Just Refl -> case readInt value of+ Just v -> VUM.unsafeWrite col i v >> return (Right True)+ Nothing -> VUM.unsafeWrite col i 0 >> return (Left value)+ Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+ Nothing -> return (Left $! value)+ Just Refl -> case readDouble value of+ Just v -> VUM.unsafeWrite col i v >> return (Right True)+ Nothing -> VUM.unsafeWrite col i 0 >> return (Left $! value)+{-# INLINE writeColumn #-}++freezeColumn' :: [(Int, T.Text)] -> MutableColumn -> IO Column+freezeColumn' nulls (MBoxedColumn col)+ | null nulls = BoxedColumn Nothing <$> VB.unsafeFreeze col+ | all (isNullish . snd) nulls = do+ frozen <- VB.unsafeFreeze col+ let n = VB.length frozen+ bm = buildBitmapFromNulls n (map fst nulls)+ return $ BoxedColumn (Just bm) frozen+ | otherwise =+ BoxedColumn Nothing+ . VB.imap+ ( \i v ->+ if i `elem` map fst nulls+ then Left (fromMaybe (error "UNEXPECTED ERROR DURING FREEZE") (lookup i nulls))+ else Right v+ )+ <$> VB.unsafeFreeze col+freezeColumn' nulls (MUnboxedColumn col)+ | null nulls = UnboxedColumn Nothing <$> VU.unsafeFreeze col+ | all (isNullish . snd) nulls = do+ c <- VU.unsafeFreeze col+ let n = VU.length c+ bm = buildBitmapFromNulls n (map fst nulls)+ return $ UnboxedColumn (Just bm) c+ | otherwise = do+ c <- VU.unsafeFreeze col+ return $+ BoxedColumn Nothing $+ VB.generate+ (VU.length c)+ ( \i ->+ if i `elem` map fst nulls+ then Left (fromMaybe (error "UNEXPECTED ERROR DURING FREEZE") (lookup i nulls))+ else Right (c VU.! i)+ )+{-# INLINE freezeColumn' #-}
src/DataFrame/IO/CSV.hs view
@@ -1,12 +1,7 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} -{- | CSV reading and writing for dataframes. The reader is a strict,-single-pass RFC 4180 scanner (cassava-compatible semantics, pinned by-the golden suite in @tests/IO/CsvGolden.hs@) that parses fields straight-into typed column builders. Ragged rows are padded with nulls (short-rows) and extra fields dropped.--}+-- | Strict single-pass RFC 4180 CSV reading and writing for dataframes. module DataFrame.IO.CSV ( -- * Reading readCsv,@@ -15,7 +10,10 @@ readSeparated, readCsvWithSchema, CsvReader,+ CsvBytesReader,+ schemaReadOptions, decodeSeparated,+ decodeSeparatedStrict, fromCsv, fromCsvBytes, @@ -36,39 +34,35 @@ import qualified Data.Map.Strict as M import qualified Data.Text as T import qualified Data.Text.Encoding as TE-import qualified Data.Text.IO as TIO+import qualified Data.Text.IO.Utf8 as TIO import Control.Exception (SomeException, catch) import Data.Maybe (fromMaybe) import DataFrame.IO.CSV.Internal.Options import DataFrame.IO.CSV.Internal.Read (decodeCsvStrict) import DataFrame.Internal.DataFrame (DataFrame (..), toSeparated)-import DataFrame.Internal.Schema (Schema, elements)--{- | Read CSV file from path and load it into a dataframe.--==== __Example__-@-ghci> D.readCsv ".\/data\/taxi.csv"+import DataFrame.Schema (Schema, elements) -@--}+-- | Read a CSV file from a path into a dataframe. readCsv :: FilePath -> IO DataFrame readCsv = readSeparated defaultReadOptions -type CsvReader = Schema -> FilePath -> IO DataFrame+-- | A reader the lazy scan can drive, such as 'readSeparated'.+type CsvReader = ReadOptions -> FilePath -> IO DataFrame -{- | Schema-driven 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.+type CsvBytesReader = ReadOptions -> BS.ByteString -> IO DataFrame -@-import qualified DataFrame as D-df <- D.readCsvWithSchema schema "input.csv"-@--}-readCsvWithSchema :: CsvReader+-- | Read options derived from a 'Schema', which sets types and columns.+schemaReadOptions :: Schema -> ReadOptions+schemaReadOptions schema =+ defaultReadOptions+ { typeSpec =+ SpecifyTypes (M.toList (elements schema)) (typeSpec defaultReadOptions)+ , readColumns = Just (M.keys (elements schema))+ }++-- | Read a CSV, coercing columns to the types declared in a 'Schema'.+readCsvWithSchema :: Schema -> FilePath -> IO DataFrame readCsvWithSchema schema = readSeparated defaultReadOptions@@ -78,54 +72,39 @@ (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" })--@--}+-- | Read a CSV file from a path using the given 'ReadOptions'. 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"--@--}+-- | Read a tab-separated file into a dataframe. 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"--@--}+-- | Read a text file with a chosen delimiter into a dataframe. readSeparated :: ReadOptions -> FilePath -> IO DataFrame readSeparated opts path = do+ validateReadOptions opts let stripUtf8Bom b = fromMaybe b (BS.stripPrefix "\xEF\xBB\xBF" b) csvData <- stripUtf8Bom <$> BS.readFile path- decodeCsvStrict opts csvData+ decodeSeparatedStrict opts csvData -{- | Decode in-memory CSV bytes into a dataframe. The result is fully-forced. (Note: unlike 'readSeparated', no UTF-8 BOM is stripped.)--}+-- | Decode in-memory CSV bytes into a fully forced dataframe. decodeSeparated :: ReadOptions -> BL.ByteString -> IO DataFrame-decodeSeparated opts csvData = decodeCsvStrict opts (BL.toStrict csvData)+decodeSeparated opts csvData = decodeSeparatedStrict opts (BL.toStrict csvData) +-- | Decode a strict in-memory CSV buffer into a fully forced dataframe.+decodeSeparatedStrict :: CsvBytesReader+decodeSeparatedStrict = decodeCsvStrict++-- | Write a dataframe to a comma-separated file. writeCsv :: FilePath -> DataFrame -> IO () writeCsv = writeSeparated ',' +-- | Write a dataframe to a tab-separated file. writeTsv :: FilePath -> DataFrame -> IO () writeTsv = writeSeparated '\t' +-- | Write a dataframe to a file using the given field separator. writeSeparated :: -- | Separator Char ->@@ -142,10 +121,11 @@ (Right <$> decodeSeparated defaultReadOptions bs) `catch` (\(e :: SomeException) -> pure (Left (show e))) --- | Parse a lazy 'ByteString' containing CSV data into a DataFrame using default options.+-- | Parse a lazy 'ByteString' of CSV data using default options. fromCsvBytes :: BL.ByteString -> IO DataFrame fromCsvBytes = decodeSeparated defaultReadOptions +-- | Strip one layer of double-quote delimiters from a field, if present. stripQuotes :: T.Text -> T.Text stripQuotes txt = case T.uncons txt of
src/DataFrame/IO/CSV/Internal/Infer.hs view
@@ -3,13 +3,7 @@ {-# 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.--}+-- | Column inference over raw CSV field bytes for the default reader. module DataFrame.IO.CSV.Internal.Infer ( inferColumnFromBS, ) where@@ -29,24 +23,43 @@ 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.Column (+ Column (UnboxedColumn),+ Columnable,+ finalizeParseResult,+ fromVector,+ )+import DataFrame.Internal.Column.Bitmap (Bitmap)+import DataFrame.Internal.Column.Builder (+ ColumnBuilder (appendNull, freezeBuilder),+ appendTextSliceFromPtr,+ newTextBuilder,+ )+import DataFrame.Internal.Parsing (+ readByteStringBool,+ readByteStringDate,+ readByteStringDouble,+ readByteStringInt,+ ) import DataFrame.Internal.Parsing.Fast ( parseBoolFieldSlice, parseDateFieldSlice, parseDoubleFieldSlice, parseIntFieldSlice, )-import DataFrame.Operations.Typing+import DataFrame.Operations.Typing (+ ParsingAssumption (..),+ SafeReadMode (EitherRead),+ makeParsingAssumptionBytes,+ promoteIntColumnIndexed,+ ) import Foreign.Ptr (Ptr, castPtr, plusPtr) inferColumnFromBS :: SafeReadMode -> ReadOptions -> Cells -> IO Column inferColumnFromBS mode opts cells = 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 sample is Maybe-wrapped but capped at 100 rows by default. samples = V.generate (min sampleN (cLen cells)) $ \i -> if validAt cells i then Just (cellAt cells i) else Nothing assumption = makeParsingAssumptionBytes dfmt samples@@ -65,27 +78,20 @@ validAt cells i = VU.unsafeIndex (cValid cells) i /= 0 {-# INLINE validAt #-} -{- | @%Y-%m-%d@ takes the WS-B slice fast path; custom formats keep the-reference parser over a materialized cell.--}+-- | Fast slice path for @%Y-%m-%d@; other formats use the reference parser. dateSliceParser :: String -> Cells -> Int -> Maybe Day dateSliceParser "%Y-%m-%d" cells i = withCell cells i parseDateFieldSlice dateSliceParser fmt cells i = readByteStringDate fmt (cellAt cells 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 \"\"@.--}+-- | Wrap a ByteString-inferred column as @Either Text a@, keeping raw bytes. handleBSEither :: String -> ParsingAssumption -> Cells -> Column handleBSEither dfmt assumption cells = 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.+ -- Empty cells become Left "" ("Left means missing/failure"). TextAssumption -> fromVector (V.generate (cLen cells) textEither) NoAssumption -> fromVector (V.generate (cLen cells) textEither) where@@ -113,11 +119,7 @@ uncurry UnboxedColumn <$> parseUnboxedCells False (\i -> withCell cells i parseBoolFieldSlice) 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).--}+-- | Int columns: one fused pass with in-place Int to Double promotion. handleBSInt :: Cells -> Maybe Column handleBSInt cells = promoteIntColumnIndexed@@ -131,9 +133,7 @@ uncurry UnboxedColumn <$> parseUnboxedCells 0 (\i -> withCell cells i parseDoubleFieldSlice) 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.+-- Dates are boxed, so fuse into a V.Vector (Maybe Day); bails on a bad cell. handleBSDate :: String -> Cells -> Maybe Column handleBSDate dfmt cells = case parseBoxedMaybeCells (dateSliceParser dfmt cells) cells of@@ -142,10 +142,7 @@ | otherwise -> Just (fromVector (V.mapMaybe id out)) Nothing -> 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).--}+-- | Text columns build through a shared-array 'TextBuilder'. textColumn :: Cells -> IO Column textColumn cells = do let n = cLen cells@@ -186,8 +183,7 @@ pure (UnboxedColumn mbm out) | otherwise = maybe (textColumn cells) pure (handleBSDate dfmt cells) --- Boxed counterpart to 'parseUnboxedCells' for types that aren't--- 'VU.Unbox' (e.g. 'Day'). Same one-pass + early-bail shape.+-- Boxed counterpart to 'parseUnboxedCells' for non-Unbox types (e.g. 'Day'). parseBoxedMaybeCells :: (Int -> Maybe a) -> Cells -> Maybe (Bool, V.Vector (Maybe a)) parseBoxedMaybeCells parser cells = runST $ do@@ -207,11 +203,7 @@ Nothing -> return Nothing loop 0 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.--}+-- | Fused parse into a typed unboxed column; Just only if every cell parses. parseUnboxedCells :: forall a. (VU.Unbox a) =>
src/DataFrame/IO/CSV/Internal/Options.hs view
@@ -1,9 +1,6 @@ {-# LANGUAGE OverloadedStrings #-} -{- | Configuration types for the CSV readers. Re-exported through-"DataFrame.IO.CSV"; shared by the default reader, @dataframe-fastcsv@-and @dataframe-lazy@.--}+-- | Configuration types for the CSV readers. module DataFrame.IO.CSV.Internal.Options ( HeaderSpec (..), TypeSpec (..),@@ -14,40 +11,47 @@ shouldInferFromSample, schemaTypeMap, typeInferenceSampleSize,+ resolveSelection,+ InvalidReadOptions (..),+ readOptionsErrors,+ validateReadOptions, ) where +import qualified Data.List as L import qualified Data.Map.Strict as M import qualified Data.Text as T -import DataFrame.Internal.Schema (SchemaType)+import Control.Exception (Exception, throw, throwIO)+import Control.Monad (unless)+import DataFrame.Errors (DataFrameException (ColumnsNotFoundException)) import DataFrame.Operations.Typing (SafeReadMode (..))+import DataFrame.Schema (SchemaType) -data HeaderSpec = NoHeader | UseFirstRow | ProvideNames [T.Text]+-- | Where a reader's column names come from.+data HeaderSpec+ = NoHeader+ | -- | Take names from the first row (the default).+ UseFirstRow+ | -- | Name the leading columns; any beyond the given list are numbered.+ ProvideNames [T.Text] deriving (Eq, Show) +-- | How a reader decides each column's type. data TypeSpec- = InferFromSample Int- | SpecifyTypes [(T.Text, SchemaType)] TypeSpec- | NoInference+ = -- | Infer each column's type from the first @n@ rows.+ InferFromSample Int+ | -- | Pin the listed columns to their given types, with a fallback.+ SpecifyTypes [(T.Text, SchemaType)] TypeSpec+ | -- | Every column is read as 'Data.Text.Text', no inference at all.+ 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 always pads short rows with nulls and drops extras.--}+-- | How the fast reader treats a row with the wrong field count. data RaggedRowPolicy- = {- | Fill missing cells with nulls; silently drop extras.- Matches pandas / polars lenient defaults.- -}+ = -- | Fill missing cells with nulls and silently drop extras. 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'.- -}+ | -- | Alias of 'PadWithNull' for callers who only care about extras. Truncate- | {- | Raise a 'DataFrame.IO.CSV.Fast.CsvParseError' on any row whose- field count differs from the header. Matches polars' strict- schema-bound mode.- -}+ | -- | Raise a 'CsvParseError' on any row whose field count differs. RaiseOnRagged deriving (Eq, Show) @@ -55,9 +59,7 @@ 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.- -}+ | -- | Return the rows parsed before the stray quote; drop the rest. BestEffort deriving (Eq, Show) @@ -68,60 +70,100 @@ , typeSpec :: TypeSpec -- ^ Whether/how to infer types. (default: InferFromSample 100) , safeRead :: SafeReadMode- {- ^ Default 'SafeReadMode' for columns without an entry in- 'safeReadOverrides'. (default: 'NoSafeRead')- -}+ -- ^ Default 'SafeReadMode' for columns without an override. , 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- @- -}+ -- ^ Format of date fields, as recognized by Data.Time.Format. , columnSeparator :: Char -- ^ Character that separates column values.- , numColumns :: Maybe Int- {- ^ Maximum number of data rows to read ('Nothing' = all). Note: despite- the legacy name, this has always capped /rows/, not columns.- -}+ , numRowsToRead :: Maybe Int+ -- ^ Maximum number of data rows to read ('Nothing' = all).+ , readColumns :: Maybe [T.Text]+ -- ^ Columns to read ('Nothing' = all); unselected ones are never decoded. , 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')- -}+ -- ^ @dataframe-fastcsv@: how to treat rows with a non-header field count. , fastCsvOnUnclosedQuote :: UnclosedQuotePolicy- {- ^ @dataframe-fastcsv@: how to treat an unclosed quoted field at EOF.- (default: 'RaiseOnUnclosedQuote')- -}+ -- ^ @dataframe-fastcsv@: how to treat an unclosed quoted field at EOF. , 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').- -}+ -- ^ @dataframe-fastcsv@: strip whitespace from unquoted fields if True. } +-- | Whether a 'TypeSpec' still infers types from a sample. shouldInferFromSample :: TypeSpec -> Bool shouldInferFromSample (InferFromSample _) = True shouldInferFromSample (SpecifyTypes _ fallback) = shouldInferFromSample fallback shouldInferFromSample _ = False +-- | The explicit column/type pairs a 'TypeSpec' declares, if any. schemaTypeMap :: TypeSpec -> M.Map T.Text SchemaType schemaTypeMap (SpecifyTypes xs _) = M.fromList xs schemaTypeMap _ = M.empty +-- | A 'ReadOptions' value that contradicts itself, listing every problem.+newtype InvalidReadOptions = InvalidReadOptions [T.Text]++instance Show InvalidReadOptions where+ show (InvalidReadOptions problems) =+ T.unpack+ (T.intercalate "\n" ("Invalid ReadOptions:" : map (" - " <>) problems))++instance Exception InvalidReadOptions++-- | Everything self-contradictory about @opts@, judged without a file.+readOptionsErrors :: ReadOptions -> [T.Text]+readOptionsErrors opts =+ concat+ [ [ "readColumns is an empty selection; omit it to read every column"+ | Just [] <- [readColumns opts]+ ]+ , [ "readColumns names "+ <> T.intercalate ", " dups+ <> " more than once"+ | not (null dups)+ ]+ , [ "numRowsToRead is negative: " <> T.pack (show n)+ | Just n <- [numRowsToRead opts]+ , n < 0+ ]+ , [ "headerSpec is ProvideNames with no names"+ | ProvideNames [] <- [headerSpec opts]+ ]+ , [ "columnSeparator is " <> T.pack (show sep) <> ", which cannot delimit fields"+ | let sep = columnSeparator opts+ , sep `elem` ['"', '\n', '\r']+ ]+ ]+ where+ dups = maybe [] (\cs -> L.nub (cs L.\\ L.nub cs)) (readColumns opts)++-- | Throw 'InvalidReadOptions' if @opts@ contradicts itself.+validateReadOptions :: ReadOptions -> IO ()+validateReadOptions opts =+ unless (null problems) (throwIO (InvalidReadOptions problems))+ where+ problems = readOptionsErrors opts++-- | Resolve 'readColumns' against header names to (name, field index) pairs.+resolveSelection :: ReadOptions -> [T.Text] -> [(T.Text, Int)]+resolveSelection opts names = case readColumns opts of+ Nothing -> indexed+ Just requested -> case filter (`M.notMember` fieldIndex) requested of+ [] -> [(n, fieldIndex M.! n) | n <- requested]+ missing -> throw (ColumnsNotFoundException missing "readCsvWithOpts" names)+ where+ indexed = zip names [0 ..]+ fieldIndex = M.fromList indexed++-- | The sample size a 'TypeSpec' infers from; 0 if it never infers. typeInferenceSampleSize :: TypeSpec -> Int typeInferenceSampleSize (InferFromSample n) = n typeInferenceSampleSize (SpecifyTypes _ fallback) = typeInferenceSampleSize fallback typeInferenceSampleSize _ = 0 +-- | The default 'ReadOptions': infer from 100 rows, header row, all data. defaultReadOptions :: ReadOptions defaultReadOptions = ReadOptions@@ -131,7 +173,8 @@ , safeReadOverrides = [] , dateFormat = "%Y-%m-%d" , columnSeparator = ','- , numColumns = Nothing+ , numRowsToRead = Nothing+ , readColumns = Nothing , missingIndicators = ["Nothing", "NULL", "", " ", "nan", "null", "N/A", "NaN", "NAN", "NA"] , fastCsvOnRaggedRow = PadWithNull
src/DataFrame/IO/CSV/Internal/Read.hs view
@@ -4,11 +4,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | Driver for the default CSV reader (Round-2 WS-D): a single strict-scan feeding per-column sinks, with cassava-parity semantics pinned by-@IO.CsvGolden@. Ragged rows pad trailing columns with null and drop-extra fields (audit D6). The result is fully forced (audit D7).--}+-- | Driver for the default CSV reader: one strict scan feeding column sinks. module DataFrame.IO.CSV.Internal.Read (decodeCsvStrict) where import qualified Data.ByteString as BS@@ -18,6 +14,7 @@ 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)@@ -30,111 +27,119 @@ 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.Column.Builder import DataFrame.Internal.DataFrame (DataFrame (..), forceDataFrame)-import DataFrame.Internal.Schema (SchemaType (..), schemaType) 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 opts bs = BSU.unsafeUseAsCStringLen bs $ \(cstr, _) -> do- let !len = BS.length bs- !sep = fromIntegral (ord (columnSeparator opts)) :: Word8+decodeCsvStrict opts bs = do+ validateReadOptions opts+ BSU.unsafeUseAsCStringLen bs $ \(cstr, _) -> do+ let !len = BS.length bs+ !sep = fromIntegral (ord (columnSeparator opts)) :: Word8 - -- Position of the next non-blank record at/after @pos@ (cassava- -- drops single-empty-field records); @len@ when none remain.- findRecord !pos- | pos >= len = len- | otherwise = withField bs len sep pos $ \cs ce _ term next ->- if cs == ce && term /= termSep- then if term == termEof then len else findRecord next- else pos+ findRecord !pos+ | pos >= len = len+ | otherwise = withField bs len sep pos $ \cs ce _ term next ->+ if cs == ce && term /= termSep+ then if term == termEof then len else findRecord next+ else pos - -- Decode the record at @pos@ as header fields; returns the- -- position after the record.- headerFields !pos = go pos id- where- go !p acc = withField bs len sep p $ \cs ce unesc term next ->- let t =- TE.decodeUtf8Lenient- (if unesc then unescapeQuotes bs cs ce else sliceBS bs cs ce)- in if term == termSep- then go next (acc . (t :))- else (acc [t], next)+ headerFields !pos = go pos id+ where+ go !p acc = withField bs len sep p $ \cs ce unesc term next ->+ let t =+ TE.decodeUtf8Lenient+ (if unesc then unescapeQuotes bs cs ce else sliceBS bs cs ce)+ in if term == termSep+ then go next (acc . (t :))+ else (acc [t], next) - let hdrPos = findRecord 0- when (hdrPos >= len) (error "Empty CSV file")- let (hdrFields, afterHdr) = headerFields hdrPos- positional n = map (T.pack . show) [0 .. n - 1 :: Int]- (names, dataPos0) = case headerSpec opts of- UseFirstRow -> (map T.strip hdrFields, afterHdr)- NoHeader -> (positional (length hdrFields), hdrPos)- ProvideNames ns ->- (ns ++ drop (length ns) (positional (length hdrFields)), hdrPos)- !ncols = length names- dataPos = findRecord dataPos0- when (dataPos >= len) (error "Empty CSV file")+ let hdrPos = findRecord 0+ when (hdrPos >= len) (error "Empty CSV file")+ let (hdrFields, afterHdr) = headerFields hdrPos+ positional n = map (T.pack . show) [0 .. n - 1 :: Int]+ (names, dataPos0) = case headerSpec opts of+ UseFirstRow -> (map T.strip hdrFields, afterHdr)+ NoHeader -> (positional (length hdrFields), hdrPos)+ ProvideNames ns ->+ (ns ++ drop (length ns) (positional (length hdrFields)), hdrPos)+ !nfields = length names+ dataPos = findRecord dataPos0+ when (dataPos >= len) (error "Empty CSV file") - let 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) names- missing = if anyEither then [] else missingIndicators opts- missMode- | null missing = MissNone- | missing == missingIndicators defaultReadOptions = MissCanonical- | otherwise = MissCustom- env = Env bs (castPtr cstr) missMode missing- rowHint = len `div` max 1 (ncols * 8) + 16+ let selection = resolveSelection opts names+ !keptNames = map fst selection+ !ncols = length selection+ !sinkOf =+ VU.replicate nfields (-1)+ VU.// [(fieldIx, slot) | (slot, (_, fieldIx)) <- zip [0 ..] selection] - sinks <- mapM (newSink opts resolveMode rowHint) names- let !sinksV = V.fromList sinks+ let resolveMode = effectiveSafeRead (safeRead opts) (safeReadOverrides opts)+ anyEither = any (\n -> resolveMode n == EitherRead) keptNames+ missing = if anyEither then [] else missingIndicators opts+ missMode+ | null missing = MissNone+ | missing == missingIndicators defaultReadOptions = MissCanonical+ | otherwise = MissCustom+ env = Env bs (castPtr cstr) missMode missing+ rowHint = len `div` max 1 (nfields * 8) + 16 - let pad !row !col- | col >= ncols = pure ()- | otherwise = nullSink (V.unsafeIndex sinksV col) row >> pad row (col + 1)- dec b = if b > 0 then b - 1 else b- rowLoop !pos !row !budget- | budget == 0 || pos >= len = pure row- | otherwise = withField bs len sep pos $ \cs ce unesc term next ->- if cs == ce && term /= termSep- then if term == termEof then pure row else rowLoop next row budget- else do- feedSink env (V.unsafeIndex sinksV 0) row cs ce unesc- if term == termSep- then fieldLoop next row 1 budget- else pad row 1 >> rowLoop next (row + 1) (dec budget)- fieldLoop !pos !row !col !budget =- withField bs len sep pos $ \cs ce unesc term next -> do- when (col < ncols) $- feedSink env (V.unsafeIndex sinksV col) row cs ce unesc- if term == termSep- then fieldLoop next row (col + 1) budget- else pad row (col + 1) >> rowLoop next (row + 1) (dec budget)+ sinks <- mapM (newSink opts resolveMode rowHint) keptNames+ let !sinksV = V.fromList sinks+ sinkAt !col+ | col >= nfields = -1+ | otherwise = VU.unsafeIndex sinkOf col - nrows <- rowLoop dataPos 0 (fromMaybe (-1) (numColumns opts))+ let pad !row !col+ | col >= nfields = pure ()+ | otherwise = do+ let !slot = VU.unsafeIndex sinkOf col+ when (slot >= 0) $ nullSink (V.unsafeIndex sinksV slot) row+ pad row (col + 1)+ feedAt !slot !row !cs !ce !unesc =+ when (slot >= 0) $+ feedSink env (V.unsafeIndex sinksV slot) row cs ce unesc+ dec b = if b > 0 then b - 1 else b+ rowLoop !pos !row !budget+ | budget == 0 || pos >= len = pure row+ | otherwise = withField bs len sep pos $ \cs ce unesc term next ->+ if cs == ce && term /= termSep+ then if term == termEof then pure row else rowLoop next row budget+ else do+ feedAt (sinkAt 0) row cs ce unesc+ if term == termSep+ then fieldLoop next row 1 budget+ else pad row 1 >> rowLoop next (row + 1) (dec budget)+ fieldLoop !pos !row !col !budget =+ withField bs len sep pos $ \cs ce unesc term next -> do+ feedAt (sinkAt col) row cs ce unesc+ if term == termSep+ then fieldLoop next row (col + 1) budget+ else pad row (col + 1) >> rowLoop next (row + 1) (dec budget) - cols <- zipWithM (finalizeSink bs opts resolveMode nrows) names sinks- let df =- DataFrame- (V.fromList cols)- (M.fromList (zip names [0 ..]))- (nrows, ncols)- M.empty- pure $!- forceDataFrame- (parseWithTypes resolveMode (schemaTypeMap (typeSpec opts)) df)+ nrows <- rowLoop dataPos 0 (fromMaybe (-1) (numRowsToRead opts)) -{- | 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).--}+ cols <- zipWithM (finalizeSink bs opts resolveMode nrows) keptNames sinks+ let df =+ DataFrame+ (V.fromList cols)+ (M.fromList (zip keptNames [0 ..]))+ (nrows, ncols)+ M.empty+ pure $!+ forceDataFrame+ (parseWithTypes resolveMode (schemaTypeMap (typeSpec opts)) df)++-- | Resolve a column's sink from its schema type or inference. newSink :: ReadOptions -> (T.Text -> SafeReadMode) -> Int -> T.Text -> IO Sink newSink opts resolveMode rowHint name
src/DataFrame/IO/CSV/Internal/Scanner.hs view
@@ -1,12 +1,6 @@ {-# LANGUAGE BangPatterns #-} -{- | RFC 4180 field scanner for the default CSV reader (Round-2 WS-D).-One pass over a strict 'BS.ByteString'; cassava-parity semantics pinned-by the @IO.CsvGolden@ suite: quotes open fields only at the first byte,-a double quote inside an unquoted field is an error, garbage after a-closing quote is an error, an unclosed quote consumes to EOF, and rows-end at @\\n@, @\\r\\n@ or a lone @\\r@.--}+-- | RFC 4180 field scanner for the default CSV reader. module DataFrame.IO.CSV.Internal.Scanner ( Term, termSep,@@ -44,11 +38,7 @@ parseError off what = error ("CSV Parse Error: " ++ what ++ " at byte " ++ show off) -{- | Scan one field starting at @pos@ and continue with-@k contentStart contentEnd needsUnescape term nextPos@. Quoted fields-yield the bytes between the outer quotes; @needsUnescape@ is set when-the content contains doubled quotes ('unescapeQuotes' collapses them).--}+-- | Scan one field at @pos@ and pass its content bounds to a continuation. {-# INLINE withField #-} withField :: BS.ByteString ->@@ -88,25 +78,20 @@ | b == quote -> parseError i "stray double quote in unquoted field" | otherwise -> unquoted (i + 1) quoted !i !unesc- | i >= len = k (pos + 1) len unesc termEof len -- unclosed: rest of input+ | i >= len = k (pos + 1) len unesc termEof len -- unclosed quote | BSU.unsafeIndex bs i == quote = if i + 1 < len && BSU.unsafeIndex bs (i + 1) == quote then quoted (i + 2) True else finish (pos + 1) i unesc (i + 1) | otherwise = quoted (i + 1) unesc -{- | Whether the record starting at @pos@ is a blank record (a single-empty field, quoted or not, then end of row) — cassava drops these.-Assumes @pos < len@. Errors propagate from 'withField'.--}+-- | Whether the record at @pos@ is a blank record, which cassava drops. {-# INLINE isBlankRecordAt #-} isBlankRecordAt :: BS.ByteString -> Int -> Word8 -> Int -> Bool isBlankRecordAt bs !len !sep !pos = withField bs len sep pos (\cs ce _ term _ -> cs == ce && term /= termSep) -{- | Collapse doubled quotes (@\"\"@ → @\"@) in @[s, e)@ into a fresh strict-'BS.ByteString'. Every quote in the range is the first of a pair.--}+-- | Collapse doubled quotes in a range into a fresh strict ByteString. unescapeQuotes :: BS.ByteString -> Int -> Int -> BS.ByteString unescapeQuotes bs !s !e = BSI.unsafeCreateUptoN (e - s) (go s 0) where@@ -117,10 +102,7 @@ pokeByteOff (dst :: Ptr Word8) j w if w == quote then go (i + 2) (j + 1) dst else go (i + 1) (j + 1) dst -{- | Strip with @Data.ByteString.Char8.strip@ semantics (Latin-1-whitespace: HT..CR, space and NBSP 0xA0) and continue with the-stripped range.--}+-- | Strip Latin-1 whitespace like ByteString strip and continue with the range. {-# INLINE withStripC8 #-} withStripC8 :: BS.ByteString -> Int -> Int -> (Int -> Int -> r) -> r withStripC8 bs !s0 !e0 k = k s e@@ -135,10 +117,7 @@ | i > s && isC8Space (BSU.unsafeIndex bs (i - 1)) = skipB (i - 1) | otherwise = i -{- | Strip ASCII whitespace (the ASCII subset of @Data.Text.strip@) and-continue with the stripped range. Callers must fall back to a full-'Data.Text.strip' when the stripped range still has non-ASCII edges.--}+-- | Strip ASCII whitespace and continue with the stripped range. {-# INLINE withStripAscii #-} withStripAscii :: BS.ByteString -> Int -> Int -> (Int -> Int -> r) -> r withStripAscii bs !s0 !e0 k = k s e
src/DataFrame/IO/CSV/Internal/Sink.hs view
@@ -1,11 +1,6 @@ {-# LANGUAGE BangPatterns #-} -{- | Per-column sinks for the default CSV reader. Typed columns (explicit-schema) parse straight from the input slice into the WS-C builders;-inferred and EitherRead columns record unboxed @(start, end)@ offsets-into the input buffer (plus a cold side map for quote-unescaped cells)-for the inference pass ("DataFrame.IO.CSV.Internal.Infer").--}+-- | Per-column sinks for the default CSV reader. module DataFrame.IO.CSV.Internal.Sink ( Env (..), MissingMode (..),@@ -33,7 +28,7 @@ import Data.IORef import Data.Word (Word8) import DataFrame.IO.CSV.Internal.Scanner-import DataFrame.Internal.ColumnBuilder+import DataFrame.Internal.Column.Builder import DataFrame.Internal.Parsing.Fast ( isMissingFieldSlice, parseDoubleFieldSlice,@@ -41,9 +36,7 @@ ) import Foreign.Ptr (Ptr, castPtr, plusPtr) -{- | How missing tokens are detected (audit D2: canonical list goes-through the WS-B memcmp fast path; custom lists decode and compare).--}+-- | How missing tokens are detected. data MissingMode = MissNone | MissCanonical | MissCustom -- | Read-time environment shared by every sink.@@ -63,11 +56,7 @@ | SinkText !(TextBuilder RealWorld) | SinkBS !SliceCol -{- | Growable cell store for inferred columns: per cell a @(start, end)@-offset pair into the input (already Char8-stripped) and a 0\/1 validity-byte. Quote-unescaped cells (cold) live in an 'IM.IntMap' keyed by row;-their offset pair is @(-1, row)@.--}+-- | Growable cell store of offset pairs for inferred columns. data SliceCol = SliceCol { scOffs :: !(IORef (VUM.IOVector Int)) , scValid :: !(IORef (VUM.IOVector Word8))@@ -83,9 +72,7 @@ , cLen :: !Int } -{- | Resolve cell @i@ as a slice and continue (no 'BS.ByteString' header-is allocated for the common in-buffer case).--}+-- | Resolve cell @i@ as a slice and continue with it. withCell :: Cells -> Int -> (BS.ByteString -> Int -> Int -> r) -> r withCell c !i k = let s = VU.unsafeIndex (cOffs c) (2 * i)@@ -147,10 +134,7 @@ sliceBS bs s e = BSU.unsafeTake (e - s) (BSU.unsafeDrop s bs) {-# INLINE sliceBS #-} -{- | Feed one field slice @[s, e)@ of the input into a sink. @unesc@-fields (quoted with embedded @\"\"@) are unescaped into a fresh strict-'BS.ByteString' first (cold path).--}+-- | Feed one field slice of the input into a sink. feedSink :: Env -> Sink -> Int -> Int -> Int -> Bool -> IO () feedSink env sink !row !s !e !unesc = case sink of SinkBS sc@@ -194,11 +178,7 @@ SinkBS _ -> error "feedTyped: SinkBS handled by feedSink" {-# INLINE feedTyped #-} -{- | Text cells reproduce @T.strip . decodeUtf8Lenient@ plus the-missing-indicator test. Fast path: ASCII-edge slices append raw bytes-(the shared-buffer freeze decodes leniently per field); anything with-non-ASCII edges or a custom missing list goes through real 'T.strip'.--}+-- | Append a text cell, with a fast path for ASCII-edge slices. appendTextCell :: Env -> TextBuilder RealWorld ->
− src/DataFrame/IO/Internal/MutableColumn.hs
@@ -1,89 +0,0 @@-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--{- |-Mutable-column ingest helpers used by CSV readers. Kept out of-@DataFrame.Internal.Column@ so that the core column type does not need to-depend on @DataFrame.Internal.Parsing@.--}-module DataFrame.IO.Internal.MutableColumn (- writeColumn,- freezeColumn',-) where--import qualified Data.Text as T-import qualified Data.Vector as VB-import qualified Data.Vector.Mutable as VBM-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM--import Data.Maybe (fromMaybe)-import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))-import Type.Reflection (typeRep)--import DataFrame.Internal.Column (- Column (..),- MutableColumn (..),- buildBitmapFromNulls,- )-import DataFrame.Internal.Parsing (isNullish, readDouble, readInt)--writeColumn :: Int -> T.Text -> MutableColumn -> IO (Either T.Text Bool)-writeColumn i value (MBoxedColumn (col :: VBM.IOVector a)) =- case testEquality (typeRep @a) (typeRep @T.Text) of- Just Refl ->- if isNullish value- then VBM.unsafeWrite col i "" >> return (Left $! value)- else VBM.unsafeWrite col i value >> return (Right True)- Nothing -> return (Left value)-writeColumn i value (MUnboxedColumn (col :: VUM.IOVector a)) =- case testEquality (typeRep @a) (typeRep @Int) of- Just Refl -> case readInt value of- Just v -> VUM.unsafeWrite col i v >> return (Right True)- Nothing -> VUM.unsafeWrite col i 0 >> return (Left value)- Nothing -> case testEquality (typeRep @a) (typeRep @Double) of- Nothing -> return (Left $! value)- Just Refl -> case readDouble value of- Just v -> VUM.unsafeWrite col i v >> return (Right True)- Nothing -> VUM.unsafeWrite col i 0 >> return (Left $! value)-{-# INLINE writeColumn #-}--freezeColumn' :: [(Int, T.Text)] -> MutableColumn -> IO Column-freezeColumn' nulls (MBoxedColumn col)- | null nulls = BoxedColumn Nothing <$> VB.unsafeFreeze col- | all (isNullish . snd) nulls = do- frozen <- VB.unsafeFreeze col- let n = VB.length frozen- bm = buildBitmapFromNulls n (map fst nulls)- return $ BoxedColumn (Just bm) frozen- | otherwise =- BoxedColumn Nothing- . VB.imap- ( \i v ->- if i `elem` map fst nulls- then Left (fromMaybe (error "UNEXPECTED ERROR DURING FREEZE") (lookup i nulls))- else Right v- )- <$> VB.unsafeFreeze col-freezeColumn' nulls (MUnboxedColumn col)- | null nulls = UnboxedColumn Nothing <$> VU.unsafeFreeze col- | all (isNullish . snd) nulls = do- c <- VU.unsafeFreeze col- let n = VU.length c- bm = buildBitmapFromNulls n (map fst nulls)- return $ UnboxedColumn (Just bm) c- | otherwise = do- c <- VU.unsafeFreeze col- return $- BoxedColumn Nothing $- VB.generate- (VU.length c)- ( \i ->- if i `elem` map fst nulls- then Left (fromMaybe (error "UNEXPECTED ERROR DURING FREEZE") (lookup i nulls))- else Right (c VU.! i)- )-{-# INLINE freezeColumn' #-}
+ src/DataFrame/Typed/IO/CSV.hs view
@@ -0,0 +1,78 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | Typed CSV reading driven entirely by the type-level schema.+module DataFrame.Typed.IO.CSV (+ readCsv,+ readCsvWithError,+ readTsv,+ readCsvWithOpts,+ writeCsv,+ writeTsv,+) where++import Control.Applicative ((<|>))+import Control.Exception (SomeException, try)+import qualified Data.Text as T++import DataFrame.IO.CSV (ReadOptions (..), TypeSpec (..), defaultReadOptions)+import qualified DataFrame.IO.CSV as CSV+import DataFrame.Schema (RuntimeSchema (..), elements)+import DataFrame.Typed.Freeze (freezeOrThrow, freezeWithError, thaw)+import DataFrame.Typed.Schema (KnownSchema)+import DataFrame.Typed.Types (TypedDataFrame)++import qualified Data.Map.Strict as M++-- | Fold a type-level schema into read options: its columns and types.+schemaOptions :: forall cols. (RuntimeSchema cols) => ReadOptions -> ReadOptions+schemaOptions opts =+ opts+ { typeSpec = SpecifyTypes (M.toList declared) (typeSpec opts)+ , readColumns = readColumns opts <|> Just (M.keys declared)+ }+ where+ declared = elements (runtimeSchema @cols)++-- | Read a CSV into a typed DataFrame, throwing on schema mismatch.+readCsv ::+ forall cols.+ (KnownSchema cols, RuntimeSchema cols) =>+ FilePath -> IO (TypedDataFrame cols)+readCsv = readCsvWithOpts @cols defaultReadOptions++-- | Read a CSV, returning an error on schema mismatch instead of throwing.+readCsvWithError ::+ forall cols.+ (KnownSchema cols, RuntimeSchema cols) =>+ FilePath -> IO (Either T.Text (TypedDataFrame cols))+readCsvWithError path = do+ -- A missing column can fail the projection before the freeze does.+ r <- try (CSV.readCsvWithOpts (schemaOptions @cols defaultReadOptions) path)+ pure $ case r of+ Left (e :: SomeException) -> Left (T.pack (show e))+ Right df -> freezeWithError df++-- | Read a tab-separated file into a typed DataFrame, throwing on mismatch.+readTsv ::+ forall cols.+ (KnownSchema cols, RuntimeSchema cols) =>+ FilePath -> IO (TypedDataFrame cols)+readTsv = readCsvWithOpts @cols defaultReadOptions{columnSeparator = '\t'}++-- | Read a CSV with custom options, throwing on schema mismatch.+readCsvWithOpts ::+ forall cols.+ (KnownSchema cols, RuntimeSchema cols) =>+ ReadOptions -> FilePath -> IO (TypedDataFrame cols)+readCsvWithOpts opts path =+ CSV.readCsvWithOpts (schemaOptions @cols opts) path >>= freezeOrThrow @cols++-- | Write a typed DataFrame to a CSV file.+writeCsv :: FilePath -> TypedDataFrame cols -> IO ()+writeCsv path = CSV.writeCsv path . thaw++-- | Write a typed DataFrame to a tab-separated file.+writeTsv :: FilePath -> TypedDataFrame cols -> IO ()+writeTsv path = CSV.writeTsv path . thaw