diff --git a/dataframe-fastcsv.cabal b/dataframe-fastcsv.cabal --- a/dataframe-fastcsv.cabal +++ b/dataframe-fastcsv.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: dataframe-fastcsv -version: 1.3.0.0 +version: 1.4.0.0 synopsis: SIMD-accelerated CSV reader for the dataframe library. description: A fast, SIMD-accelerated CSV/TSV reader using memory-mapped I/O @@ -58,7 +58,7 @@ bytestring >= 0.11 && < 0.14, containers >= 0.6.7 && < 0.10, dataframe-core >= 2.1 && < 2.2, - dataframe-csv >= 2.1 && < 2.2, + dataframe-csv >= 2.2 && < 2.3, dataframe-operations >= 2.1 && < 2.2, dataframe-parsing >= 2.1 && < 2.2, dataframe-parsing >= 2.1 && < 2.2, @@ -85,13 +85,14 @@ -- capabilities (the library itself does not choose the RTS). ghc-options: -threaded -rtsopts "-with-rtsopts=-N" other-modules: Operations.ChunkParallel + Operations.Projection Operations.ReadCsv Operations.TypedExtraction Properties.Csv build-depends: base >= 4 && < 5, containers >= 0.6.7 && < 0.10, dataframe-core >= 2.1 && < 2.2, - dataframe-csv >= 2.1 && < 2.2, + dataframe-csv >= 2.2 && < 2.3, dataframe-fastcsv, dataframe-operations >= 2.1 && < 2.2, dataframe-parsing >= 2.1 && < 2.2, diff --git a/src/DataFrame/IO/CSV/Fast.hs b/src/DataFrame/IO/CSV/Fast.hs --- a/src/DataFrame/IO/CSV/Fast.hs +++ b/src/DataFrame/IO/CSV/Fast.hs @@ -23,58 +23,114 @@ import qualified Data.ByteString as BS import qualified Data.ByteString.Internal as BSI import qualified Data.Map as M -import qualified Data.Vector as Vector import qualified Data.Vector.Storable as VS +import Data.Char (ord) import Data.Text (Text) import Data.Word (Word8) import System.IO.MMap (Mode (WriteCopy), mmapFileForeignPtr) import DataFrame.IO.CSV ( + CsvReader, ReadOptions (..), TypeSpec (..), defaultReadOptions, + validateReadOptions, ) -import DataFrame.IO.CSV.Fast.Core (readSeparatedCore) +import DataFrame.IO.CSV.Fast.Core (checkSeparatorAgrees, readSeparatedCore) import DataFrame.IO.CSV.Fast.Index ( CsvParseError (..), comma, getDelimiterIndices, tab, ) -import DataFrame.Internal.DataFrame (DataFrame (..)) +import DataFrame.Internal.DataFrame (DataFrame) import DataFrame.Schema (Schema (..)) -readSeparatedDefault :: Word8 -> FilePath -> IO DataFrame -readSeparatedDefault separator = - readSeparated separator defaultReadOptions +-- | The separator a 'ReadOptions' asks for, as the scanner's byte. +separatorByte :: ReadOptions -> Word8 +separatorByte = fromIntegral . ord . columnSeparator +{- | Read a CSV file with the default options. + +==== __Example__ +@ +ghci> D.fastReadCsv ".\/data\/taxi.csv" + +@ +-} fastReadCsv :: FilePath -> IO DataFrame -fastReadCsv = readSeparatedDefault comma +fastReadCsv = fastReadCsvWithOpts defaultReadOptions +{- | Alias for 'fastReadCsv'. + +==== __Example__ +@ +ghci> D.readCsvFast ".\/data\/taxi.csv" + +@ +-} readCsvFast :: FilePath -> IO DataFrame readCsvFast = fastReadCsv +{- | Read a TSV file with the default options. + +==== __Example__ +@ +ghci> D.fastReadTsv ".\/data\/taxi.tsv" + +@ +-} fastReadTsv :: FilePath -> IO DataFrame -fastReadTsv = readSeparatedDefault tab +fastReadTsv = fastReadTsvWithOpts defaultReadOptions +{- | Alias for 'fastReadTsv'. + +==== __Example__ +@ +ghci> D.readTsvFast ".\/data\/taxi.tsv" + +@ +-} readTsvFast :: FilePath -> IO DataFrame readTsvFast = fastReadTsv {- | Like 'fastReadCsv' but takes a 'ReadOptions' record. Use this when -you need to tune ragged-row handling, unclosed-quote handling, or the -whitespace-trimming knob; otherwise stick with 'fastReadCsv'. +you need to tune ragged-row handling, unclosed-quote handling, the +whitespace-trimming knob, or the column selection; otherwise stick with +'fastReadCsv'. The separator comes from 'columnSeparator', so this is a +'CsvReader' the lazy scan can drive. + +==== __Example__ +@ +ghci> D.fastReadCsvWithOpts D.defaultReadOptions{D.readColumns = Just ["id", "name"]} ".\/data\/customers.csv" + +@ -} -fastReadCsvWithOpts :: ReadOptions -> FilePath -> IO DataFrame -fastReadCsvWithOpts = readSeparated comma +fastReadCsvWithOpts :: CsvReader +fastReadCsvWithOpts opts = readSeparated (separatorByte opts) opts --- | TSV counterpart to 'fastReadCsvWithOpts'. +{- | TSV counterpart to 'fastReadCsvWithOpts'; pins 'columnSeparator' to tab. + +==== __Example__ +@ +ghci> D.fastReadTsvWithOpts D.defaultReadOptions{D.safeRead = D.MaybeRead} ".\/data\/taxi.tsv" + +@ +-} fastReadTsvWithOpts :: ReadOptions -> FilePath -> IO DataFrame -fastReadTsvWithOpts = readSeparated tab +fastReadTsvWithOpts opts = fastReadCsvWithOpts opts{columnSeparator = '\t'} {- | Read a CSV and coerce each column to the type declared in the supplied 'Schema'. Schema columns bypass inference: parsed straight from file bytes as the declared type, avoiding row-1 misclassification. + +==== __Example__ +@ +ghci> schema = D.makeSchema [("id", D.schemaType \@Int)] +ghci> D.fastReadCsvWithSchema schema ".\/data\/customers.csv" + +@ -} fastReadCsvWithSchema :: Schema -> FilePath -> IO DataFrame fastReadCsvWithSchema schema = @@ -85,10 +141,26 @@ SpecifyTypes (M.toList (elements schema)) (typeSpec defaultReadOptions) } --- | In-memory version of 'fastReadCsv'. +{- | In-memory version of 'fastReadCsv'. + +==== __Example__ +@ +ghci> D.fastReadCsvFromBytes "id,name\\n1,Ada\\n" + +@ +-} fastReadCsvFromBytes :: BS.ByteString -> IO DataFrame fastReadCsvFromBytes = readSeparatedFromBytes comma defaultReadOptions +{- | In-memory version of 'fastReadCsvWithSchema'. + +==== __Example__ +@ +ghci> schema = D.makeSchema [("id", D.schemaType \@Int)] +ghci> D.fastReadCsvFromBytesWithSchema schema "id,name\\n1,Ada\\n" + +@ +-} fastReadCsvFromBytesWithSchema :: Schema -> BS.ByteString -> IO DataFrame fastReadCsvFromBytesWithSchema schema = readSeparatedFromBytes @@ -99,33 +171,39 @@ } {- | Read a CSV and return only the named columns, in the order given. -The SIMD scan and delimiter classification still run over the whole -file, but unreferenced columns are dropped from the result. +Shorthand for 'readColumns'. The SIMD scan and delimiter +classification still run over the whole file, but unnamed columns are +never extracted, parsed or allocated. Naming a column the file does not +have raises a 'ColumnsNotFoundException'. + +==== __Example__ +@ +ghci> D.fastReadCsvProj ["id", "name"] ".\/data\/customers.csv" + +@ -} fastReadCsvProj :: [Text] -> FilePath -> IO DataFrame -fastReadCsvProj projection path = do - df <- fastReadCsv path - pure (projectColumns projection df) +fastReadCsvProj projection = + readSeparated comma defaultReadOptions{readColumns = Just projection} -{- | Filter a DataFrame's columns down to (and in the order of) the -supplied names. Missing names are silently dropped; check the result's -'columnIndices' map for strict semantics. --} -projectColumns :: [Text] -> DataFrame -> DataFrame -projectColumns names df = - let idxs = [(n, i) | n <- names, Just i <- [M.lookup n (columnIndices df)]] - newCols = Vector.fromListN (length idxs) [columns df Vector.! i | (_, i) <- idxs] - newIndices = M.fromList (zip (map fst idxs) [0 ..]) - (rows, _) = dataframeDimensions df - in DataFrame newCols newIndices (rows, length idxs) M.empty +{- | Reads via a private ('WriteCopy') mmap so the buffer is never copied. +The @separator@ byte and 'columnSeparator' in @opts@ must agree; use +'fastReadCsvWithOpts' unless you need to pass a delimiter byte directly. --- | Reads via a private ('WriteCopy') mmap so the buffer is never copied. +==== __Example__ +@ +ghci> D.readSeparated 44 D.defaultReadOptions ".\/data\/taxi.csv" -- 44 = ',' + +@ +-} readSeparated :: Word8 -> ReadOptions -> FilePath -> IO DataFrame readSeparated separator opts filePath = do + validateReadOptions opts + checkSeparatorAgrees separator opts (bufferPtr, offset, len) <- mmapFileForeignPtr filePath @@ -140,6 +218,12 @@ {- | Identical to 'readSeparated' but reads from an already-loaded byte buffer (e.g. a slice of a memory-mapped file). Zero-copy: the bytes are only read, never modified or grown. + +==== __Example__ +@ +ghci> D.readSeparatedFromBytes 44 D.defaultReadOptions "id,name\\n1,Ada\\n" -- 44 = ',' + +@ -} readSeparatedFromBytes :: Word8 -> @@ -153,6 +237,12 @@ {- | 'readSeparatedFromBytes' with a forced chunk count, bypassing the capability/size gate. Intended for tests and benchmarks that need to pin the parallel split; @1@ forces the sequential path. + +==== __Example__ +@ +ghci> D.readSeparatedFromBytesChunks 1 44 D.defaultReadOptions "id,name\\n1,Ada\\n" -- 44 = ',' + +@ -} readSeparatedFromBytesChunks :: Int -> diff --git a/src/DataFrame/IO/CSV/Fast/Core.hs b/src/DataFrame/IO/CSV/Fast/Core.hs --- a/src/DataFrame/IO/CSV/Fast/Core.hs +++ b/src/DataFrame/IO/CSV/Fast/Core.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE OverloadedStrings #-} + {- | The shared read pipeline behind every "DataFrame.IO.CSV.Fast" entry point: scan for delimiters, classify rows, then build every column with the typed extraction passes — chunk-parallel when the input is large and @@ -6,6 +8,7 @@ -} module DataFrame.IO.CSV.Fast.Core ( readSeparatedCore, + checkSeparatorAgrees, ) where import qualified Data.Map as M @@ -15,16 +18,21 @@ import qualified Data.Vector.Storable as VS import Control.Exception (throwIO) -import Control.Monad (when) +import Control.Monad (unless, when) +import Data.Char (ord) +import qualified Data.Text as T import Data.Word (Word8) import Debug.Trace (traceMarkerIO) import DataFrame.IO.CSV ( HeaderSpec (..), + InvalidReadOptions (..), RaggedRowPolicy (..), ReadOptions (..), UnclosedQuotePolicy (..), + resolveSelection, schemaTypeMap, + validateReadOptions, ) import DataFrame.IO.CSV.Fast.Index ( CsvParseError (..), @@ -55,6 +63,8 @@ VS.Vector Word8 -> IO DataFrame readSeparatedCore chunkSpec separator opts fileRaw = do + validateReadOptions opts + checkSeparatorAgrees separator opts let (file, _bomLen) = stripBom fileRaw contentLen = VS.length file nChunks <- maybe (autoChunkCount contentLen) pure chunkSpec @@ -96,21 +106,22 @@ , fcTrim = fastCsvTrimUnquoted opts } headerNumCol = fieldsInRow rowEnds 0 - (columnNames, dataStartRow, numCol) = case headerSpec opts of + (headerNames, dataStartRow, numCol) = case headerSpec opts of NoHeader -> - ( Vector.fromList $ - map (Text.pack . show) [0 .. headerNumCol - 1] + ( map (Text.pack . show) [0 .. headerNumCol - 1] , 0 , headerNumCol ) UseFirstRow -> - ( Vector.fromList $ - map (extractField ctx 0) [0 .. headerNumCol - 1] + ( map (extractField ctx 0) [0 .. headerNumCol - 1] , 1 , headerNumCol ) ProvideNames ns -> - (Vector.fromList ns, 0, length ns) + (ns, 0, length ns) + wanted = resolveSelection opts headerNames + columnNames = Vector.fromList (map fst wanted) + numKept = length wanted if numCol == 0 then return emptyDataFrame else do @@ -137,9 +148,9 @@ max 64 (contentLen `div` numCol) , ceOpts = opts } - results <- buildAllColumns nChunks env columnNames numCol + results <- buildAllColumns nChunks env wanted traceMarkerIO "fastcsv:build-done" - let cols = Vector.fromListN numCol (map fst results) + let cols = Vector.fromListN numKept (map fst results) legacySchema = S.fromList [ columnNames Vector.! c @@ -152,7 +163,7 @@ DataFrame cols colIndices - (numRow, numCol) + (numRow, numKept) M.empty resolveMode = effectiveSafeRead @@ -178,6 +189,21 @@ Vector.foldl' (\() c -> c `seq` ()) () cols `seq` return df else return $! forceDataFrame df' + +{- | The scanner is handed its separator as a byte, so a 'columnSeparator' +saying anything else would be silently ignored. Reject that up front. +-} +checkSeparatorAgrees :: Word8 -> ReadOptions -> IO () +checkSeparatorAgrees separator opts = + unless (separator == fromIntegral (ord asked)) $ + throwIO . InvalidReadOptions $ + [ "columnSeparator is " + <> T.pack (show asked) + <> " but this reader was given " + <> T.pack (show (toEnum (fromIntegral separator) :: Char)) + ] + where + asked = columnSeparator opts {- | An empty 'DataFrame' — returned when the input has no delimiters (empty file or single line with no separator and no newline). Guards diff --git a/src/DataFrame/IO/CSV/Fast/Parallel.hs b/src/DataFrame/IO/CSV/Fast/Parallel.hs --- a/src/DataFrame/IO/CSV/Fast/Parallel.hs +++ b/src/DataFrame/IO/CSV/Fast/Parallel.hs @@ -19,12 +19,10 @@ ) where import qualified Data.Text as T -import qualified Data.Vector as V import qualified Data.Vector.Storable as VS import Control.Concurrent (getNumCapabilities) import Control.Exception (throwIO) -import Control.Monad (zipWithM) import Data.Either (partitionEithers) import Debug.Trace (traceMarkerIO) import System.Mem (performMajorGC) @@ -73,16 +71,17 @@ tests and property. -} buildAllColumns :: - Int -> ColumnEnv -> V.Vector T.Text -> Int -> IO [(Column, Bool)] -buildAllColumns nChunks env names numCol + Int -> ColumnEnv -> [(T.Text, Int)] -> IO [(Column, Bool)] +buildAllColumns nChunks env wanted | chunks <= 1 = - mapM (\c -> buildColumn env (names V.! c) c) [0 .. numCol - 1] + mapM (uncurry (buildColumn env)) wanted | otherwise = do width <- getNumCapabilities - let plans = [planColumn env (names V.! c) c | c <- [0 .. numCol - 1]] + let plans = [planColumn env name fieldIx | (name, fieldIx) <- wanted] + fields = map snd wanted ranges = chunkRanges chunks (ceNumRow env) traceMarkerIO "fastcsv:fanout" - perChunk <- pooledRun width [runChunk env plans r | r <- ranges] + perChunk <- pooledRun width [runChunk env (zip fields plans) r | r <- ranges] traceMarkerIO "fastcsv:join-done" -- Reset the major-GC trigger here, where the copyable live set is -- small (chunk payloads are large objects). Otherwise the heap @@ -100,12 +99,12 @@ mergeColumn width env - (names V.! c) - (plans !! c) - c - [(rng, cols !! c) | (rng, cols) <- zip ranges perChunk] + name + plan + fieldIx + [(rng, cols !! slot) | (rng, cols) <- zip ranges perChunk] forceColumn col `seq` pure r - | c <- [0 .. numCol - 1] + | (slot, ((name, fieldIx), plan)) <- zip [0 ..] (zip wanted plans) ] traceMarkerIO "fastcsv:merge-done" pure merged @@ -132,8 +131,8 @@ } -- | One worker: run every column's plan over one chunk of rows. -runChunk :: ColumnEnv -> [ColumnPlan] -> (Int, Int) -> IO [ChunkCol] -runChunk env plans range@(off, _) = zipWithM build [0 ..] plans +runChunk :: ColumnEnv -> [(Int, ColumnPlan)] -> (Int, Int) -> IO [ChunkCol] +runChunk env plans range@(off, _) = mapM (uncurry build) plans where cenv = chunkEnv env range build col plan = case plan of diff --git a/tests/Main.hs b/tests/Main.hs --- a/tests/Main.hs +++ b/tests/Main.hs @@ -8,6 +8,7 @@ import Test.QuickCheck import qualified Operations.ChunkParallel +import qualified Operations.Projection import qualified Operations.ReadCsv import qualified Operations.TypedExtraction import qualified Properties.Csv @@ -18,6 +19,7 @@ ( Operations.ReadCsv.tests <> Operations.TypedExtraction.tests <> Operations.ChunkParallel.tests + <> Operations.Projection.tests ) isSuccessful :: Result -> Bool diff --git a/tests/Operations/Projection.hs b/tests/Operations/Projection.hs new file mode 100644 --- /dev/null +++ b/tests/Operations/Projection.hs @@ -0,0 +1,90 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +{- | 'readColumns' semantics for the SIMD reader, pinned against the +pure reader. Every case asserts both readers agree: the two disagreed +over 'ProvideNames' for a long time, and a projection only one of them +honours is worse than no projection at all. +-} +module Operations.Projection (tests) where + +import qualified Data.Map as M +import qualified Data.Text as T +import qualified Data.Text.IO as TIO + +import Control.Exception (SomeException, evaluate, try) +import Data.List (isInfixOf) +import qualified DataFrame.IO.CSV as CSV +import qualified DataFrame.IO.CSV.Fast as Fast +import DataFrame.Internal.DataFrame ( + DataFrame, + columnIndices, + dataframeDimensions, + forceDataFrame, + ) +import System.Directory (removeFile) +import Test.HUnit + +withCsv :: String -> T.Text -> (FilePath -> IO a) -> IO a +withCsv label contents k = do + let path = "./tests/data/unstable_csv/projection_" <> label <> ".csv" + TIO.writeFile path contents + r <- k path + removeFile path + pure r + +{- | Read @contents@ through both readers with @opts@, check they agree, +then hand the frame to the assertion. +-} +bothReaders :: + String -> CSV.ReadOptions -> T.Text -> (DataFrame -> Assertion) -> Test +bothReaders label opts contents check = TestLabel label $ + TestCase $ + withCsv label contents $ \path -> do + pureDf <- forceDataFrame <$> CSV.readSeparated opts path + fastDf <- forceDataFrame <$> Fast.fastReadCsvWithOpts opts path + assertEqual (label <> ": readers agree") pureDf fastDf + check pureDf + +-- | Both readers must reject a name the file does not have. +bothReadersFail :: String -> CSV.ReadOptions -> T.Text -> String -> Test +bothReadersFail label opts contents needle = TestLabel label $ + TestCase $ + withCsv label contents $ \path -> do + let expectFail which act = do + r <- try (act >>= evaluate . forceDataFrame) + case r of + Left (e :: SomeException) -> + assertBool + (label <> ": " <> which <> " should mention " <> show needle) + (needle `isInfixOf` show e) + Right _ -> + assertFailure (label <> ": " <> which <> " should have failed") + expectFail "pure reader" (CSV.readSeparated opts path) + expectFail "fast reader" (Fast.fastReadCsvWithOpts opts path) + +sample :: T.Text +sample = "id,name,surname,dob\n1,Ada,Lovelace,1815\n2,Alan,Turing,1912\n" + +select :: [T.Text] -> CSV.ReadOptions +select cs = CSV.defaultReadOptions{CSV.readColumns = Just cs} + +tests :: [Test] +tests = + [ bothReaders "prefix_subset" (select ["id", "name"]) sample $ + assertEqual "kept" (2, 2) . dataframeDimensions + , bothReaders "non_prefix_subset" (select ["id", "dob"]) sample $ \df -> do + assertEqual "kept" (2, 2) (dataframeDimensions df) + assertEqual "names" [("dob", 1), ("id", 0)] (M.toList (columnIndices df)) + , bothReaders "requested_order_wins" (select ["dob", "id"]) sample $ + assertEqual "order" [("dob", 0), ("id", 1)] . M.toList . columnIndices + , bothReaders "single_column" (select ["surname"]) sample $ + assertEqual "kept" (2, 1) . dataframeDimensions + , bothReaders "no_selection_keeps_all" CSV.defaultReadOptions sample $ + assertEqual "kept" (2, 4) . dataframeDimensions + , bothReadersFail + "missing_column" + (select ["id", "nope"]) + sample + "Column not found" + ]