dataframe-csv 1.0.2.0 → 2.3.0.0
raw patch · 7 files changed
Files
- dataframe-csv.cabal +13/−11
- src-internal/DataFrame/IO/Internal/MutableColumn.hs +89/−0
- src/DataFrame/IO/CSV.hs +95/−14
- src/DataFrame/IO/CSV/Internal/Options.hs +184/−9
- src/DataFrame/IO/CSV/Internal/Read.hs +95/−84
- src/DataFrame/IO/Internal/MutableColumn.hs +0/−89
- src/DataFrame/Typed/IO/CSV.hs +136/−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.0.0 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.2 && < 2.3,+ dataframe-operations >= 2.2 && < 2.3,+ 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,89 @@+{-# 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/IO/CSV.hs view
@@ -1,11 +1,9 @@ {-# 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.+{- | CSV reading and writing for dataframes. A strict, single-pass+RFC 4180 scanner (cassava-compatible) parses fields into typed column+builders; ragged rows are padded with nulls and extra fields dropped. -} module DataFrame.IO.CSV ( -- * Reading@@ -15,6 +13,7 @@ readSeparated, readCsvWithSchema, CsvReader,+ schemaReadOptions, decodeSeparated, fromCsv, fromCsvBytes,@@ -43,7 +42,7 @@ 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)+import DataFrame.Schema (Schema, elements) {- | Read CSV file from path and load it into a dataframe. @@ -56,19 +55,49 @@ readCsv :: FilePath -> IO DataFrame readCsv = readSeparated defaultReadOptions -type CsvReader = Schema -> FilePath -> IO DataFrame+{- | A reader the lazy scan can drive: 'readSeparated' here, or+@fastReadCsvWithOpts@ from @dataframe-fastcsv@. The scan owns the+'ReadOptions', so every reader it accepts honours the projection it asks+for. -{- | 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.+==== __Example__+@+myReader :: CsvReader+myReader opts path = D.readCsvWithOpts opts path+@+-}+type CsvReader = ReadOptions -> FilePath -> IO DataFrame +{- | Read options a scan derives from its 'Schema': the schema assigns the+column types /and/ selects the columns, so a scan reads only what its+schema names — the same contract as @scanParquet@.++==== __Example__ @+ghci> schema = D.makeSchema [("id", D.schemaType \@Int), ("name", D.schemaType \@Text)]+ghci> D.readCsvWithOpts (schemaReadOptions schema) ".\/data\/customers.csv"++@+-}+schemaReadOptions :: Schema -> ReadOptions+schemaReadOptions schema =+ defaultReadOptions+ { typeSpec =+ SpecifyTypes (M.toList (elements schema)) (typeSpec defaultReadOptions)+ , readColumns = Just (M.keys (elements schema))+ }++{- | Schema-driven CSV reader. Coerces each column to the type declared+in 'Schema'; columns absent from the schema fall back to inference and are+still returned. To read /only/ the schema's columns, pass+'schemaReadOptions' to 'readCsvWithOpts'.++@ import qualified DataFrame as D df <- D.readCsvWithSchema schema "input.csv" @ -}-readCsvWithSchema :: CsvReader+readCsvWithSchema :: Schema -> FilePath -> IO DataFrame readCsvWithSchema schema = readSeparated defaultReadOptions@@ -110,22 +139,50 @@ -} 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 {- | Decode in-memory CSV bytes into a dataframe. The result is fully forced. (Note: unlike 'readSeparated', no UTF-8 BOM is stripped.)++==== __Example__+@+ghci> D.decodeSeparated D.defaultReadOptions "id,name\\n1,Ada\\n"++@ -} decodeSeparated :: ReadOptions -> BL.ByteString -> IO DataFrame decodeSeparated opts csvData = decodeCsvStrict opts (BL.toStrict csvData) +{- | Write a dataframe to a comma-separated file.++==== __Example__+@+ghci> D.writeCsv ".\/out.csv" df+@+-} writeCsv :: FilePath -> DataFrame -> IO () writeCsv = writeSeparated ',' +{- | Write a dataframe to a tab-separated file.++==== __Example__+@+ghci> D.writeTsv ".\/out.tsv" df+@+-} writeTsv :: FilePath -> DataFrame -> IO () writeTsv = writeSeparated '\t' +{- | Write a dataframe to a file using the given field separator.++==== __Example__+@+ghci> D.writeSeparated ';' ".\/out.txt" df+@+-} writeSeparated :: -- | Separator Char ->@@ -135,17 +192,41 @@ IO () writeSeparated c filepath df = TIO.writeFile filepath (toSeparated c df) --- | Parse a CSV string into a DataFrame using default options.+{- | Parse a CSV string into a DataFrame using default options.++==== __Example__+@+ghci> D.fromCsv "id,name\\n1,Ada\\n"+Right (DataFrame ...)++@+-} 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.+{- | Parse a lazy 'ByteString' containing CSV data into a DataFrame using+default options.++==== __Example__+@+ghci> D.fromCsvBytes "id,name\\n1,Ada\\n"++@+-} fromCsvBytes :: BL.ByteString -> IO DataFrame fromCsvBytes = decodeSeparated defaultReadOptions +{- | Strip one layer of double-quote delimiters from a field, if present.++==== __Example__+>>> stripQuotes "\"hello\""+"hello"+>>> stripQuotes "hello"+"hello"+-} stripQuotes :: T.Text -> T.Text stripQuotes txt = case T.uncons txt of
src/DataFrame/IO/CSV/Internal/Options.hs view
@@ -14,21 +14,55 @@ 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.++==== __Example__+@+ghci> D.readCsvWithOpts D.defaultReadOptions{D.headerSpec = D.ProvideNames ["a", "b"]} "no_header.csv"++@+-}+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.++==== __Example__+@+ghci> D.readCsvWithOpts D.defaultReadOptions{D.typeSpec = D.InferFromSample 500} "wide.csv"++@+-} 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; fall back to the+ nested 'TypeSpec' for the rest.+ -}+ 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@@ -87,10 +121,27 @@ -} , 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).++ __Example:__++ @+ ghci> D.readCsvWithOpts D.defaultReadOptions{D.numRowsToRead = Just 5} "big.csv"+ @ -}+ , readColumns :: Maybe [T.Text]+ {- ^ Columns to read ('Nothing' = all). Unselected columns are never+ decoded, parsed or allocated; the returned frame carries the selected+ columns in the order given here. Naming a column the file does not have+ raises a 'ColumnsNotFoundException'.++ __Example:__++ @+ ghci> D.readCsvWithOpts D.defaultReadOptions{D.readColumns = Just ["id", "name"]} "customers.csv"+ @+ -} , missingIndicators :: [T.Text] -- ^ Values that should be read as `Nothing`. , fastCsvOnRaggedRow :: RaggedRowPolicy@@ -108,20 +159,143 @@ -} } +{- | Whether a 'TypeSpec' still infers types from a sample, once any+'SpecifyTypes' fallback chain is followed to its end.++==== __Example__+>>> shouldInferFromSample (InferFromSample 100)+True+>>> shouldInferFromSample NoInference+False+-} shouldInferFromSample :: TypeSpec -> Bool shouldInferFromSample (InferFromSample _) = True shouldInferFromSample (SpecifyTypes _ fallback) = shouldInferFromSample fallback shouldInferFromSample _ = False +{- | The explicit column\/type pairs a 'TypeSpec' declares, if any.++==== __Example__+>>> :set -XTypeApplications+>>> M.toList (schemaTypeMap (SpecifyTypes [("id", schemaType @Int)] (InferFromSample 100)))+[("id",Int)]+-} schemaTypeMap :: TypeSpec -> M.Map T.Text SchemaType schemaTypeMap (SpecifyTypes xs _) = M.fromList xs schemaTypeMap _ = M.empty +{- | A 'ReadOptions' value that contradicts itself, with every problem listed.+Thrown by 'validateReadOptions'; see 'readOptionsErrors' for the check+without the exception.++==== __Example__+>>> show (InvalidReadOptions ["readColumns is an empty selection; omit it to read every column"])+"Invalid ReadOptions:\n - readColumns is an empty selection; omit it to read every column"+-}+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 looking at a+file. A type or safe-read override naming an unread column is redundant, not+contradictory, and is not reported.++==== __Example__+>>> readOptionsErrors defaultReadOptions{readColumns = Just []}+["readColumns is an empty selection; omit it to read every column"]+>>> readOptionsErrors defaultReadOptions+[]+-}+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. Both readers+call this before opening the file.++==== __Example__+@+ghci> D.validateReadOptions D.defaultReadOptions{D.readColumns = Just []}+*** Exception: Invalid ReadOptions:+ - readColumns is an empty selection; omit it to read every column+@+-}+validateReadOptions :: ReadOptions -> IO ()+validateReadOptions opts =+ unless (null problems) (throwIO (InvalidReadOptions problems))+ where+ problems = readOptionsErrors opts++{- | Resolve 'readColumns' against a file's header names: the columns to+build, each paired with its field index within a row. 'Nothing' keeps every+column in file order; a selection keeps the order it was written in.++==== __Example__+>>> resolveSelection defaultReadOptions{readColumns = Just ["b", "a"]} ["a", "b", "c"]+[("b",1),("a",0)]+>>> resolveSelection defaultReadOptions ["a", "b"]+[("a",0),("b",1)]+-}+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, once any 'SpecifyTypes'+fallback chain is followed to its end. @0@ when the chain ends in+'NoInference'.++==== __Example__+>>> typeInferenceSampleSize (InferFromSample 100)+100+>>> typeInferenceSampleSize NoInference+0+-} typeInferenceSampleSize :: TypeSpec -> Int typeInferenceSampleSize (InferFromSample n) = n typeInferenceSampleSize (SpecifyTypes _ fallback) = typeInferenceSampleSize fallback typeInferenceSampleSize _ = 0 +{- | The default 'ReadOptions': infer types from a 100-row sample, treat the+first row as a header, read every row and every column.++==== __Example__+@+ghci> D.readCsvWithOpts D.defaultReadOptions{D.columnSeparator = ';'} "data.csv"+@+-} defaultReadOptions :: ReadOptions defaultReadOptions = ReadOptions@@ -131,7 +305,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,10 +4,9 @@ {-# 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: a single strict scan feeding+per-column sinks, with cassava-parity semantics. Ragged rows pad+trailing columns with null, drop extra fields, and force the result. -} module DataFrame.IO.CSV.Internal.Read (decodeCsvStrict) where @@ -18,6 +17,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)@@ -32,104 +32,115 @@ import DataFrame.Internal.Column (Column, ensureOptional) import DataFrame.Internal.ColumnBuilder 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))++ 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: EitherRead and to-be-inferred columns collect raw slices; schema'd Int/Double parse straight into builders;
− 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,136 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | Typed CSV reading.++The schema is the whole specification of the read: it names the columns to+fetch and gives their types, so these readers touch only the columns @cols@+declares and never infer a type the schema already knows.++@+type Customer = '[ '(\"customer_id\", Int), '(\"customer_name\", Text)]++customers <- readCsv \@Customer \"customers.csv\" -- reads 2 columns, whatever+ -- else the file holds+@+-}+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: fetch exactly its columns,+and parse them at its types. Anything else the caller asked for is kept,+including a wider 'typeSpec' for columns the schema does not name.+-}+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 file into a typed DataFrame, throwing on schema mismatch.+Reads only the columns @cols@ names, typed as @cols@ says.++==== __Example__+@+ghci> type Customer = '[ '(\"id\", Int), '(\"name\", Text)]+ghci> customers <- readCsv \@Customer \"customers.csv\"+@+-}+readCsv ::+ forall cols.+ (KnownSchema cols, RuntimeSchema cols) =>+ FilePath -> IO (TypedDataFrame cols)+readCsv = readCsvWithOpts @cols defaultReadOptions++{- | Read a CSV file, returning a descriptive error on schema mismatch or a+missing column instead of throwing.++==== __Example__+@+ghci> readCsvWithError \@Customer \"customers.csv\"+Right (TDF ...)+@+-}+readCsvWithError ::+ forall cols.+ (KnownSchema cols, RuntimeSchema cols) =>+ FilePath -> IO (Either T.Text (TypedDataFrame cols))+readCsvWithError path = do+ -- The projection can fail before the freeze does (a column the file does+ -- not have); this reader reports rather than throws either way.+ 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.++==== __Example__+@+ghci> customers <- readTsv \@Customer \"customers.tsv\"+@+-}+readTsv ::+ forall cols.+ (KnownSchema cols, RuntimeSchema cols) =>+ FilePath -> IO (TypedDataFrame cols)+readTsv = readCsvWithOpts @cols defaultReadOptions{columnSeparator = '\t'}++{- | Read a CSV file with custom options, throwing on schema mismatch. The+schema still supplies the column selection and types; an explicit+'readColumns' takes precedence, and is then checked by the freeze.++==== __Example__+@+ghci> customers <- readCsvWithOpts \@Customer defaultReadOptions{safeRead = MaybeRead} \"customers.csv\"+@+-}+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.++==== __Example__+@+ghci> writeCsv \"customers.csv\" customers+@+-}+writeCsv :: FilePath -> TypedDataFrame cols -> IO ()+writeCsv path = CSV.writeCsv path . thaw++{- | Write a typed DataFrame to a tab-separated file.++==== __Example__+@+ghci> writeTsv \"customers.tsv\" customers+@+-}+writeTsv :: FilePath -> TypedDataFrame cols -> IO ()+writeTsv path = CSV.writeTsv path . thaw