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.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 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,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' #-} 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,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 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 @@ -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 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,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; 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,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