diff --git a/dataframe-csv.cabal b/dataframe-csv.cabal --- a/dataframe-csv.cabal +++ b/dataframe-csv.cabal @@ -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 diff --git a/src-internal/DataFrame/IO/Internal/MutableColumn.hs b/src-internal/DataFrame/IO/Internal/MutableColumn.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/IO/Internal/MutableColumn.hs @@ -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' #-} diff --git a/src/DataFrame/IO/CSV.hs b/src/DataFrame/IO/CSV.hs --- a/src/DataFrame/IO/CSV.hs +++ b/src/DataFrame/IO/CSV.hs @@ -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 diff --git a/src/DataFrame/IO/CSV/Internal/Infer.hs b/src/DataFrame/IO/CSV/Internal/Infer.hs --- a/src/DataFrame/IO/CSV/Internal/Infer.hs +++ b/src/DataFrame/IO/CSV/Internal/Infer.hs @@ -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) => diff --git a/src/DataFrame/IO/CSV/Internal/Options.hs b/src/DataFrame/IO/CSV/Internal/Options.hs --- a/src/DataFrame/IO/CSV/Internal/Options.hs +++ b/src/DataFrame/IO/CSV/Internal/Options.hs @@ -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 diff --git a/src/DataFrame/IO/CSV/Internal/Read.hs b/src/DataFrame/IO/CSV/Internal/Read.hs --- a/src/DataFrame/IO/CSV/Internal/Read.hs +++ b/src/DataFrame/IO/CSV/Internal/Read.hs @@ -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 diff --git a/src/DataFrame/IO/CSV/Internal/Scanner.hs b/src/DataFrame/IO/CSV/Internal/Scanner.hs --- a/src/DataFrame/IO/CSV/Internal/Scanner.hs +++ b/src/DataFrame/IO/CSV/Internal/Scanner.hs @@ -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 diff --git a/src/DataFrame/IO/CSV/Internal/Sink.hs b/src/DataFrame/IO/CSV/Internal/Sink.hs --- a/src/DataFrame/IO/CSV/Internal/Sink.hs +++ b/src/DataFrame/IO/CSV/Internal/Sink.hs @@ -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 -> diff --git a/src/DataFrame/IO/Internal/MutableColumn.hs b/src/DataFrame/IO/Internal/MutableColumn.hs deleted file mode 100644 --- a/src/DataFrame/IO/Internal/MutableColumn.hs +++ /dev/null @@ -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' #-} diff --git a/src/DataFrame/Typed/IO/CSV.hs b/src/DataFrame/Typed/IO/CSV.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/IO/CSV.hs @@ -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