diff --git a/cbits/process_csv.c b/cbits/process_csv.c --- a/cbits/process_csv.c +++ b/cbits/process_csv.c @@ -119,8 +119,12 @@ // I owe a debt to https://github.com/geofflangdale/simdcsv // Let's go ahead and assume `in` will only ever get 64 bytes // initial_quoted will be either all_ones ~0ULL or all_zeros 0ULL +// Also reports the unquoted-newline bits through `newline_bits` so the +// chunked scan can classify row ends without re-touching file bytes. #ifdef HAS_SIMD_CSV -static uint64_t parse_chunk(uint8_t *in, uint8_t separator, uint64_t *initial_quoted) { +static uint64_t parse_chunk_nl(uint8_t *in, uint8_t separator, + uint64_t *initial_quoted, + uint64_t *newline_bits) { uint64_t quotebits = find_character_in_chunk(in, QUOTE_CHAR); // See https://wunkolo.github.io/post/2020/05/pclmulqdq-tricks/ // Also, section 3.1.1 of Parsing Gigabytes of JSON per Second, @@ -144,9 +148,16 @@ uint64_t commabits = find_character_in_chunk(in, separator); uint64_t newlinebits = find_character_in_chunk(in, NEWLINE_CHAR); - uint64_t delimiter_bits = (commabits | newlinebits) & ~quotemask; + uint64_t not_quoted = ~quotemask; + *newline_bits = newlinebits & not_quoted; + uint64_t delimiter_bits = (commabits | newlinebits) & not_quoted; return delimiter_bits; } + +static uint64_t parse_chunk(uint8_t *in, uint8_t separator, uint64_t *initial_quoted) { + uint64_t newline_bits; + return parse_chunk_nl(in, separator, initial_quoted, &newline_bits); +} #endif #ifdef HAS_SIMD_CSV @@ -201,6 +212,74 @@ // SIMD not available or carryless multiplication not supported. // Signal fallback to Haskell implementation. (void)buf; (void)len; (void)separator; (void)indices; + return GDI_SIMD_UNAVAILABLE; +#endif +} + +// Chunked entry points for the parallel scan (Round-2 WS-E2). A worker +// owns a 64-byte-aligned byte range [base, base+len) of the file; the +// caller supplies the quote parity at `base` (from a prefix-xor of +// per-range quote counts, which is exact because every quote byte +// toggles the PCLMUL parity chain). + +// Pass 1: exact delimiter / row-end counts for one range, so the flat +// output arrays can be right-sized and each worker gets an exact write +// offset. `len` must be a multiple of 64. +size_t count_delimiters_chunk(const uint8_t *buf, size_t base, size_t len, + uint8_t separator, int start_in_quote, + size_t *newline_count) { +#ifdef HAS_SIMD_CSV + uint64_t initial_quoted = start_in_quote ? ALL_ONES_MASK : ALL_ZEROS_MASK; + size_t count = 0; + size_t newlines = 0; + for (size_t i = 0; i + 64 <= len; i += 64) { + uint64_t newline_bits; + uint64_t bits = parse_chunk_nl((uint8_t *)buf + base + i, separator, + &initial_quoted, &newline_bits); + count += (size_t)__builtin_popcountll(bits); + newlines += (size_t)__builtin_popcountll(newline_bits); + } + *newline_count = newlines; + return count; +#else + (void)buf; (void)base; (void)len; (void)separator; (void)start_in_quote; + (void)newline_count; + return GDI_SIMD_UNAVAILABLE; +#endif +} + +// Pass 2: emit absolute delimiter byte positions into `indices` and, for +// every unquoted newline, its delimiter ordinal (`ordinal_base` + local +// position) into `rowends` -- the exact shape Haskell's classifyRowEnds +// used to recompute by re-reading one file byte per delimiter. +size_t emit_delimiter_indices_chunk(const uint8_t *buf, size_t base, + size_t len, uint8_t separator, + int start_in_quote, size_t ordinal_base, + size_t *indices, size_t *rowends, + size_t *rowend_count) { +#ifdef HAS_SIMD_CSV + uint64_t initial_quoted = start_in_quote ? ALL_ONES_MASK : ALL_ZEROS_MASK; + size_t pos = 0; + size_t rc = 0; + for (size_t i = 0; i + 64 <= len; i += 64) { + uint64_t newline_bits; + uint64_t bits = parse_chunk_nl((uint8_t *)buf + base + i, separator, + &initial_quoted, &newline_bits); + while (bits != 0) { + size_t r = (size_t)__builtin_ctzll(bits); + indices[pos] = base + i + r; + if ((newline_bits >> r) & 1ULL) { + rowends[rc++] = ordinal_base + pos; + } + pos++; + bits &= bits - 1; + } + } + *rowend_count = rc; + return pos; +#else + (void)buf; (void)base; (void)len; (void)separator; (void)start_in_quote; + (void)ordinal_base; (void)indices; (void)rowends; (void)rowend_count; return GDI_SIMD_UNAVAILABLE; #endif } diff --git a/cbits/process_csv.h b/cbits/process_csv.h --- a/cbits/process_csv.h +++ b/cbits/process_csv.h @@ -42,4 +42,16 @@ size_t get_delimiter_indices(uint8_t *buf, size_t len, uint8_t separator, size_t *indices, int *status); +// Chunked parallel-scan entry points (see process_csv.c). Both return +// GDI_SIMD_UNAVAILABLE when the build/CPU lacks SIMD support. +size_t count_delimiters_chunk(const uint8_t *buf, size_t base, size_t len, + uint8_t separator, int start_in_quote, + size_t *newline_count); + +size_t emit_delimiter_indices_chunk(const uint8_t *buf, size_t base, + size_t len, uint8_t separator, + int start_in_quote, size_t ordinal_base, + size_t *indices, size_t *rowends, + size_t *rowend_count); + #endif diff --git a/dataframe-fastcsv.cabal b/dataframe-fastcsv.cabal --- a/dataframe-fastcsv.cabal +++ b/dataframe-fastcsv.cabal @@ -1,7 +1,6 @@ cabal-version: 2.4 name: dataframe-fastcsv -version: 1.1.0.1 - +version: 1.1.1.0 synopsis: SIMD-accelerated CSV reader for the dataframe library. description: A fast, SIMD-accelerated CSV/TSV reader using memory-mapped I/O @@ -46,17 +45,25 @@ import: warnings default-extensions: Strict exposed-modules: DataFrame.IO.CSV.Fast + other-modules: DataFrame.IO.CSV.Fast.Columns + DataFrame.IO.CSV.Fast.Core + DataFrame.IO.CSV.Fast.Index + DataFrame.IO.CSV.Fast.IndexPar + DataFrame.IO.CSV.Fast.Parallel + DataFrame.IO.CSV.Fast.Passes + DataFrame.IO.CSV.Fast.Slice + DataFrame.IO.CSV.Fast.TextMerge + DataFrame.IO.CSV.Fast.Workers build-depends: base >= 4 && < 5, - array >= 0.5.4.0 && < 0.6, bytestring >= 0.11 && < 0.13, containers >= 0.6.7 && < 0.9, - dataframe-core ^>= 1.0, - dataframe-csv ^>= 1.0, - dataframe-operations ^>= 1.1, - dataframe-parsing ^>= 1.0, + dataframe-core ^>= 1.1, + dataframe-csv ^>= 1.0.2, + dataframe-operations ^>= 1.1.1, + dataframe-parsing ^>= 1.0.2, mmap >= 0.5.8 && < 0.6, - parallel >= 3.2.2.0 && < 5, text >= 2.0 && < 3, + time >= 1.12 && < 2, vector ^>= 0.13 hs-source-dirs: src c-sources: cbits/process_csv.c @@ -73,19 +80,26 @@ import: warnings type: exitcode-stdio-1.0 main-is: Main.hs - other-modules: Operations.ReadCsv + -- -threaded so the chunk-parallel path actually runs on multiple + -- capabilities (the library itself does not choose the RTS). + ghc-options: -threaded -rtsopts "-with-rtsopts=-N" + other-modules: Operations.ChunkParallel + Operations.ReadCsv + Operations.TypedExtraction Properties.Csv build-depends: base >= 4 && < 5, containers >= 0.6.7 && < 0.9, dataframe >= 1 && < 3, - dataframe-core ^>= 1.0, - dataframe-csv ^>= 1.0, + dataframe-core ^>= 1.1, + dataframe-csv ^>= 1.0.2, dataframe-fastcsv, - dataframe-parsing ^>= 1.0, + dataframe-operations ^>= 1.1.1, + dataframe-parsing ^>= 1.0.2, directory >= 1.3.0.0 && < 2, HUnit ^>= 1.6, QuickCheck >= 2 && < 3, text >= 2.0 && < 3, + time >= 1.12 && < 2, vector ^>= 0.13 hs-source-dirs: tests default-language: Haskell2010 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 @@ -1,7 +1,14 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CApiFFI #-} -{-# LANGUAGE ForeignFunctionInterface #-} +{- | SIMD-accelerated CSV reader: a C scanner finds delimiter positions, +then each column is parsed directly from the mmap'd byte slices into +unboxed column builders (Round-2 typed extraction). Reads are strict: the +returned 'DataFrame' is fully built, with no deferred parse work. +Extraction is chunk-parallel: inputs of 4 MB and up are split into one +row chunk per capability and parsed by 'Control.Concurrent.forkIO' +workers. The library never chooses RTS settings — to actually get +parallel reads, link the executable with @-threaded@ and run with +@+RTS -N@ (otherwise the reader falls back to the sequential path). +-} module DataFrame.IO.CSV.Fast ( fastReadCsv, readCsvFast, @@ -15,68 +22,35 @@ fastReadCsvFromBytesWithSchema, readSeparated, readSeparatedFromBytes, + readSeparatedFromBytesChunks, getDelimiterIndices, CsvParseError (..), ) where -import qualified Data.Vector as Vector -import qualified Data.Vector.Storable as VS -import Data.Vector.Storable.Mutable ( - grow, - unsafeFromForeignPtr, - ) -import qualified Data.Vector.Storable.Mutable as VSM -import System.IO.MMap ( - Mode (WriteCopy), - mmapFileForeignPtr, - ) - -import Control.Exception (Exception, throwIO) -import Control.Monad (when) -import Foreign ( - Ptr, - castForeignPtr, - castPtr, - ) -import qualified Foreign -import Foreign.C.Types -import Foreign.Marshal.Alloc (alloca) -import Foreign.Marshal.Utils (copyBytes) -import Foreign.Storable (peek, poke) - import qualified Data.ByteString as BS -import Data.ByteString.Internal (ByteString (PS)) 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.Text (Text) -import qualified Data.Text as Text -import qualified Data.Text.Encoding as TextEncoding import Data.Word (Word8) - -import Control.Parallel.Strategies (parList, rpar, using) -import Data.Array.IArray (array, (!)) -import Data.Array.Unboxed (UArray) -import Data.Ix (range) +import System.IO.MMap (Mode (WriteCopy), mmapFileForeignPtr) import DataFrame.IO.CSV ( - HeaderSpec (..), - RaggedRowPolicy (..), ReadOptions (..), TypeSpec (..), - UnclosedQuotePolicy (..), defaultReadOptions, - schemaTypeMap, - shouldInferFromSample, - typeInferenceSampleSize, ) +import DataFrame.IO.CSV.Fast.Core (readSeparatedCore) +import DataFrame.IO.CSV.Fast.Index ( + CsvParseError (..), + comma, + getDelimiterIndices, + tab, + ) import DataFrame.Internal.DataFrame (DataFrame (..)) import DataFrame.Internal.Schema (Schema (..)) -import DataFrame.Operations.Typing ( - ParseOptions (..), - effectiveSafeRead, - parseFromExamples, - parseWithTypes, - ) readSeparatedDefault :: Word8 -> FilePath -> IO DataFrame readSeparatedDefault separator = @@ -106,9 +80,8 @@ fastReadTsvWithOpts = readSeparated tab {- | Read a CSV and coerce each column to the type declared in the -supplied 'Schema'. Columns mentioned in the schema bypass inference; -columns absent from the schema fall back to the default inference -path. Use when the schema is known (DDL, prior runs) — it's both +supplied 'Schema'. Schema columns bypass inference entirely: they are +parsed straight from the file bytes as the declared type, which is both faster than inference and guards against the row-1 \"looks like Int\" misclassification trap. -} @@ -121,13 +94,6 @@ SpecifyTypes (M.toList (elements schema)) (typeSpec defaultReadOptions) } -{- | 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 'extractField' and 'parseFromExamples' are skipped for -unreferenced columns, so the end-to-end cost scales with the size of -the projection rather than the row width. --} - -- | In-memory version of 'fastReadCsv'. fastReadCsvFromBytes :: BS.ByteString -> IO DataFrame fastReadCsvFromBytes = readSeparatedFromBytes comma defaultReadOptions @@ -141,6 +107,10 @@ SpecifyTypes (M.toList (elements schema)) (typeSpec defaultReadOptions) } +{- | 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. +-} fastReadCsvProj :: [Text] -> FilePath -> IO DataFrame fastReadCsvProj projection path = do df <- fastReadCsv path @@ -165,20 +135,22 @@ FilePath -> IO DataFrame readSeparated separator opts filePath = do - -- We use write copy mode so that we can append - -- padding to the end of the memory space + -- WriteCopy keeps the mapping private; the scanner never needs to + -- grow or pad the buffer, so the file is no longer copied (audit F6). (bufferPtr, offset, len) <- mmapFileForeignPtr filePath WriteCopy Nothing - let mutableFile = unsafeFromForeignPtr bufferPtr offset len - readSeparatedFromMVec separator opts mutableFile len + readSeparatedCore + Nothing + separator + opts + (VS.unsafeFromForeignPtr bufferPtr offset len) -{- | Identical to 'readSeparated' but takes an already-loaded byte buffer -(e.g. a slice of a memory-mapped file). The bytes are copied once into -a fresh mutable storable vector so the SIMD code can grow it by 64 -bytes of trailing padding without disturbing the caller's input. +{- | 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. -} readSeparatedFromBytes :: Word8 -> @@ -186,487 +158,23 @@ BS.ByteString -> IO DataFrame readSeparatedFromBytes separator opts bs = do - let len = BS.length bs - mvec <- VSM.unsafeNew len - -- Copy ByteString → mutable vector (one pass, ~1 ns/byte). - let (fp, off, _) = BSI.toForeignPtr bs - Foreign.withForeignPtr fp $ \srcPtr -> - VSM.unsafeWith mvec $ \dstPtr -> - copyBytes (castPtr dstPtr) (srcPtr `Foreign.plusPtr` off) len - readSeparatedFromMVec separator opts mvec len + let (fp, off, len) = BSI.toForeignPtr bs + readSeparatedCore Nothing separator opts (VS.unsafeFromForeignPtr fp off len) -{- | Shared core: takes a mutable storable vector containing exactly @len@ -bytes of CSV (no trailing padding required by the caller — it is added -here) and produces a 'DataFrame'. +{- | '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. -} -readSeparatedFromMVec :: +readSeparatedFromBytesChunks :: + Int -> Word8 -> ReadOptions -> - VSM.IOVector Word8 -> - Int -> + BS.ByteString -> IO DataFrame -readSeparatedFromMVec separator opts mutableFile len = do - paddedMutableFile <- grow mutableFile 64 - paddedCSVFileRaw <- VS.unsafeFreeze paddedMutableFile - let (paddedCSVFile, bomLen) = stripBom paddedCSVFileRaw - contentLen = len - bomLen - indices <- - getDelimiterIndicesPolicy - (fastCsvOnUnclosedQuote opts) - separator - contentLen - paddedCSVFile - let rowEnds = classifyRowEnds paddedCSVFile contentLen indices - totalRows = VS.length rowEnds - if totalRows == 0 - then return emptyDataFrame - else do - let headerNumCol = fieldsInRow rowEnds 0 - trim = fastCsvTrimUnquoted opts - extractAt = - extractField trim paddedCSVFile indices rowEnds contentLen - (columnNames, dataStartRow, numCol) = case headerSpec opts of - NoHeader -> - ( Vector.fromList $ - map (Text.pack . show) [0 .. headerNumCol - 1] - , 0 - , headerNumCol - ) - UseFirstRow -> - ( Vector.fromList $ - map (extractAt 0) [0 .. headerNumCol - 1] - , 1 - , headerNumCol - ) - ProvideNames ns -> - (Vector.fromList ns, 0, length ns) - if numCol == 0 - then return emptyDataFrame - else do - let dataRows = - collectDataRows - paddedCSVFile - indices - rowEnds - contentLen - dataStartRow - totalRows - numRow = VS.length dataRows - when (fastCsvOnRaggedRow opts == RaiseOnRagged) $ - checkNoRaggedRows rowEnds dataRows numCol - let parseTypes name col = - let n = - if shouldInferFromSample (typeSpec opts) - then typeInferenceSampleSize (typeSpec opts) - else 0 - mode = - effectiveSafeRead - (safeRead opts) - (safeReadOverrides opts) - name - parseOpts = - ParseOptions - { missingValues = missingIndicators opts - , sampleSize = n - , parseSafe = mode - , parseSafeOverrides = [] - , parseDateFormat = dateFormat opts - } - in parseFromExamples parseOpts col - generateColumn col = - parseTypes (columnNames Vector.! col) $ - Vector.generate numRow $ \i -> - extractAt (dataRows VS.! i) col - cols = - Vector.fromListN - numCol - ( map generateColumn [0 .. numCol - 1] - `using` parList rpar - ) - colIndices = - M.fromList $ - zip (Vector.toList columnNames) [0 ..] - dfDims = (numRow, numCol) - let rawDf = - DataFrame - cols - colIndices - dfDims - M.empty - schemaMap = schemaTypeMap (typeSpec opts) - resolveMode = - effectiveSafeRead - (safeRead opts) - (safeReadOverrides opts) - return $! - if M.null schemaMap - then rawDf - else parseWithTypes resolveMode schemaMap rawDf - -{- | An empty 'DataFrame' — returned when the input has no delimiters -(empty file or single line with no separator and no newline). Guards -against a divide-by-zero in the row-stride math when 'numCol == 0'. --} -emptyDataFrame :: DataFrame -emptyDataFrame = DataFrame Vector.empty M.empty (0, 0) M.empty - -{- | Strip a leading UTF-8 BOM (EF BB BF) if present. Returns the trimmed -vector and the number of bytes removed (0 or 3). --} -{-# INLINE stripBom #-} -stripBom :: VS.Vector Word8 -> (VS.Vector Word8, Int) -stripBom v - | VS.length v >= 3 - , VS.unsafeIndex v 0 == 0xEF - , VS.unsafeIndex v 1 == 0xBB - , VS.unsafeIndex v 2 == 0xBF = - (VS.drop 3 v, 3) - | otherwise = (v, 0) - -{- | Classify each entry in the flat delimiter vector as either a row -terminator or a field terminator, and return the indices-into-the-vector -of every row terminator. A position counts as a row break if either - - * the byte at that position is @\\n@, or - * the position is beyond the original content length, which only - happens for the synthetic end-of-file delimiter written when a file - does not end in a newline. - -Field terminators (commas / tabs / the configured separator) are left -implicit: anything between consecutive row terminators is a field. --} -{-# INLINE classifyRowEnds #-} -classifyRowEnds :: VS.Vector Word8 -> Int -> VS.Vector CSize -> VS.Vector Int -classifyRowEnds file contentLen = - VS.findIndices isRowBreak - where - isRowBreak pos = - let p = fromIntegral pos :: Int - in p >= contentLen || VS.unsafeIndex file p == lf - -{- | Number of fields that row @r@ contains, derived directly from the -gap between consecutive entries in @rowEnds@. --} -{-# INLINE fieldsInRow #-} -fieldsInRow :: VS.Vector Int -> Int -> Int -fieldsInRow rowEnds r = - let endIdx = VS.unsafeIndex rowEnds r - startIdx = if r == 0 then 0 else VS.unsafeIndex rowEnds (r - 1) + 1 - in endIdx - startIdx + 1 - -{- | Does row @r@ contain exactly one, empty field? That is the case for -a bare @\\n@ or a @\\r\\n@ \u2014 blank lines, which we skip by default to -match pandas / polars @skip_blank_lines@ semantics. We also treat a -row whose only byte is @\\r@ as blank, so CRLF-terminated blank lines -don't leak through when only the @\\n@ counts as a row break. --} -{-# INLINE isBlankRow #-} -isBlankRow :: - VS.Vector Word8 -> - VS.Vector CSize -> - VS.Vector Int -> - Int -> - Int -> - Bool -isBlankRow file delimiters rowEnds contentLen r = - let endIdx = VS.unsafeIndex rowEnds r - startIdx = if r == 0 then 0 else VS.unsafeIndex rowEnds (r - 1) + 1 - numFields = endIdx - startIdx + 1 - in numFields == 1 - && let fieldEndRaw = - fromIntegral (VS.unsafeIndex delimiters startIdx) :: Int - fieldEnd = min fieldEndRaw contentLen - fieldStart = - if startIdx == 0 - then 0 - else - fromIntegral - (VS.unsafeIndex delimiters (startIdx - 1)) - + 1 - fieldLen = fieldEnd - fieldStart - in fieldLen == 0 - || ( fieldLen == 1 - && VS.unsafeIndex file fieldStart == cr - ) - -{- | Select the row indices that contain actual data: skip @[0 .. skip - 1]@ -and drop any blank rows from the remainder. --} -{-# INLINE collectDataRows #-} -collectDataRows :: - VS.Vector Word8 -> - VS.Vector CSize -> - VS.Vector Int -> - Int -> - Int -> - Int -> - VS.Vector Int -collectDataRows file delimiters rowEnds contentLen skip total = - VS.filter - (not . isBlankRow file delimiters rowEnds contentLen) - (VS.generate (total - skip) (+ skip)) - -{- | Extract field @col@ of row @r@ as a 'Text'. - -Semantics follow RFC 4180: - - * If the raw field is wrapped in @\"..\"@, the outer quotes are - stripped and any embedded @\"\"@ is unescaped to a single @\"@. - Whitespace inside the quotes is preserved verbatim. - - * If the raw field is unquoted, it is returned as-is; whitespace - is preserved by default (matching pandas / polars). Callers that - want legacy trim-everything behaviour pass @trimUnquoted = True@. - - * A trailing @\\r@ is dropped before decoding so CRLF files produce - clean Text on every column, not only when @Text.strip@ happens to - sweep it up. - -If @col@ is out of range (a ragged short row), returns the empty text; -callers convert that into a null via the missing-indicator list that -'parseFromExamples' honours. --} -{-# INLINE extractField #-} -extractField :: - Bool -> - VS.Vector Word8 -> - VS.Vector CSize -> - VS.Vector Int -> - Int -> - Int -> - Int -> - Text -extractField trimUnquoted file delimiters rowEnds contentLen r col - | col >= numFields = Text.empty - | fieldEnd - fieldStart >= 2 - , VS.unsafeIndex file fieldStart == quote - , VS.unsafeIndex file (fieldEnd - 1) == quote = - unescapeDoubledQuotes - . TextEncoding.decodeUtf8Lenient - . unsafeToByteString - $ VS.slice (fieldStart + 1) (fieldEnd - fieldStart - 2) file - | otherwise = - (if trimUnquoted then Text.strip else id) - . TextEncoding.decodeUtf8Lenient - . unsafeToByteString - $ VS.slice fieldStart (fieldEnd - fieldStart) file - where - endIdx = VS.unsafeIndex rowEnds r - startIdx = if r == 0 then 0 else VS.unsafeIndex rowEnds (r - 1) + 1 - numFields = endIdx - startIdx + 1 - boundaryIdx = startIdx + col - fieldEndRaw = - fromIntegral (VS.unsafeIndex delimiters boundaryIdx) :: Int - fieldEndClamped = min fieldEndRaw contentLen - fieldStart = - if boundaryIdx == 0 - then 0 - else - fromIntegral - (VS.unsafeIndex delimiters (boundaryIdx - 1)) - + 1 - -- Strip a trailing \r (CRLF line endings) from the last field of a row - -- when the SIMD scanner, which only tracks \n, leaves it behind. - fieldEnd = - if fieldEndClamped > fieldStart - && VS.unsafeIndex file (fieldEndClamped - 1) == cr - then fieldEndClamped - 1 - else fieldEndClamped - unsafeToByteString :: VS.Vector Word8 -> BS.ByteString - unsafeToByteString v = PS (castForeignPtr ptr) 0 n - where - (ptr, n) = VS.unsafeToForeignPtr0 v - -{- | Walk every data row and throw 'CsvRaggedRow' on the first one whose -field count differs from the header. Only called when the user opts -into 'RaiseOnRagged' via 'fastCsvOnRaggedRow'. --} -{-# INLINE checkNoRaggedRows #-} -checkNoRaggedRows :: - VS.Vector Int -> - VS.Vector Int -> - Int -> - IO () -checkNoRaggedRows rowEnds dataRows numCol = - VS.mapM_ - ( \r -> - let actual = fieldsInRow rowEnds r - in when (actual /= numCol) $ - throwIO (CsvRaggedRow r numCol actual) - ) - dataRows - -{- | RFC 4180 inner-quote unescape: @\"\"@ → @\"@. 'Text.replace' on a -two-char needle does a single linear pass and is allocation-free when -no doubled quote is present (short-circuits at the first miss). --} -{-# INLINE unescapeDoubledQuotes #-} -unescapeDoubledQuotes :: Text -> Text -unescapeDoubledQuotes t - | Text.isInfixOf doubledQuote t = Text.replace doubledQuote singleQuote t - | otherwise = t - where - doubledQuote = Text.pack "\"\"" - singleQuote = Text.singleton '"' - -{- | Exceptions raised by the fast CSV parser. Catchable with -'Control.Exception.catch' or 'Control.Exception.try'. --} -data CsvParseError - = {- | The input ends with a quoted field that was never closed. The - PCLMUL quote-parity chain in the SIMD scanner treats an unmatched - @\"@ as if the rest of the file were inside quotes, so we refuse - to return a silently corrupted 'DataFrame' and raise instead. - -} - CsvUnclosedQuote - | {- | A row has a different number of fields from the header. Only - raised when 'fastCsvOnRaggedRow' is set to 'RaiseOnRagged'. - Carries the 0-based row index, the expected field count, and the - actual field count. - -} - CsvRaggedRow !Int !Int !Int - deriving (Eq, Show) - -instance Exception CsvParseError - --- Status codes reported via the out-parameter of 'get_delimiter_indices'. --- Must stay in sync with @GDI_*@ in @cbits/process_csv.h@. -gdiOk, gdiUnclosedQuote :: CInt -gdiOk = 0 -gdiUnclosedQuote = 1 - -{- | Return value that the C helper uses to tell Haskell \"SIMD isn't -available on this build / CPU, run the pure-Haskell state machine.\" -Mirrors @GDI_SIMD_UNAVAILABLE@ in @cbits/process_csv.h@. 'CSize' is -unsigned so this has the same bit pattern as 'maxBound'. --} -simdUnavailable :: CSize -simdUnavailable = maxBound - -foreign import capi "process_csv.h get_delimiter_indices" - get_delimiter_indices :: - Ptr CUChar -> -- input - CSize -> -- input size - CUChar -> -- separator character - Ptr CSize -> -- result array - Ptr CInt -> -- status out-parameter - IO CSize -- occupancy of result array - -{- | Locate delimiter byte positions in @csvFile@. Treats an unclosed -quoted field at EOF as a hard error; callers that want to suppress the -exception can use 'getDelimiterIndicesPolicy' with 'BestEffort'. --} -{-# INLINE getDelimiterIndices #-} -getDelimiterIndices :: - Word8 -> - Int -> - VS.Vector Word8 -> - IO (VS.Vector CSize) -getDelimiterIndices = getDelimiterIndicesPolicy RaiseOnUnclosedQuote - -{-# INLINE getDelimiterIndicesPolicy #-} -getDelimiterIndicesPolicy :: - UnclosedQuotePolicy -> - Word8 -> - Int -> - VS.Vector Word8 -> - IO (VS.Vector CSize) -getDelimiterIndicesPolicy policy separator originalLen csvFile = - VS.unsafeWith csvFile $ \buffer -> do - let paddedLen = VS.length csvFile - -- GC-managed pinned memory: freed automatically, no leak in streaming use. - resultMV <- VSM.unsafeNew paddedLen - (num_fields, status) <- alloca $ \statusPtr -> do - poke statusPtr gdiOk - n <- - VSM.unsafeWith resultMV $ \indicesPtr -> - get_delimiter_indices - (castPtr buffer) - (fromIntegral paddedLen) - (fromIntegral separator) - (castPtr indicesPtr) - statusPtr - s <- peek statusPtr - return (n, s) - if num_fields == simdUnavailable - then do - -- Haskell state-machine fallback, writing directly into resultMV. - let trans = stateTransitionTable separator - processChar (!state, !idx) i byte = - case state of - UnEscaped -> - if byte == lf || byte == separator - then do - VSM.unsafeWrite resultMV idx (fromIntegral i) - return (toEnum (trans ! (fromEnum state, byte)), idx + 1) - else return (toEnum (trans ! (fromEnum state, byte)), idx) - Escaped -> - return (toEnum (trans ! (fromEnum state, byte)), idx) - (finalState, finalIdx) <- - VS.ifoldM' processChar (UnEscaped, 0 :: Int) csvFile - case finalState of - Escaped - | policy == RaiseOnUnclosedQuote -> - throwIO CsvUnclosedQuote - _ -> return () - finalLen <- - if originalLen > 0 && csvFile VS.! (originalLen - 1) /= lf - then do - VSM.unsafeWrite resultMV finalIdx (fromIntegral originalLen) - return (finalIdx + 1) - else return finalIdx - VS.unsafeFreeze (VSM.slice 0 finalLen resultMV) - else do - when - ( status == gdiUnclosedQuote - && policy == RaiseOnUnclosedQuote - ) - (throwIO CsvUnclosedQuote) - let n = fromIntegral num_fields - finalLen <- - if originalLen > 0 && csvFile VS.! (originalLen - 1) /= lf - then do - VSM.write resultMV n (fromIntegral originalLen) - return (n + 1) - else return n - VS.unsafeFreeze (VSM.slice 0 finalLen resultMV) - --- We have a Native version in case the C version --- cannot be used. For example if neither ARM_NEON --- nor AVX2 are available - -lf, cr, comma, tab, quote :: Word8 -lf = 0x0A -cr = 0x0D -comma = 0x2C -tab = 0x09 -quote = 0x22 - --- We parse using a state machine -data State - = UnEscaped -- non quoted - | Escaped -- quoted - deriving (Enum) - -{-# INLINE stateTransitionTable #-} -stateTransitionTable :: Word8 -> UArray (Int, Word8) Int -stateTransitionTable separator = array ((0, 0), (1, 255)) [(i, f i) | i <- range ((0, 0), (1, 255))] - where - f (0, character) - -- Unescaped newline - | character == 0x0A = fromEnum UnEscaped - -- Unescaped separator - | character == separator = fromEnum UnEscaped - -- Unescaped quote - | character == 0x22 = fromEnum Escaped - | otherwise = fromEnum UnEscaped - -- Escaped quote - -- escaped quote in fields are dealt as - -- consecutive quoted sections of a field - -- example: If we have - -- field1, "abc""def""ghi, field3 - -- we end up processing abc, def, and ghi - -- as consecutive quoted strings. - f (1, 0x22) = fromEnum UnEscaped - -- Everything else - f (state, _) = state +readSeparatedFromBytesChunks nChunks separator opts bs = do + let (fp, off, len) = BSI.toForeignPtr bs + readSeparatedCore + (Just nChunks) + separator + opts + (VS.unsafeFromForeignPtr fp off len) diff --git a/src/DataFrame/IO/CSV/Fast/Columns.hs b/src/DataFrame/IO/CSV/Fast/Columns.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/CSV/Fast/Columns.hs @@ -0,0 +1,254 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeApplications #-} + +{- | Column-level dispatch for the fast CSV reader: resolve each column's +target type up front into a 'ColumnPlan' (schema hit -> typed pass directly; +miss -> sample classification + typed fallback chain; EitherRead \/ exotic +schema types -> the legacy Text pipeline). Plans are chunk-runnable so the +parallel reader (WS-E2) can execute them over row sub-ranges. +-} +module DataFrame.IO.CSV.Fast.Columns ( + ColumnEnv (..), + ColumnPlan (..), + Step, + TargetType, + ChunkOutcome (..), + buildColumn, + planColumn, + runChainChunk, + runSchemaChunk, + runStep, + legacyColumn, + allNullColumn, + finishMode, + schemaError, +) where + +import qualified Data.Map as M +import qualified Data.Proxy as P +import qualified Data.Text as T +import qualified Data.Vector as V + +import Control.Exception (ErrorCall (..), throwIO) +import Data.Maybe (isJust) +import Data.Time (Day) +import Data.Type.Equality (testEquality) +import Type.Reflection (typeRep) + +import DataFrame.IO.CSV ( + ReadOptions (..), + schemaTypeMap, + shouldInferFromSample, + typeInferenceSampleSize, + ) +import DataFrame.IO.CSV.Fast.Passes +import DataFrame.IO.CSV.Fast.Slice (extractField) +import DataFrame.Internal.Column (Column, ensureOptional, fromVector) +import DataFrame.Internal.Schema (SchemaType (..)) +import DataFrame.Operations.Typing ( + ParseOptions (..), + ParsingAssumption (..), + SafeReadMode (..), + effectiveSafeRead, + isNullishOrMissing, + makeParsingAssumption, + parseFromExamples, + ) + +-- | One step of a type-fallback chain (always ends in 'StepText'). +data Step = StepBool | StepInt | StepDouble | StepDate | StepText + +-- | Run one chain step over the env's row range as a single typed pass. +runStep :: ColumnEnv -> NullSpec -> Int -> Bool -> Step -> IO PassResult +runStep env nspec col nullOnFail step = case step of + StepBool -> boolPass env nspec col nullOnFail + StepInt -> intPass env nspec col nullOnFail + StepDouble -> doublePass env nspec col nullOnFail + StepDate -> datePass env nspec col nullOnFail + StepText -> Right . PassText <$> textPass env nspec col + +-- | The documented fallback chain for each inference assumption. +chainFor :: ParsingAssumption -> [Step] +chainFor BoolAssumption = [StepBool, StepText] +chainFor IntAssumption = [StepInt, StepDouble, StepText] +chainFor DoubleAssumption = [StepDouble, StepText] +chainFor DateAssumption = [StepDate, StepText] +chainFor TextAssumption = [StepText] +chainFor NoAssumption = [StepBool, StepInt, StepDouble, StepDate, StepText] + +-- | Schema target types served by the direct typed passes. +data TargetType = TInt | TDouble | TText | TBool | TDate + +targetName :: TargetType -> String +targetName TInt = "Int" +targetName TDouble = "Double" +targetName TText = "Text" +targetName TBool = "Bool" +targetName TDate = "Day" + +stepForTarget :: TargetType -> Step +stepForTarget TInt = StepInt +stepForTarget TDouble = StepDouble +stepForTarget TText = StepText +stepForTarget TBool = StepBool +stepForTarget TDate = StepDate + +supportedType :: SchemaType -> Maybe TargetType +supportedType (SType (_ :: P.Proxy a)) + | isJust (testEquality (typeRep @a) (typeRep @Int)) = Just TInt + | isJust (testEquality (typeRep @a) (typeRep @Double)) = Just TDouble + | isJust (testEquality (typeRep @a) (typeRep @T.Text)) = Just TText + | isJust (testEquality (typeRep @a) (typeRep @Bool)) = Just TBool + | isJust (testEquality (typeRep @a) (typeRep @Day)) = Just TDate + | otherwise = Nothing + +{- | How to build one column, resolved once per read (before any chunk +fan-out). 'PlanLegacy' carries the \"re-apply 'parseWithTypes'\" flag. +-} +data ColumnPlan + = {- | Inference: fallback chain; the 'Bool' asks for the all-null + pre-check ('NoAssumption' only). + -} + PlanChain !SafeReadMode !NullSpec ![Step] !Bool + | PlanSchema !SafeReadMode !NullSpec !TargetType + | PlanLegacy !SafeReadMode !Bool + +{- | Resolve a column's plan. For inferred columns the sample is always +classified against the global (full-range) env, so chunked and sequential +reads share one assumption. +-} +planColumn :: ColumnEnv -> T.Text -> Int -> ColumnPlan +planColumn env name col = case mode of + EitherRead -> PlanLegacy mode (isJust schemaEntry) + _ -> case schemaEntry of + Nothing -> chainPlan env mode col + Just st -> case supportedType st of + Just t -> PlanSchema mode (nullSpecFor opts mode) t + Nothing -> PlanLegacy mode True + where + opts = ceOpts env + mode = effectiveSafeRead (safeRead opts) (safeReadOverrides opts) name + schemaEntry = M.lookup name (schemaTypeMap (typeSpec opts)) + +-- | Classify a row sample with the existing assumption logic. +chainPlan :: ColumnEnv -> SafeReadMode -> Int -> ColumnPlan +chainPlan env mode col = + PlanChain mode (nullSpecFor opts mode) (chainFor assumption) noAssumption + where + opts = ceOpts env + ts = typeSpec opts + sampleN = + if shouldInferFromSample ts then typeInferenceSampleSize ts else 0 + isNullT = case mode of + NoSafeRead -> T.null + _ -> isNullishOrMissing (missingIndicators opts) + sample = + V.generate + (min sampleN (ceNumRow env)) + (\i -> extractField (ceCtx env) (rowAt env i) col) + classified = + V.map (\t -> if isNullT t then Nothing else Just t) sample + assumption = makeParsingAssumption (dateFormat opts) classified + noAssumption = case assumption of + NoAssumption -> True + _ -> False + +{- | What a chain plan produced for one row range: every cell null (only +reported when the all-null pre-check is on), or the index of the first +chain step whose pass succeeded plus its column. +-} +data ChunkOutcome = AllNullChunk | Resolved !Int !PassCol + +-- | Run a fallback chain over the env's row range. +runChainChunk :: + ColumnEnv -> NullSpec -> [Step] -> Bool -> Int -> IO ChunkOutcome +runChainChunk env nspec steps checkAllNull col = do + allNull <- if checkAllNull then allNullPass env nspec col else pure False + if allNull then pure AllNullChunk else go 0 steps + where + go i (s : rest) = do + r <- runStep env nspec col False s + case r of + Right c -> pure (Resolved i c) + Left _ -> go (i + 1) rest + go _ [] = error "fastcsv: type-fallback chain exhausted (StepText must succeed)" + +{- | Run a schema-typed pass over the env's row range. @rowOff@ maps a +chunk-local failure index to the global data-row index. +-} +runSchemaChunk :: + ColumnEnv -> + NullSpec -> + SafeReadMode -> + TargetType -> + Int -> + Int -> + IO (Either Int PassCol) +runSchemaChunk env nspec mode t col rowOff = do + r <- runStep env nspec col (mode == MaybeRead) (stepForTarget t) + pure (either (Left . (+ rowOff)) Right r) + +-- | The error a schema column raises when a cell fails under 'NoSafeRead'. +schemaError :: T.Text -> TargetType -> Int -> ErrorCall +schemaError name t i = + ErrorCall $ + "fastcsv: schema column " + <> show name + <> ": row " + <> show i + <> " does not parse as " + <> targetName t + +-- | The all-null column an all-null range resolves to (Maybe Text). +allNullColumn :: Int -> Column +allNullColumn n = fromVector (V.replicate n (Nothing :: Maybe T.Text)) + +-- | 'MaybeRead' columns keep their Maybe-ness even when fully valid. +finishMode :: SafeReadMode -> Column -> Column +finishMode MaybeRead = ensureOptional +finishMode _ = id + +{- | Build one column sequentially (single chunk over the full row range). +Returns the column plus a flag telling the caller to re-apply the legacy +schema conversion ('parseWithTypes') afterwards — only set for schema'd +columns the typed passes cannot serve. +-} +buildColumn :: ColumnEnv -> T.Text -> Int -> IO (Column, Bool) +buildColumn env name col = case planColumn env name col of + PlanLegacy mode pwt -> (,pwt) <$> legacyColumn env mode col + PlanSchema mode nspec t -> do + r <- runSchemaChunk env nspec mode t col 0 + c <- either (throwIO . schemaError name t) (pure . passColumn) r + pure (finishMode mode c, False) + PlanChain mode nspec steps checkAllNull -> do + oc <- runChainChunk env nspec steps checkAllNull col + let c = case oc of + AllNullChunk -> allNullColumn (ceNumRow env) + Resolved _ c' -> passColumn c' + pure (finishMode mode c, False) + +{- | The original Text-materializing pipeline, kept for 'EitherRead' and +exotic schema types (cold paths). +-} +legacyColumn :: ColumnEnv -> SafeReadMode -> Int -> IO Column +legacyColumn env mode col = do + let opts = ceOpts env + ts = typeSpec opts + parseOpts = + ParseOptions + { missingValues = missingIndicators opts + , sampleSize = + if shouldInferFromSample ts + then typeInferenceSampleSize ts + else 0 + , parseSafe = mode + , parseSafeOverrides = [] + , parseDateFormat = dateFormat opts + } + cells = + V.generate + (ceNumRow env) + (\i -> extractField (ceCtx env) (rowAt env i) col) + pure (parseFromExamples parseOpts cells) diff --git a/src/DataFrame/IO/CSV/Fast/Core.hs b/src/DataFrame/IO/CSV/Fast/Core.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/CSV/Fast/Core.hs @@ -0,0 +1,187 @@ +{- | 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 +capabilities allow ('autoChunkCount'), unless the caller pins a chunk +count. The result is fully forced before it is returned. +-} +module DataFrame.IO.CSV.Fast.Core ( + readSeparatedCore, +) where + +import qualified Data.Map as M +import qualified Data.Set as S +import qualified Data.Text as Text +import qualified Data.Vector as Vector +import qualified Data.Vector.Storable as VS + +import Control.Exception (throwIO) +import Control.Monad (when) +import Data.Word (Word8) +import Debug.Trace (traceMarkerIO) + +import DataFrame.IO.CSV ( + HeaderSpec (..), + RaggedRowPolicy (..), + ReadOptions (..), + UnclosedQuotePolicy (..), + schemaTypeMap, + ) +import DataFrame.IO.CSV.Fast.Index ( + CsvParseError (..), + byteStringView, + checkNoRaggedRows, + classifyRowEnds, + collectDataRows, + fieldsInRow, + getDelimiterIndicesPolicy, + ) +import DataFrame.IO.CSV.Fast.IndexPar ( + ParIndex (..), + getDelimiterIndicesPar, + ) +import DataFrame.IO.CSV.Fast.Parallel (autoChunkCount, buildAllColumns) +import DataFrame.IO.CSV.Fast.Passes (ColumnEnv (..)) +import DataFrame.IO.CSV.Fast.Slice (FieldCtx (..), extractField, stripBom) +import DataFrame.Internal.DataFrame (DataFrame (..), forceDataFrame) +import DataFrame.Operations.Typing (effectiveSafeRead, parseWithTypes) + +{- | @readSeparatedCore chunkSpec separator opts buffer@; @chunkSpec@ +pins the chunk count ('Nothing' = capability\/size-gated automatic). +-} +readSeparatedCore :: + Maybe Int -> + Word8 -> + ReadOptions -> + VS.Vector Word8 -> + IO DataFrame +readSeparatedCore chunkSpec separator opts fileRaw = do + let (file, _bomLen) = stripBom fileRaw + contentLen = VS.length file + nChunks <- maybe (autoChunkCount contentLen) pure chunkSpec + -- Parallel reads scan in parallel too; sequential reads (or builds + -- without SIMD) keep the original scan + row classification. + mPar <- + if nChunks > 1 + then getDelimiterIndicesPar nChunks separator contentLen file + else pure Nothing + (indices, rowEnds) <- case mPar of + Just par -> do + when + ( piEndsInQuote par + && fastCsvOnUnclosedQuote opts == RaiseOnUnclosedQuote + ) + (throwIO CsvUnclosedQuote) + pure (piDelims par, piRowEnds par) + Nothing -> do + is <- + getDelimiterIndicesPolicy + (fastCsvOnUnclosedQuote opts) + separator + contentLen + file + pure (is, classifyRowEnds file contentLen is) + -- Eventlog phase markers (visible under +RTS -l; no-ops otherwise). + traceMarkerIO "fastcsv:scan-done" + let totalRows = VS.length rowEnds + if totalRows == 0 + then return emptyDataFrame + else do + let ctx = + FieldCtx + { fcFile = file + , fcBS = byteStringView contentLen file + , fcDelims = indices + , fcRowEnds = rowEnds + , fcContentLen = contentLen + , fcTrim = fastCsvTrimUnquoted opts + } + headerNumCol = fieldsInRow rowEnds 0 + (columnNames, dataStartRow, numCol) = case headerSpec opts of + NoHeader -> + ( Vector.fromList $ + map (Text.pack . show) [0 .. headerNumCol - 1] + , 0 + , headerNumCol + ) + UseFirstRow -> + ( Vector.fromList $ + map (extractField ctx 0) [0 .. headerNumCol - 1] + , 1 + , headerNumCol + ) + ProvideNames ns -> + (Vector.fromList ns, 0, length ns) + if numCol == 0 + then return emptyDataFrame + else do + let dataRows = + collectDataRows + file + indices + rowEnds + contentLen + dataStartRow + totalRows + numRow = VS.length dataRows + when (fastCsvOnRaggedRow opts == RaiseOnRagged) $ + checkNoRaggedRows rowEnds dataRows numCol + traceMarkerIO "fastcsv:rows-done" + VS.unsafeWith file $ \filePtr -> do + let env = + ColumnEnv + { ceCtx = ctx + , ceRows = dataRows + , ceNumRow = numRow + , cePtr = filePtr + , ceTextHint = + max 64 (contentLen `div` numCol) + , ceOpts = opts + } + results <- buildAllColumns nChunks env columnNames numCol + traceMarkerIO "fastcsv:build-done" + let cols = Vector.fromListN numCol (map fst results) + legacySchema = + S.fromList + [ columnNames Vector.! c + | (c, (_, True)) <- zip [0 ..] results + ] + colIndices = + M.fromList $ + zip (Vector.toList columnNames) [0 ..] + df = + DataFrame + cols + colIndices + (numRow, numCol) + M.empty + resolveMode = + effectiveSafeRead + (safeRead opts) + (safeReadOverrides opts) + df' = + if S.null legacySchema + then df + else + parseWithTypes + resolveMode + ( M.restrictKeys + (schemaTypeMap (typeSpec opts)) + legacySchema + ) + df + -- The parallel merge already forced every column + -- payload, so only the column thunks themselves + -- need forcing; the deep walk would re-chase 6M+ + -- pointers per boxed column for nothing. + if nChunks > 1 && S.null legacySchema + then do + Vector.foldl' (\() c -> c `seq` ()) () cols `seq` + return df + else return $! forceDataFrame df' + +{- | An empty 'DataFrame' — returned when the input has no delimiters +(empty file or single line with no separator and no newline). Guards +against a divide-by-zero in the row-stride math when 'numCol == 0'. +-} +emptyDataFrame :: DataFrame +emptyDataFrame = DataFrame Vector.empty M.empty (0, 0) M.empty diff --git a/src/DataFrame/IO/CSV/Fast/Index.hs b/src/DataFrame/IO/CSV/Fast/Index.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/CSV/Fast/Index.hs @@ -0,0 +1,298 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CApiFFI #-} +{-# LANGUAGE ForeignFunctionInterface #-} + +{- | Delimiter-index construction for the fast CSV reader. The index buffer +is sized from an exact byte count (not 8 bytes per input byte), the SIMD +scan runs over the 64-byte-aligned prefix of the unpadded mmap, and a +scalar Haskell tail finishes the final partial chunk — so the input is +never copied or grown. +-} +module DataFrame.IO.CSV.Fast.Index ( + CsvParseError (..), + getDelimiterIndices, + getDelimiterIndicesPolicy, + scalarScan, + classifyRowEnds, + fieldsInRow, + isBlankRow, + collectDataRows, + checkNoRaggedRows, + byteStringView, + lf, + cr, + comma, + tab, + quote, +) where + +import qualified Data.ByteString as BS +import qualified Data.ByteString.Internal as BSI +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Storable.Mutable as VSM + +import Control.Exception (Exception, throwIO) +import Control.Monad (when) +import Data.Word (Word8) +import Foreign (Ptr, castPtr) +import Foreign.C.Types (CInt (..), CSize (..), CUChar (..)) +import Foreign.Marshal.Alloc (alloca) +import Foreign.Storable (peek, poke) + +import DataFrame.IO.CSV (UnclosedQuotePolicy (..)) + +{- | Exceptions raised by the fast CSV parser. Catchable with +'Control.Exception.catch' or 'Control.Exception.try'. +-} +data CsvParseError + = {- | The input ends with a quoted field that was never closed. The + PCLMUL quote-parity chain in the SIMD scanner treats an unmatched + @\"@ as if the rest of the file were inside quotes, so we refuse + to return a silently corrupted 'DataFrame.Internal.DataFrame.DataFrame' + and raise instead. + -} + CsvUnclosedQuote + | {- | A row has a different number of fields from the header. Only + raised when 'DataFrame.IO.CSV.fastCsvOnRaggedRow' is set to + 'RaiseOnRagged'. Carries the 0-based row index, the expected field + count, and the actual field count. + -} + CsvRaggedRow !Int !Int !Int + deriving (Eq, Show) + +instance Exception CsvParseError + +lf, cr, comma, tab, quote :: Word8 +lf = 0x0A +cr = 0x0D +comma = 0x2C +tab = 0x09 +quote = 0x22 + +-- Status codes reported via the out-parameter of 'get_delimiter_indices'. +-- Must stay in sync with @GDI_*@ in @cbits/process_csv.h@. +gdiOk, gdiUnclosedQuote :: CInt +gdiOk = 0 +gdiUnclosedQuote = 1 + +{- | Return value that the C helper uses to tell Haskell \"SIMD isn't +available on this build / CPU, run the pure-Haskell state machine.\" +Mirrors @GDI_SIMD_UNAVAILABLE@ in @cbits/process_csv.h@. 'CSize' is +unsigned so this has the same bit pattern as 'maxBound'. +-} +simdUnavailable :: CSize +simdUnavailable = maxBound + +foreign import capi "process_csv.h get_delimiter_indices" + get_delimiter_indices :: + Ptr CUChar -> -- input + CSize -> -- input size + CUChar -> -- separator character + Ptr CSize -> -- result array + Ptr CInt -> -- status out-parameter + IO CSize -- occupancy of result array + +-- | Zero-copy 'BS.ByteString' view of the first @n@ bytes of a vector. +byteStringView :: Int -> VS.Vector Word8 -> BS.ByteString +byteStringView n v = BSI.fromForeignPtr (fst (VS.unsafeToForeignPtr0 v)) 0 n +{-# INLINE byteStringView #-} + +{- | Locate delimiter byte positions in the first @originalLen@ bytes of +@csvFile@. Treats an unclosed quoted field at EOF as a hard error; callers +that want to suppress the exception can use 'getDelimiterIndicesPolicy' +with 'BestEffort'. +-} +{-# INLINE getDelimiterIndices #-} +getDelimiterIndices :: + Word8 -> + Int -> + VS.Vector Word8 -> + IO (VS.Vector CSize) +getDelimiterIndices = getDelimiterIndicesPolicy RaiseOnUnclosedQuote + +getDelimiterIndicesPolicy :: + UnclosedQuotePolicy -> + Word8 -> + Int -> + VS.Vector Word8 -> + IO (VS.Vector CSize) +getDelimiterIndicesPolicy policy separator contentLen csvFile = do + -- Exact upper bound on the index count: every output index is an + -- unquoted separator or newline byte, so two memchr-speed counts + -- (plus one synthetic EOF slot) right-size the buffer (audit F4). + let contentBS = byteStringView contentLen csvFile + capacity = BS.count separator contentBS + BS.count lf contentBS + 1 + resultMV <- VSM.unsafeNew capacity + -- The C scanner reads whole 64-byte chunks; cap it at the largest + -- chunk boundary inside the content so no padding is needed and the + -- input never has to be grown or copied (audit F6). The loop in C + -- runs while i < len - 64, hence the + 64. + let simdLen = (contentLen `div` 64) * 64 + (numFound, status) <- VS.unsafeWith csvFile $ \buffer -> + alloca $ \statusPtr -> do + poke statusPtr gdiOk + n <- VSM.unsafeWith resultMV $ \indicesPtr -> + get_delimiter_indices + (castPtr buffer) + (fromIntegral (simdLen + 64)) + (fromIntegral separator) + (castPtr indicesPtr) + statusPtr + s <- peek statusPtr + pure (n, s) + -- The status reports the quote parity at the end of the SIMD region, + -- which is exactly the starting state for the scalar tail. + (found, inQuote) <- + if numFound == simdUnavailable + then scalarScan separator csvFile resultMV 0 0 contentLen False + else + scalarScan + separator + csvFile + resultMV + (fromIntegral numFound) + simdLen + contentLen + (status == gdiUnclosedQuote) + when (inQuote && policy == RaiseOnUnclosedQuote) (throwIO CsvUnclosedQuote) + -- Synthetic EOF delimiter when the file does not end in a newline. + finalLen <- + if contentLen > 0 && VS.unsafeIndex csvFile (contentLen - 1) /= lf + then do + VSM.unsafeWrite resultMV found (fromIntegral contentLen) + pure (found + 1) + else pure found + VS.unsafeFreeze (VSM.slice 0 finalLen resultMV) + +{- | RFC 4180 quote-parity state machine over bytes @[from, to)@, appending +delimiter indices at @writeIdx@. Returns the next write index and whether +the scan ended inside a quoted region. Serves as both the SIMD tail and the +full fallback when SIMD is unavailable. +-} +scalarScan :: + Word8 -> + VS.Vector Word8 -> + VSM.IOVector CSize -> + Int -> + Int -> + Int -> + Bool -> + IO (Int, Bool) +scalarScan separator file out writeIdx from to = go writeIdx from + where + go !idx !i !inQ + | i >= to = pure (idx, inQ) + | otherwise = + let b = VS.unsafeIndex file i + in if inQ + then go idx (i + 1) (b /= quote) + else + if b == quote + then go idx (i + 1) True + else + if b == lf || b == separator + then do + VSM.unsafeWrite out idx (fromIntegral i) + go (idx + 1) (i + 1) False + else go idx (i + 1) False + +{- | Walk every data row and throw 'CsvRaggedRow' on the first one whose +field count differs from the header. Only called when the user opts +into 'RaiseOnRagged' via 'DataFrame.IO.CSV.fastCsvOnRaggedRow'. +-} +checkNoRaggedRows :: + VS.Vector Int -> + VS.Vector Int -> + Int -> + IO () +checkNoRaggedRows rowEnds dataRows numCol = + VS.mapM_ + ( \r -> + let actual = fieldsInRow rowEnds r + in when (actual /= numCol) $ + throwIO (CsvRaggedRow r numCol actual) + ) + dataRows + +{- | Classify each entry in the flat delimiter vector as either a row +terminator or a field terminator, and return the indices-into-the-vector +of every row terminator. A position counts as a row break if either + + * the byte at that position is @\\n@, or + * the position is beyond the original content length, which only + happens for the synthetic end-of-file delimiter written when a file + does not end in a newline. + +Field terminators (commas / tabs / the configured separator) are left +implicit: anything between consecutive row terminators is a field. +-} +{-# INLINE classifyRowEnds #-} +classifyRowEnds :: VS.Vector Word8 -> Int -> VS.Vector CSize -> VS.Vector Int +classifyRowEnds file contentLen = + VS.findIndices isRowBreak + where + isRowBreak pos = + let p = fromIntegral pos :: Int + in p >= contentLen || VS.unsafeIndex file p == lf + +{- | Number of fields that row @r@ contains, derived directly from the +gap between consecutive entries in @rowEnds@. +-} +{-# INLINE fieldsInRow #-} +fieldsInRow :: VS.Vector Int -> Int -> Int +fieldsInRow rowEnds r = + let endIdx = VS.unsafeIndex rowEnds r + startIdx = if r == 0 then 0 else VS.unsafeIndex rowEnds (r - 1) + 1 + in endIdx - startIdx + 1 + +{- | Does row @r@ contain exactly one, empty field? That is the case for +a bare @\\n@ or a @\\r\\n@ \x2014 blank lines, which we skip by default to +match pandas / polars @skip_blank_lines@ semantics. We also treat a +row whose only byte is @\\r@ as blank, so CRLF-terminated blank lines +don't leak through when only the @\\n@ counts as a row break. +-} +{-# INLINE isBlankRow #-} +isBlankRow :: + VS.Vector Word8 -> + VS.Vector CSize -> + VS.Vector Int -> + Int -> + Int -> + Bool +isBlankRow file delimiters rowEnds contentLen r = + let endIdx = VS.unsafeIndex rowEnds r + startIdx = if r == 0 then 0 else VS.unsafeIndex rowEnds (r - 1) + 1 + numFields = endIdx - startIdx + 1 + in numFields == 1 + && let fieldEndRaw = + fromIntegral (VS.unsafeIndex delimiters startIdx) :: Int + fieldEnd = min fieldEndRaw contentLen + fieldStart = + if startIdx == 0 + then 0 + else + fromIntegral + (VS.unsafeIndex delimiters (startIdx - 1)) + + 1 + fieldLen = fieldEnd - fieldStart + in fieldLen == 0 + || ( fieldLen == 1 + && VS.unsafeIndex file fieldStart == cr + ) + +{- | Select the row indices that contain actual data: skip @[0 .. skip - 1]@ +and drop any blank rows from the remainder. +-} +{-# INLINE collectDataRows #-} +collectDataRows :: + VS.Vector Word8 -> + VS.Vector CSize -> + VS.Vector Int -> + Int -> + Int -> + Int -> + VS.Vector Int +collectDataRows file delimiters rowEnds contentLen skip total = + VS.filter + (not . isBlankRow file delimiters rowEnds contentLen) + (VS.generate (total - skip) (+ skip)) diff --git a/src/DataFrame/IO/CSV/Fast/IndexPar.hs b/src/DataFrame/IO/CSV/Fast/IndexPar.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/CSV/Fast/IndexPar.hs @@ -0,0 +1,187 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CApiFFI #-} + +{- | Parallel delimiter-index construction (Round-2 WS-E2): the SIMD region +is split into 64-byte-aligned ranges, one per worker. Quote parity at each +range start comes from a prefix-xor of per-range quote counts (exact: every +quote byte toggles the scanner's PCLMUL parity chain), a parallel counting +pass right-sizes the flat output arrays and gives every worker an exact +write offset, and a parallel emit pass writes delimiter positions plus +row-end ordinals directly (no per-delimiter file-byte re-read). The scalar +tail and the unclosed-quote policy are unchanged from the sequential path. +-} +module DataFrame.IO.CSV.Fast.IndexPar ( + ParIndex (..), + getDelimiterIndicesPar, +) where + +import qualified Data.Bits as Bits +import qualified Data.ByteString as BS +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Storable.Mutable as VSM + +import Data.Word (Word8) +import Foreign (Ptr, castPtr) +import Foreign.C.Types (CInt (..), CSize (..), CUChar (..)) +import Foreign.Marshal.Alloc (alloca) +import Foreign.Marshal.Array (advancePtr) +import Foreign.Storable (peek) + +import Control.Concurrent (getNumCapabilities) +import DataFrame.IO.CSV.Fast.Index (byteStringView, lf, quote, scalarScan) +import DataFrame.IO.CSV.Fast.Workers (pooledRun) + +foreign import capi "process_csv.h count_delimiters_chunk" + count_delimiters_chunk :: + Ptr CUChar -> CSize -> CSize -> CUChar -> CInt -> Ptr CSize -> IO CSize + +foreign import capi "process_csv.h emit_delimiter_indices_chunk" + emit_delimiter_indices_chunk :: + Ptr CUChar -> + CSize -> + CSize -> + CUChar -> + CInt -> + CSize -> + Ptr CSize -> + Ptr CSize -> + Ptr CSize -> + IO CSize + +-- | Mirrors @GDI_SIMD_UNAVAILABLE@ in @cbits/process_csv.h@. +simdUnavailable :: CSize +simdUnavailable = maxBound + +{- | Result of the parallel scan: the same @(delimiters, rowEnds)@ pair the +sequential 'DataFrame.IO.CSV.Fast.Index.classifyRowEnds' path produces. +-} +data ParIndex = ParIndex + { piDelims :: !(VS.Vector CSize) + , piRowEnds :: !(VS.Vector Int) + , piEndsInQuote :: !Bool + } + +{- | Scan with @nWorkers@ parallel range scans. Returns 'Nothing' when the +build\/CPU lacks SIMD support (callers fall back to the sequential scan). +-} +getDelimiterIndicesPar :: + Int -> Word8 -> Int -> VS.Vector Word8 -> IO (Maybe ParIndex) +getDelimiterIndicesPar nWorkers separator contentLen file = + VS.unsafeWith file $ \filePtr -> do + probe <- + alloca $ \nlPtr -> + count_delimiters_chunk (castPtr filePtr) 0 0 sep 0 nlPtr + if probe == simdUnavailable + then pure Nothing + else Just <$> scanParallel nWorkers separator contentLen file filePtr + where + sep = fromIntegral separator + +scanParallel :: + Int -> Word8 -> Int -> VS.Vector Word8 -> Ptr Word8 -> IO ParIndex +scanParallel nWorkers separator contentLen file filePtr = do + width <- getNumCapabilities + let sep = fromIntegral separator :: CUChar + simdLen = (contentLen `div` 64) * 64 + tailLen = contentLen - simdLen + ranges = alignedRanges nWorkers simdLen + contentBS = byteStringView contentLen file + cBool b = if b then 1 else 0 :: CInt + -- Quote parity at each range start (prefix-xor of quote counts). + quoteCounts <- + pooledRun + width + [ pure $! BS.count quote (BS.take len (BS.drop off contentBS)) + | (off, len) <- ranges + ] + let parities = scanl Bits.xor False (map odd quoteCounts) + starts = zip ranges parities + -- Pass 1 (parallel): exact per-range delimiter / row-end counts. + counts <- + pooledRun + width + [ alloca $ \nlPtr -> do + d <- + count_delimiters_chunk + (castPtr filePtr) + (fromIntegral off) + (fromIntegral len) + sep + (cBool p) + nlPtr + nl <- peek nlPtr + pure (fromIntegral d :: Int, fromIntegral nl :: Int) + | ((off, len), p) <- starts + ] + let dOffs = scanl (+) 0 (map fst counts) + nlOffs = scanl (+) 0 (map snd counts) + dTot = last dOffs + nlTot = last nlOffs + -- Right-sized flat outputs (+ scalar-tail upper bound + synthetic EOF). + delimsMV <- VSM.unsafeNew (dTot + tailLen + 1) + rowEndsMV <- VSM.unsafeNew (nlTot + tailLen + 1) :: IO (VSM.IOVector Int) + -- Pass 2 (parallel): emit positions + row-end ordinals at exact offsets. + VSM.unsafeWith delimsMV $ \dPtr -> + VSM.unsafeWith rowEndsMV $ \rePtr -> do + _ <- + pooledRun + width + [ alloca $ \rcPtr -> + emit_delimiter_indices_chunk + (castPtr filePtr) + (fromIntegral off) + (fromIntegral len) + sep + (cBool p) + (fromIntegral dOff) + (dPtr `advancePtr` dOff) + (castPtr (rePtr `advancePtr` nlOff)) + rcPtr + | (((off, len), p), dOff, nlOff) <- zip3 starts dOffs nlOffs + ] + pure () + -- Scalar tail (final partial 64-byte chunk), parity carried in. + (found, inQuote) <- + scalarScan separator file delimsMV dTot simdLen contentLen (last parities) + nlAfterTail <- appendTailRowEnds file rowEndsMV delimsMV nlTot dTot found + -- Synthetic EOF delimiter when the file does not end in a newline. + (found', nlFinal) <- + if contentLen > 0 && VS.unsafeIndex file (contentLen - 1) /= lf + then do + VSM.unsafeWrite delimsMV found (fromIntegral contentLen) + VSM.unsafeWrite rowEndsMV nlAfterTail found + pure (found + 1, nlAfterTail + 1) + else pure (found, nlAfterTail) + delims <- VS.unsafeFreeze (VSM.slice 0 found' delimsMV) + rowEnds <- VS.unsafeFreeze (VSM.slice 0 nlFinal rowEndsMV) + pure (ParIndex delims rowEnds inQuote) + +{- | Classify the scalar tail's delimiters: append the ordinal of every +newline entry in @[from, to)@ of the delimiter vector. +-} +appendTailRowEnds :: + VS.Vector Word8 -> + VSM.IOVector Int -> + VSM.IOVector CSize -> + Int -> + Int -> + Int -> + IO Int +appendTailRowEnds file rowEndsMV delimsMV = go + where + go !w !i !to + | i >= to = pure w + | otherwise = do + p <- VSM.unsafeRead delimsMV i + if VS.unsafeIndex file (fromIntegral p) == lf + then VSM.unsafeWrite rowEndsMV w i >> go (w + 1) (i + 1) to + else go w (i + 1) to + +-- | Split @[0, n)@ (a multiple of 64) into 64-byte-aligned ranges. +alignedRanges :: Int -> Int -> [(Int, Int)] +alignedRanges k n = filter ((> 0) . snd) (zip offs sizes) + where + blocks = n `div` 64 + (q, r) = blocks `divMod` max 1 k + sizes = [64 * (q + if i < r then 1 else 0) | i <- [0 .. max 1 k - 1]] + offs = scanl (+) 0 sizes diff --git a/src/DataFrame/IO/CSV/Fast/Parallel.hs b/src/DataFrame/IO/CSV/Fast/Parallel.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/CSV/Fast/Parallel.hs @@ -0,0 +1,197 @@ +{-# LANGUAGE TupleSections #-} + +{- | Chunk-parallel extraction (Round-2 WS-E2): the data-row range is split +into contiguous, newline-aligned chunks (alignment is by construction — +chunks are row ranges over the quote-resolved delimiter index), each +worker runs every column's 'ColumnPlan' over its chunk with right-sized +builders, and the merger splices the results ('mergeColumns' / +'mergeTextChunksPar'). Inference classifies its sample once, globally, +before fan-out; when chunks resolve different types the merger re-parses +only the narrower-typed chunks. + +Parallelism requires the executable to be linked with @-threaded@ and run +with @+RTS -N@; otherwise ('getNumCapabilities' == 1) the sequential path +runs. Inputs under 4 MB are always read sequentially. +-} +module DataFrame.IO.CSV.Fast.Parallel ( + buildAllColumns, + autoChunkCount, +) 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) + +import DataFrame.IO.CSV.Fast.Columns +import DataFrame.IO.CSV.Fast.Passes (NullSpec, PassCol (..)) +import DataFrame.IO.CSV.Fast.TextMerge (mergeTextChunksPar) +import DataFrame.IO.CSV.Fast.Workers (pooledRun) +import DataFrame.Internal.Column (Column, forceColumn) +import DataFrame.Internal.ColumnBuilder (mergeColumns) + +-- | Below this input size the fan-out overhead outweighs the parallelism. +parallelThresholdBytes :: Int +parallelThresholdBytes = 4 * 1024 * 1024 + +{- | Chunk count for an automatic read: four chunks per capability (the +worker pool bounds concurrency at the capability count; the extra chunks +absorb straggler imbalance), or the sequential path (1) on a single +capability or a small input. +-} +autoChunkCount :: Int -> IO Int +autoChunkCount contentLen = do + caps <- getNumCapabilities + pure $ + if caps <= 1 || contentLen < parallelThresholdBytes + then 1 + else caps * 4 + +-- | What one worker produced for one column over its chunk. +data ChunkCol + = CCChain !ChunkOutcome + | -- | @Left@ carries the global data-row index of the first failure. + CCSchema !(Either Int PassCol) + | -- | Legacy (cold-path) columns are built sequentially by the merger. + CCSkip + +-- | Splice homogeneous chunk passes (text chunks merge at the byte level). +mergePassCols :: Int -> [PassCol] -> IO Column +mergePassCols width ps@(PassText{} : _) = + mergeTextChunksPar width [tc | PassText tc <- ps] +mergePassCols _ ps = pure $! mergeColumns [c | PassFull c <- ps] + +{- | Build every column, fanning the row range out over @nChunks@ chunks +on a capability-wide worker pool (sequentially when @nChunks <= 1@). Same +result as the sequential path by construction; see the chunk-determinism +tests and property. +-} +buildAllColumns :: + Int -> ColumnEnv -> V.Vector T.Text -> Int -> IO [(Column, Bool)] +buildAllColumns nChunks env names numCol + | chunks <= 1 = + mapM (\c -> buildColumn env (names V.! c) c) [0 .. numCol - 1] + | otherwise = do + width <- getNumCapabilities + let plans = [planColumn env (names V.! c) c | c <- [0 .. numCol - 1]] + ranges = chunkRanges chunks (ceNumRow env) + traceMarkerIO "fastcsv:fanout" + perChunk <- pooledRun width [runChunk env 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 + -- doubling crosses its threshold mid-merge and a ~0.5s gen-1 + -- collection lands at the worst point of the read. + performMajorGC + -- Merge in parallel too (chunk-level inside each column, pooled + -- across columns), forcing each spliced column here so no memcpy + -- is deferred to the final forceDataFrame walk. + merged <- + pooledRun + width + [ do + r@(col, _) <- + mergeColumn + width + env + (names V.! c) + (plans !! c) + c + [(rng, cols !! c) | (rng, cols) <- zip ranges perChunk] + forceColumn col `seq` pure r + | c <- [0 .. numCol - 1] + ] + traceMarkerIO "fastcsv:merge-done" + pure merged + where + chunks = min nChunks (ceNumRow env) + +-- | Even split of @n@ rows into @k@ @(offset, length)@ ranges. +chunkRanges :: Int -> Int -> [(Int, Int)] +chunkRanges k n = + [ (q * i + min i r, q + (if i < r then 1 else 0)) + | i <- [0 .. k - 1] + ] + where + (q, r) = n `divMod` k + +-- | Restrict an env to a row sub-range (builder hints scaled to match). +chunkEnv :: ColumnEnv -> (Int, Int) -> ColumnEnv +chunkEnv env (off, len) = + env + { ceRows = VS.slice off len (ceRows env) + , ceNumRow = len + , ceTextHint = + max 64 ((ceTextHint env * len) `div` max 1 (ceNumRow env)) + } + +-- | 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 + where + cenv = chunkEnv env range + build col plan = case plan of + PlanLegacy{} -> pure CCSkip + PlanSchema mode nspec t -> + CCSchema <$> runSchemaChunk cenv nspec mode t col off + PlanChain _ nspec steps checkAllNull -> + CCChain <$> runChainChunk cenv nspec steps checkAllNull col + +{- | Splice one column's chunk results back together, resolving chunk-level +type promotion. Mirrors 'buildColumn' for the single-chunk case. +-} +mergeColumn :: + Int -> + ColumnEnv -> + T.Text -> + ColumnPlan -> + Int -> + [((Int, Int), ChunkCol)] -> + IO (Column, Bool) +mergeColumn width env name plan col parts = case plan of + PlanLegacy mode pwt -> (,pwt) <$> legacyColumn env mode col + PlanSchema mode _ t -> + case partitionEithers [r | (_, CCSchema r) <- parts] of + ([], cs) -> do + c <- mergePassCols width cs + pure (finishMode mode c, False) + (failed, _) -> throwIO (schemaError name t (minimum failed)) + PlanChain mode nspec steps _ -> do + c <- mergeChain width env nspec steps col [(r, oc) | (r, CCChain oc) <- parts] + pure (finishMode mode c, False) + +{- | Merge chain-plan chunks. The candidate type starts at the widest chunk +resolution; any chunk that resolved narrower is re-parsed at the candidate +step, and a re-parse failure escalates the candidate further down the +chain (exactly the rows the sequential pass would have failed on). The +final 'StepText' never fails, so this terminates. +-} +mergeChain :: + Int -> + ColumnEnv -> + NullSpec -> + [Step] -> + Int -> + [((Int, Int), ChunkOutcome)] -> + IO Column +mergeChain width env nspec steps col parts + | null resolved = pure (allNullColumn (ceNumRow env)) + | otherwise = settle (maximum resolved) + where + resolved = [i | (_, Resolved i _) <- parts] + settle cand = do + results <- mapM (reparse cand) parts + case sequence results of + Just cs -> mergePassCols width cs + Nothing -> settle (cand + 1) + reparse cand (range, oc) = case oc of + Resolved i c | i == cand -> pure (Just c) + _ -> do + r <- runStep (chunkEnv env range) nspec col False (steps !! cand) + pure (either (const Nothing) Just r) diff --git a/src/DataFrame/IO/CSV/Fast/Passes.hs b/src/DataFrame/IO/CSV/Fast/Passes.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/CSV/Fast/Passes.hs @@ -0,0 +1,274 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE ScopedTypeVariables #-} + +{- | The typed column passes (Round-2 WS-E1): each pass walks one column's +delimiter slices and appends parsed values straight into a WS-C builder — +no per-field 'Data.Text.Text' and no deferred re-parse. Text bytes are +memcpy'd into a shared buffer, only for Text-typed columns. +-} +module DataFrame.IO.CSV.Fast.Passes ( + ColumnEnv (..), + NullSpec (..), + PassCol (..), + PassResult, + passColumn, + nullSpecFor, + isNullSlice, + isNullText, + rowAt, + intPass, + doublePass, + boolPass, + textPass, + datePass, + allNullPass, +) where + +import qualified Data.ByteString as BS +import qualified Data.ByteString.Internal as BSI +import qualified Data.Text as T +import qualified Data.Text.Encoding as TE +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as VM +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Unboxed as VU + +import Control.Monad.ST (stToIO) +import Data.Word (Word8) +import Foreign.Ptr (Ptr, nullPtr, plusPtr) +import GHC.Exts (Int (..), Int#, isTrue#, (+#), (>=#)) + +import DataFrame.IO.CSV (ReadOptions (..)) +import DataFrame.IO.CSV.Fast.Index (quote) +import DataFrame.IO.CSV.Fast.Slice +import DataFrame.Internal.Column (Column, Columnable, fromVector) +import DataFrame.Internal.ColumnBuilder ( + ColumnBuilder (..), + TextChunk, + appendNum, + appendText, + appendTextSliceFromPtr, + freezeTextChunk, + mergeTextChunks, + newNumBuilder, + newTextBuilder, + ) +import DataFrame.Internal.Parsing.Fast ( + isMissingField, + isMissingFieldIn, + isMissingFieldSlice, + parseBoolFieldSlice, + parseDateFieldSlice, + parseDoubleFieldSlice, + parseIntFieldSlice, + ) +import DataFrame.Operations.Typing ( + SafeReadMode (..), + isNullishOrMissing, + parseTimeOpt, + ) + +-- | Per-read context shared by every column pass. +data ColumnEnv = ColumnEnv + { ceCtx :: !FieldCtx + , ceRows :: !(VS.Vector Int) + -- ^ Row indices (into row-ends) of the data rows, blanks skipped. + , ceNumRow :: !Int + , cePtr :: !(Ptr Word8) + -- ^ Base pointer of the file content (valid for the whole read). + , ceTextHint :: !Int + -- ^ Initial byte capacity for a Text column builder. + , ceOpts :: !ReadOptions + } + +-- | Byte-level null test: mode-dependent (audit T3, behaviors list #1). +data NullSpec + = -- | 'NoSafeRead': only the empty field is null. + NullEmpty + | -- | 'MaybeRead': canonical missing set plus custom extras. + NullMissing ![BS.ByteString] ![T.Text] + +nullSpecFor :: ReadOptions -> SafeReadMode -> NullSpec +nullSpecFor _ NoSafeRead = NullEmpty +nullSpecFor opts _ = + NullMissing + (filter (not . isMissingField) (map TE.encodeUtf8 toks)) + toks + where + toks = missingIndicators opts + +{-# INLINE isNullSlice #-} +isNullSlice :: NullSpec -> BS.ByteString -> Int -> Int -> Bool +isNullSlice NullEmpty _ s e = e <= s +isNullSlice (NullMissing extras _) bs s e = + isMissingFieldSlice bs s e + || (not (null extras) && isMissingFieldIn extras (sliceBS bs s e)) + +-- | Text-level twin of 'isNullSlice' for the decode-and-strip trim path. +isNullText :: NullSpec -> T.Text -> Bool +isNullText NullEmpty t = T.null t +isNullText (NullMissing _ toks) t = isNullishOrMissing toks t + +{- | What a pass produced: a finished 'Column', or a raw text chunk that the +merger splices at the byte level (no per-chunk 'T.Text' materialization). +-} +data PassCol = PassFull !Column | PassText !TextChunk + +-- | Finish one 'PassCol' on its own (the single-chunk \/ sequential case). +passColumn :: PassCol -> Column +passColumn (PassFull c) = c +passColumn (PassText tc) = mergeTextChunks [tc] + +-- | Result of a fallible pass: @Left rowIndex@ on the first parse failure. +type PassResult = Either Int PassCol + +rowAt :: ColumnEnv -> Int -> Int +rowAt env = VS.unsafeIndex (ceRows env) +{-# INLINE rowAt #-} + +{- | One pass appending unboxed values; @nullOnFail@ turns parse failures +into nulls (schema + 'MaybeRead') instead of aborting (inference demotion). +-} +{-# INLINE unboxedPass #-} +unboxedPass :: + (Columnable a, VU.Unbox a) => + ColumnEnv -> + NullSpec -> + Int -> + Bool -> + a -> + (BS.ByteString -> Int -> Int -> Maybe a) -> + IO PassResult +unboxedPass env nspec col nullOnFail nullVal parser = do + let ctx = ceCtx env + bs = fcBS ctx + !(I# n) = ceNumRow env + b <- stToIO (newNumBuilder nullVal (ceNumRow env)) + -- Loop counter is Int#: join points are not worker/wrapper'd, so a + -- boxed counter would allocate on every iteration. + let go :: Int# -> IO PassResult + go i + | isTrue# (i >=# n) = + Right . PassFull <$> stToIO (freezeBuilder b) + | otherwise = + withParseSlice ctx (rowAt env (I# i)) col $ \s e -> + if isNullSlice nspec bs s e + then stToIO (appendNull b) >> go (i +# 1#) + else case parser bs s e of + Just v -> stToIO (appendNum b v) >> go (i +# 1#) + Nothing + | nullOnFail -> + stToIO (appendNull b) >> go (i +# 1#) + | otherwise -> pure (Left (I# i)) + go 0# + +intPass :: ColumnEnv -> NullSpec -> Int -> Bool -> IO PassResult +intPass env nspec col nullOnFail = + unboxedPass env nspec col nullOnFail (0 :: Int) parseIntFieldSlice + +doublePass :: ColumnEnv -> NullSpec -> Int -> Bool -> IO PassResult +doublePass env nspec col nullOnFail = + unboxedPass env nspec col nullOnFail (0 :: Double) parseDoubleFieldSlice + +boolPass :: ColumnEnv -> NullSpec -> Int -> Bool -> IO PassResult +boolPass env nspec col nullOnFail = + unboxedPass env nspec col nullOnFail False parseBoolFieldSlice + +{- | Text pass (never fails): raw bytes are memcpy'd into the shared +builder buffer; only quoted fields containing a doubled quote, and the +legacy trim knob, take the decode path. Returns the raw chunk; 'T.Text' +values materialize once, when the merger splices the chunks. +-} +textPass :: ColumnEnv -> NullSpec -> Int -> IO TextChunk +textPass env nspec col = do + let ctx = ceCtx env + bs = fcBS ctx + ptr = cePtr env + !(I# n) = ceNumRow env + b <- stToIO (newTextBuilder (ceNumRow env) (ceTextHint env)) + let go :: Int# -> IO TextChunk + go i + | isTrue# (i >=# n) = stToIO (freezeTextChunk b) + | otherwise = + withFieldSlice ctx (rowAt env (I# i)) col $ \s e q -> + if fcTrim ctx && not q + then do + let t = T.strip (decodeSlice bs s e) + if isNullText nspec t + then stToIO (appendNull b) + else stToIO (appendText b t) + go (i +# 1#) + else + if isNullSlice nspec bs s e + then stToIO (appendNull b) >> go (i +# 1#) + else do + hasQ <- + if q + then containsQuote ptr s e + else pure False + if hasQ + then + stToIO . appendText b $ + unescapeDoubledQuotes + (decodeSlice bs s e) + else + stToIO $ + appendTextSliceFromPtr + b + (ptr `plusPtr` s) + (e - s) + go (i +# 1#) + go 0# + +-- | memchr for a doubled-quote candidate inside a quoted field's content. +containsQuote :: Ptr Word8 -> Int -> Int -> IO Bool +containsQuote ptr s e = do + p <- BSI.memchr (ptr `plusPtr` s) quote (fromIntegral (e - s)) + pure (p /= nullPtr) +{-# INLINE containsQuote #-} + +-- | Date pass: boxed (Day is not unboxable), mirrors @handleDateAssumption@. +datePass :: ColumnEnv -> NullSpec -> Int -> Bool -> IO PassResult +datePass env nspec col nullOnFail = do + let ctx = ceCtx env + bs = fcBS ctx + n = ceNumRow env + dfmt = dateFormat (ceOpts env) + parseAt s e + | dfmt == "%Y-%m-%d" = parseDateFieldSlice bs s e + | otherwise = parseTimeOpt dfmt (decodeSlice bs s e) + out <- VM.new n + let go !i !anyNull + | i >= n = do + vec <- V.unsafeFreeze out + pure . Right . PassFull $ + if anyNull + then fromVector vec + else fromVector (V.mapMaybe id vec) + | otherwise = + withParseSlice ctx (rowAt env i) col $ \s e -> + if isNullSlice nspec bs s e + then VM.unsafeWrite out i Nothing >> go (i + 1) True + else case parseAt s e of + Just d -> + VM.unsafeWrite out i (Just d) >> go (i + 1) anyNull + Nothing + | nullOnFail -> + VM.unsafeWrite out i Nothing >> go (i + 1) True + | otherwise -> pure (Left i) + go (0 :: Int) False + +-- | Is every data cell of the column null? (@handleNoAssumption@ head case.) +allNullPass :: ColumnEnv -> NullSpec -> Int -> IO Bool +allNullPass env nspec col = go 0 + where + ctx = ceCtx env + bs = fcBS ctx + n = ceNumRow env + go !i + | i >= n = pure True + | otherwise = + withParseSlice ctx (rowAt env i) col $ \s e -> + if isNullSlice nspec bs s e then go (i + 1) else pure False diff --git a/src/DataFrame/IO/CSV/Fast/Slice.hs b/src/DataFrame/IO/CSV/Fast/Slice.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/CSV/Fast/Slice.hs @@ -0,0 +1,171 @@ +{-# LANGUAGE BangPatterns #-} + +{- | Field-slice resolution for the fast CSV reader: maps @(row, col)@ to a +byte range of the input (quote-stripped, CR-stripped, optionally trimmed) +without materializing any intermediate value. 'extractField' is the legacy +Text materialization used for headers, inference samples and cold paths. +-} +module DataFrame.IO.CSV.Fast.Slice ( + FieldCtx (..), + withFieldSlice, + withParseSlice, + sliceBS, + decodeSlice, + extractField, + unescapeDoubledQuotes, + stripBom, +) where + +import qualified Data.ByteString as BS +import qualified Data.ByteString.Unsafe as BSU +import qualified Data.Text as Text +import qualified Data.Text.Encoding as TextEncoding +import qualified Data.Vector.Storable as VS + +import Data.Text (Text) +import Data.Word (Word8) +import Foreign.C.Types (CSize) + +import DataFrame.IO.CSV.Fast.Index (cr, quote) + +-- | Everything needed to resolve a field slice, shared by all columns. +data FieldCtx = FieldCtx + { fcFile :: !(VS.Vector Word8) + -- ^ File content (BOM already stripped), no padding. + , fcBS :: !BS.ByteString + -- ^ Zero-copy 'BS.ByteString' view of 'fcFile'. + , fcDelims :: !(VS.Vector CSize) + -- ^ Delimiter byte positions (field and row terminators, flat). + , fcRowEnds :: !(VS.Vector Int) + -- ^ Indices into 'fcDelims' that terminate a row. + , fcContentLen :: !Int + , fcTrim :: !Bool + -- ^ 'DataFrame.IO.CSV.fastCsvTrimUnquoted'. + } + +{- | Resolve field @col@ of row @r@ and continue with @k start end quoted@. +Quoted fields yield the bytes between the outer quotes (embedded @\"\"@ is +NOT unescaped here); a trailing @\\r@ is stripped first; a column beyond the +row's field count (ragged short row) yields the empty slice. CPS so the hot +loops never box the bounds. +-} +{-# INLINE withFieldSlice #-} +withFieldSlice :: FieldCtx -> Int -> Int -> (Int -> Int -> Bool -> r) -> r +withFieldSlice ctx r col k = + let rowEnds = fcRowEnds ctx + endIdx = VS.unsafeIndex rowEnds r + startIdx = if r == 0 then 0 else VS.unsafeIndex rowEnds (r - 1) + 1 + numFields = endIdx - startIdx + 1 + in if col >= numFields + then k 0 0 False + else + let boundaryIdx = startIdx + col + fieldEndRaw = + fromIntegral (VS.unsafeIndex (fcDelims ctx) boundaryIdx) :: Int + fieldEndClamped = min fieldEndRaw (fcContentLen ctx) + fieldStart = + if boundaryIdx == 0 + then 0 + else + fromIntegral + (VS.unsafeIndex (fcDelims ctx) (boundaryIdx - 1)) + + 1 + file = fcFile ctx + fieldEnd = + if fieldEndClamped > fieldStart + && VS.unsafeIndex file (fieldEndClamped - 1) == cr + then fieldEndClamped - 1 + else fieldEndClamped + in if fieldEnd - fieldStart >= 2 + && VS.unsafeIndex file fieldStart == quote + && VS.unsafeIndex file (fieldEnd - 1) == quote + then k (fieldStart + 1) (fieldEnd - 1) True + else k fieldStart fieldEnd False + +{- | 'withFieldSlice' with the trim knob applied: when 'fcTrim' is set, +unquoted slices have ASCII whitespace stripped from both ends before the +continuation runs. This is the slice the typed parsers and null tests see. +-} +{-# INLINE withParseSlice #-} +withParseSlice :: FieldCtx -> Int -> Int -> (Int -> Int -> r) -> r +withParseSlice ctx r col k = + withFieldSlice ctx r col $ \s e q -> + if fcTrim ctx && not q + then + let file = fcFile ctx + s' = skipWsForward file s e + e' = skipWsBackward file s' e + in k s' e' + else k s e + +{-# INLINE isAsciiWs #-} +isAsciiWs :: Word8 -> Bool +isAsciiWs w = w == 0x20 || (w - 0x09) <= 4 + +{-# INLINE skipWsForward #-} +skipWsForward :: VS.Vector Word8 -> Int -> Int -> Int +skipWsForward file = go + where + go !i !e + | i < e && isAsciiWs (VS.unsafeIndex file i) = go (i + 1) e + | otherwise = i + +{-# INLINE skipWsBackward #-} +skipWsBackward :: VS.Vector Word8 -> Int -> Int -> Int +skipWsBackward file = go + where + go !s !e + | e > s && isAsciiWs (VS.unsafeIndex file (e - 1)) = go s (e - 1) + | otherwise = e + +-- | O(1) sub-'BS.ByteString' of the file view (no copy). +{-# INLINE sliceBS #-} +sliceBS :: BS.ByteString -> Int -> Int -> BS.ByteString +sliceBS bs s e = BSU.unsafeTake (e - s) (BSU.unsafeDrop s bs) + +-- | Lenient UTF-8 decode of a slice (allocates; cold paths only). +{-# INLINE decodeSlice #-} +decodeSlice :: BS.ByteString -> Int -> Int -> Text +decodeSlice bs s e = TextEncoding.decodeUtf8Lenient (sliceBS bs s e) + +{- | Extract field @col@ of row @r@ as a 'Text' with the original fastcsv +semantics: quotes stripped and @\"\"@ unescaped for quoted fields, optional +'Text.strip' (full Unicode) for unquoted fields when 'fcTrim' is set. +Used for header names, inference samples and the legacy (EitherRead / +unsupported-schema) column paths. +-} +{-# INLINE extractField #-} +extractField :: FieldCtx -> Int -> Int -> Text +extractField ctx r col = + withFieldSlice ctx r col $ \s e q -> + if q + then unescapeDoubledQuotes (decodeSlice (fcBS ctx) s e) + else + (if fcTrim ctx then Text.strip else id) + (decodeSlice (fcBS ctx) s e) + +{- | RFC 4180 inner-quote unescape: @\"\"@ → @\"@. 'Text.replace' on a +two-char needle does a single linear pass and is allocation-free when +no doubled quote is present (short-circuits at the first miss). +-} +{-# INLINE unescapeDoubledQuotes #-} +unescapeDoubledQuotes :: Text -> Text +unescapeDoubledQuotes t + | Text.isInfixOf doubledQuote t = Text.replace doubledQuote singleQuote t + | otherwise = t + where + doubledQuote = Text.pack "\"\"" + singleQuote = Text.singleton '"' + +{- | Strip a leading UTF-8 BOM (EF BB BF) if present. Returns the trimmed +vector and the number of bytes removed (0 or 3). +-} +{-# INLINE stripBom #-} +stripBom :: VS.Vector Word8 -> (VS.Vector Word8, Int) +stripBom v + | VS.length v >= 3 + , VS.unsafeIndex v 0 == 0xEF + , VS.unsafeIndex v 1 == 0xBB + , VS.unsafeIndex v 2 == 0xBF = + (VS.drop 3 v, 3) + | otherwise = (v, 0) diff --git a/src/DataFrame/IO/CSV/Fast/TextMerge.hs b/src/DataFrame/IO/CSV/Fast/TextMerge.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/CSV/Fast/TextMerge.hs @@ -0,0 +1,64 @@ +{-# LANGUAGE BangPatterns #-} + +{- | Parallel byte-level merge of per-chunk text columns: the byte copies +and offset rebase fan out over the chunk list, so a large text column does +not serialize the merge phase. Wraps the merged shared buffer + offsets as +'PackedText' (no 'Data.Text.Text' spine, no eager UTF-8 validation), exactly +as the pure 'mergeTextChunks' now produces. +-} +module DataFrame.IO.CSV.Fast.TextMerge ( + mergeTextChunksPar, +) where + +import qualified Data.Text.Array as A +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM + +import Control.Monad (void) +import Control.Monad.ST (stToIO) + +import DataFrame.IO.CSV.Fast.Workers (pooledRun) +import DataFrame.Internal.Column (Column (..)) +import DataFrame.Internal.ColumnMerge ( + TextChunk (..), + mergeTextChunks, + spliceBitmaps, + tcRows, + ) +import DataFrame.Internal.PackedText (mkPackedContiguous) + +{- | Merge text chunks with @width@-way parallel byte copies + offset +rebase, then wrap the shared buffer as 'PackedText'. Single chunks take +the pure (zero-copy) path. +-} +mergeTextChunksPar :: Int -> [TextChunk] -> IO Column +mergeTextChunksPar _ [c] = pure $! mergeTextChunks [c] +mergeTextChunksPar width cs = do + let byteOffs = scanl (+) 0 (map tcUsed cs) + rowOffs = scanl (+) 0 (map tcRows cs) + totalBytes = last byteOffs + totalRows = last rowOffs + marr <- stToIO (A.new (max 1 totalBytes)) + offsMV <- VUM.unsafeNew (totalRows + 1) + VUM.unsafeWrite offsMV 0 0 + -- Byte copy + offset rebase, parallel over chunks (disjoint ranges). + void . pooledRun width $ + [ do + stToIO (A.copyI (tcUsed c) marr bOff (tcBytes c) 0) + let co = tcOffsets c + n = tcRows c + fill !i + | i > n = pure () + | otherwise = do + VUM.unsafeWrite + offsMV + (rOff + i) + (bOff + VU.unsafeIndex co i) + fill (i + 1) + fill 1 + | (c, bOff, rOff) <- zip3 cs byteOffs rowOffs + ] + arr <- stToIO (A.unsafeFreeze marr) + offs <- VU.unsafeFreeze offsMV + let !bm = spliceBitmaps [(tcBitmap c, tcRows c) | c <- cs] + pure (PackedText bm (mkPackedContiguous arr offs)) diff --git a/src/DataFrame/IO/CSV/Fast/Workers.hs b/src/DataFrame/IO/CSV/Fast/Workers.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/CSV/Fast/Workers.hs @@ -0,0 +1,51 @@ +{- | Thread fan-out primitives shared by the parallel scan, the chunk +parse and the parallel merge: plain 'forkIO' workers joined through +'MVar's (no sparks), with a counter-based pool for finer-grained chunks. +-} +module DataFrame.IO.CSV.Fast.Workers ( + forkJoin, + pooledRun, +) where + +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as VM + +import Control.Concurrent (forkIO) +import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar) +import Control.Exception (SomeException, throwIO, try) +import Control.Monad (when) +import Data.IORef (atomicModifyIORef', newIORef) + +-- | Run each action in its own thread; rethrow the first failure in order. +forkJoin :: [IO a] -> IO [a] +forkJoin actions = do + vars <- mapM spawn actions + results <- mapM takeMVar vars + either (throwIO :: SomeException -> IO [a]) pure (sequence results) + where + spawn act = do + var <- newEmptyMVar + _ <- forkIO (try act >>= putMVar var) + pure var + +{- | Run the actions on a pool of @width@ threads (work-stealing via a +shared counter), so finer-grained chunks balance load without running +every chunk's builders concurrently. Results keep their input order. +-} +pooledRun :: Int -> [IO a] -> IO [a] +pooledRun width actions + | width >= n = forkJoin actions + | otherwise = do + next <- newIORef 0 + out <- VM.unsafeNew n + let acts = V.fromListN n actions + worker = do + i <- atomicModifyIORef' next (\j -> (j + 1, j)) + when (i < n) $ do + r <- acts V.! i + VM.write out i r + worker + _ <- forkJoin (replicate width worker) + V.toList <$> V.freeze out + where + n = length actions diff --git a/tests/Main.hs b/tests/Main.hs --- a/tests/Main.hs +++ b/tests/Main.hs @@ -7,11 +7,18 @@ import Test.HUnit import Test.QuickCheck +import qualified Operations.ChunkParallel import qualified Operations.ReadCsv +import qualified Operations.TypedExtraction import qualified Properties.Csv tests :: Test -tests = TestList Operations.ReadCsv.tests +tests = + TestList + ( Operations.ReadCsv.tests + <> Operations.TypedExtraction.tests + <> Operations.ChunkParallel.tests + ) isSuccessful :: Result -> Bool isSuccessful (Success{}) = True diff --git a/tests/Operations/ChunkParallel.hs b/tests/Operations/ChunkParallel.hs new file mode 100644 --- /dev/null +++ b/tests/Operations/ChunkParallel.hs @@ -0,0 +1,277 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +{- | Round-2 (WS-E2) behavior pins for chunk-parallel extraction: a chunked +read must be indistinguishable from the sequential one — same column types +(promotion resolved across chunks), same values, same errors. +-} +module Operations.ChunkParallel (tests) where + +import qualified Data.Map as M +import qualified Data.Proxy as P +import qualified Data.Text as T +import qualified Data.Text.Encoding as TE + +import Control.Exception (ErrorCall, try) +import Data.Time (Day) +import Data.Word (Word8) +import DataFrame.IO.CSV ( + ReadOptions (..), + TypeSpec (..), + defaultReadOptions, + ) +import qualified DataFrame.IO.CSV.Fast as D +import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.DataFrame (DataFrame, getColumn) +import DataFrame.Internal.Schema (Schema (..), SchemaType (..), elements) +import DataFrame.Operations.Typing (SafeReadMode (..)) +import Test.HUnit + +comma :: Word8 +comma = 0x2C + +-- | Read @body@ with a forced chunk count (bypasses the size/caps gate). +readChunked :: Int -> ReadOptions -> T.Text -> IO DataFrame +readChunked n opts body = + D.readSeparatedFromBytesChunks n comma opts (TE.encodeUtf8 body) + +readSequential :: ReadOptions -> T.Text -> IO DataFrame +readSequential opts body = + D.readSeparatedFromBytes comma opts (TE.encodeUtf8 body) + +-- | Assert chunked == sequential for several chunk counts. +assertChunkInvariant :: String -> ReadOptions -> T.Text -> Assertion +assertChunkInvariant what opts body = do + seqDf <- readSequential opts body + mapM_ + ( \n -> do + parDf <- readChunked n opts body + assertEqual (what <> " (chunks=" <> show n <> ")") seqDf parDf + ) + [2, 3, 5, 8] + +expectColumn :: String -> T.Text -> DataFrame -> DI.Column -> Assertion +expectColumn what name df expected = case getColumn name df of + Nothing -> assertFailure (what <> ": column missing") + Just col -> assertEqual what expected col + +-- A 200-row mixed dataset: nullable ints, doubles, quoted text with +-- embedded separators/newlines/doubled quotes, bools, nullable dates. +mixedBody :: T.Text +mixedBody = "i,d,t,b,dt\n" <> T.concat (map row [0 .. 199 :: Int]) + where + row k = + T.intercalate + "," + [ if k `mod` 7 == 3 then "" else T.pack (show k) + , T.pack (show (fromIntegral k * 0.25 :: Double)) + , textCell k + , if even k then "True" else "false" + , if k `mod` 11 == 5 + then "" + else "2024-01-" <> T.justifyRight 2 '0' (T.pack (show (1 + k `mod` 28))) + ] + <> "\n" + textCell k = case k `mod` 5 of + 0 -> "\"a,b" <> T.pack (show k) <> "\"" + 1 -> "\"line\nbreak" <> T.pack (show k) <> "\"" + 2 -> "\"quote\"\"inside" <> T.pack (show k) <> "\"" + 3 -> "" + _ -> "plain" <> T.pack (show k) + +testParallelMatchesSequentialMixed :: Test +testParallelMatchesSequentialMixed = + TestLabel "par_matches_seq_mixed" . TestCase $ + assertChunkInvariant "mixed nulls+quoted" defaultReadOptions mixedBody + +-- A Double cell far past the sample, in the last chunk: every chunk that +-- resolved Int must be re-parsed as Double by the merger. +testPromotionAcrossChunks :: Test +testPromotionAcrossChunks = TestLabel "par_promotion_int_double" . TestCase $ do + let body = + "a\n" + <> T.concat [T.pack (show k) <> "\n" | k <- [0 .. 189 :: Int]] + <> "190.5\n" + assertChunkInvariant "late double promotes ints" defaultReadOptions body + df <- readChunked 4 defaultReadOptions body + expectColumn + "merged column is Double" + "a" + df + (DI.fromList @Double ([fromIntegral k | k <- [0 .. 189 :: Int]] <> [190.5])) + +-- A Bool cell in the last chunk of an Int column: candidate escalates past +-- Double (the bool cell fails it) all the way to Text. +testDemotionToTextAcrossChunks :: Test +testDemotionToTextAcrossChunks = TestLabel "par_demotion_text" . TestCase $ do + let body = + "a\n" + <> T.concat [T.pack (show k) <> "\n" | k <- [0 .. 189 :: Int]] + <> "True\n" + assertChunkInvariant "late bool demotes ints to Text" defaultReadOptions body + df <- readChunked 4 defaultReadOptions body + expectColumn + "merged column is Text" + "a" + df + (DI.fromList @T.Text ([T.pack (show k) | k <- [0 .. 189 :: Int]] <> ["True"])) + +-- An all-null column must merge to the same all-null Maybe Text column. +testAllNullColumnAcrossChunks :: Test +testAllNullColumnAcrossChunks = TestLabel "par_all_null" . TestCase $ do + let body = + "a,b\n" <> T.concat [",x" <> T.pack (show k) <> "\n" | k <- [0 .. 199 :: Int]] + assertChunkInvariant "all-null column" defaultReadOptions body + df <- readChunked 4 defaultReadOptions body + expectColumn + "all-null column is Maybe Text" + "a" + df + (DI.fromList @(Maybe T.Text) (replicate 200 Nothing)) + +-- All-null sample (NoAssumption), data only in the last chunk: all-null +-- chunks must be re-parsed at the type the data chunk resolved. +testLateDataAfterAllNullSample :: Test +testLateDataAfterAllNullSample = TestLabel "par_late_bool" . TestCase $ do + let cell k + | k < 150 = "" + | even k = "True" + | otherwise = "false" + body = + "a,b\n" + <> T.concat [cell k <> ",x\n" | k <- [0 .. 199 :: Int]] + assertChunkInvariant "late bools after null sample" defaultReadOptions body + df <- readChunked 4 defaultReadOptions body + expectColumn + "column is Maybe Bool" + "a" + df + ( DI.fromList @(Maybe Bool) + [ if k < 150 then Nothing else Just (even k) + | k <- [0 .. 199 :: Int] + ] + ) + +-- Schema + NoSafeRead failure in a late chunk: same error, same row index +-- as the sequential read. +testSchemaErrorSameRowAsSequential :: Test +testSchemaErrorSameRowAsSequential = TestLabel "par_schema_error" . TestCase $ do + let body = + "a\n" + <> T.concat [T.pack (show k) <> "\n" | k <- [0 .. 179 :: Int]] + <> "boom\n" + <> T.concat [T.pack (show k) <> "\n" | k <- [181 .. 199 :: Int]] + opts = + defaultReadOptions + { typeSpec = + SpecifyTypes + (M.toList (elements (Schema (M.fromList [("a", SType (P.Proxy @Int))])))) + (typeSpec defaultReadOptions) + } + seqErr <- try @ErrorCall (readSequential opts body) + parErr <- try @ErrorCall (readChunked 4 opts body) + case (seqErr, parErr) of + (Left e1, Left e2) -> assertEqual "same schema error" (show e1) (show e2) + _ -> assertFailure "both reads should raise on the unparseable schema cell" + +-- MaybeRead missing tokens with chunked extraction. +testMaybeReadAcrossChunks :: Test +testMaybeReadAcrossChunks = TestLabel "par_mayberead" . TestCase $ do + let cell k = if k `mod` 13 == 7 then "NA" else T.pack (show k) + body = "a\n" <> T.concat [cell k <> "\n" | k <- [0 .. 199 :: Int]] + opts = defaultReadOptions{safeRead = MaybeRead} + assertChunkInvariant "MaybeRead missing tokens" opts body + df <- readChunked 4 opts body + expectColumn + "column is Maybe Int" + "a" + df + ( DI.fromList @(Maybe Int) + [ if k `mod` 13 == 7 then Nothing else Just k + | k <- [0 .. 199 :: Int] + ] + ) + +-- Ragged rows (PadWithNull default) must behave identically chunked. +testRaggedRowsAcrossChunks :: Test +testRaggedRowsAcrossChunks = TestLabel "par_ragged" . TestCase $ do + let row k + | k `mod` 9 == 4 = T.pack (show k) <> "\n" + | k `mod` 9 == 8 = T.pack (show k) <> ",x,extra\n" + | otherwise = T.pack (show k) <> ",x" <> T.pack (show k) <> "\n" + body = "a,b\n" <> T.concat (map row [0 .. 199 :: Int]) + assertChunkInvariant "ragged rows pad with null" defaultReadOptions body + +-- Date column with nulls: boxed chunk columns must splice correctly. +testDateNullsAcrossChunks :: Test +testDateNullsAcrossChunks = TestLabel "par_date_nulls" . TestCase $ do + -- A second column keeps null rows from looking like blank lines + -- (blank lines are skipped by documented behavior). + let cell k = + if k `mod` 17 == 9 + then "" + else "2023-05-" <> T.justifyRight 2 '0' (T.pack (show (1 + k `mod` 28))) + body = "a,b\n" <> T.concat [cell k <> ",x\n" | k <- [0 .. 199 :: Int]] + assertChunkInvariant "nullable dates" defaultReadOptions body + df <- readChunked 4 defaultReadOptions body + expectColumn + "column is Maybe Day" + "a" + df + ( DI.fromList @(Maybe Day) + [ if k `mod` 17 == 9 + then Nothing + else Just (read ("2023-05-" <> pad (1 + k `mod` 28))) + | k <- [0 .. 199 :: Int] + ] + ) + where + pad n = if n < 10 then '0' : show n else show n + +-- An unclosed quote must still raise under the chunked (parallel) scan. +testUnclosedQuoteAcrossChunks :: Test +testUnclosedQuoteAcrossChunks = TestLabel "par_unclosed_quote" . TestCase $ do + let body = + "a\n" + <> T.concat [T.pack (show k) <> "\n" | k <- [0 .. 199 :: Int]] + <> "\"dangling" + r <- try @D.CsvParseError (readChunked 4 defaultReadOptions body) + case r of + Left D.CsvUnclosedQuote -> pure () + other -> assertFailure ("expected CsvUnclosedQuote, got " <> show other) + +-- WS-E2 ingest pin: the parallel byte-level merge wraps the merged shared +-- buffer as 'PackedText' (no Text spine), so the parallel path is also an +-- RSS win. Values stay identical to the sequential read. +testParallelTextFreezesPacked :: Test +testParallelTextFreezesPacked = TestLabel "par_text_freezes_packed" . TestCase $ do + let body = + "a,b\n" + <> T.concat + [ "row" <> T.pack (show k) <> ",t" <> T.pack (show k) <> "\n" + | k <- [0 .. 199 :: Int] + ] + df <- readChunked 4 defaultReadOptions body + case getColumn "b" df of + Nothing -> assertFailure "missing text column" + Just col -> do + assertBool "parallel text column is PackedText" (DI.isPackedText col) + assertEqual + "parallel text values byte-identical" + (DI.fromList @T.Text ["t" <> T.pack (show k) | k <- [0 .. 199 :: Int]]) + col + +tests :: [Test] +tests = + [ testParallelMatchesSequentialMixed + , testPromotionAcrossChunks + , testDemotionToTextAcrossChunks + , testAllNullColumnAcrossChunks + , testLateDataAfterAllNullSample + , testSchemaErrorSameRowAsSequential + , testMaybeReadAcrossChunks + , testRaggedRowsAcrossChunks + , testDateNullsAcrossChunks + , testUnclosedQuoteAcrossChunks + , testParallelTextFreezesPacked + ] diff --git a/tests/Operations/ReadCsv.hs b/tests/Operations/ReadCsv.hs --- a/tests/Operations/ReadCsv.hs +++ b/tests/Operations/ReadCsv.hs @@ -85,6 +85,7 @@ getRowEscaped sep df i = V.ifoldr go [] (columns df) where go :: Int -> Column -> [T.Text] -> [T.Text] + go idx c@(PackedText _ _) acc = go idx (DI.materializePacked c) acc go _ (BoxedColumn bm (c :: V.Vector a)) acc = case c V.!? i of Just e -> escapeField sep textRep : acc where diff --git a/tests/Operations/TypedExtraction.hs b/tests/Operations/TypedExtraction.hs new file mode 100644 --- /dev/null +++ b/tests/Operations/TypedExtraction.hs @@ -0,0 +1,322 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Round-2 (WS-E1) behavior pins for the typed-extraction fastcsv path: +schema bypass, unified WS-B parser semantics, byte-level missing tests, +and the chunk-boundary scalar tail that replaced the full-file copy. +-} +module Operations.TypedExtraction (tests) where + +import qualified Data.Map as M +import qualified Data.Proxy as P +import qualified Data.Text as T +import qualified Data.Text.Encoding as TE +import qualified Data.Text.IO as TIO + +import Control.Exception (ErrorCall, evaluate, try) +import Data.Time (Day) +import DataFrame.IO.CSV ( + ReadOptions (..), + TypeSpec (..), + defaultReadOptions, + ) +import qualified DataFrame.IO.CSV.Fast as D +import DataFrame.Internal.Column (Column) +import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.DataFrame (DataFrame, dataframeDimensions, getColumn) +import DataFrame.Internal.Schema (Schema (..), SchemaType (..), elements) +import DataFrame.Operations.Typing (SafeReadMode (..)) +import System.Directory (removeFile) +import Test.HUnit + +tmpDir :: FilePath +tmpDir = "./tests/data/unstable_csv/" + +withCsv :: String -> T.Text -> (FilePath -> IO a) -> IO a +withCsv name body action = do + let path = tmpDir <> "typed_" <> name <> ".csv" + TIO.writeFile path body + r <- action path + removeFile path + pure r + +expectColumn :: String -> T.Text -> DataFrame -> Column -> Assertion +expectColumn what name df expected = case getColumn name df of + Nothing -> assertFailure (what <> ": column missing") + Just col -> assertEqual what expected col + +-- Divergence #1 (pinned new behavior): Int overflow rejects instead of +-- wrapping, so an overflowing cell demotes the column to Double. +testIntOverflowDemotesToDouble :: Test +testIntOverflowDemotesToDouble = TestLabel "typed_int_overflow_demotes" $ + TestCase $ + withCsv "overflow" "a\n1\n9223372036854775808\n" $ \path -> do + df <- D.fastReadCsv path + expectColumn + "overflowing int column becomes Double" + "a" + df + (DI.fromList @Double [1.0, 9.223372036854776e18]) + +-- Divergence #2 (pinned new behavior): doubles parse with strip-equivalent +-- semantics, so a padded double past the 100-row inference sample no longer +-- demotes the whole column to Text. (Inside the sample the Text-level +-- classifier still rejects padding, by design.) +testPaddedDoubleStaysDouble :: Test +testPaddedDoubleStaysDouble = TestLabel "typed_padded_double" $ + TestCase $ do + let clean = [fromIntegral i + 0.5 :: Double | i <- [1 .. 120 :: Int]] + cleanRows = T.unlines (map (T.pack . show) clean) + withCsv "padded_double" ("a\n" <> cleanRows <> " 3.5\n") $ \path -> do + df <- D.fastReadCsv path + expectColumn + "padded double cell past the sample stays Double" + "a" + df + (DI.fromList @Double (clean <> [3.5])) + +-- Schema bypasses inference: a numeric-looking column declared Text in the +-- schema must come back as Text (old code silently kept the inferred Int). +testSchemaTextBypassesInference :: Test +testSchemaTextBypassesInference = TestLabel "typed_schema_text" $ + TestCase $ + withCsv "schema_text" "a\n1\n2\n" $ \path -> do + let schema = Schema (M.fromList [("a", SType (P.Proxy @T.Text))]) + df <- D.fastReadCsvWithSchema schema path + expectColumn + "schema Text wins over inferred Int" + "a" + df + (DI.fromList @T.Text ["1", "2"]) + +-- Schema + NoSafeRead: a cell that cannot parse as the declared type +-- raises during the read (strict read), instead of hiding an error thunk. +testSchemaStrictFailureThrows :: Test +testSchemaStrictFailureThrows = TestLabel "typed_schema_strict_throws" $ + TestCase $ + withCsv "schema_strict" "a\n1\nx\n" $ \path -> do + let schema = Schema (M.fromList [("a", SType (P.Proxy @Int))]) + result <- try @ErrorCall $ do + df <- D.fastReadCsvWithSchema schema path + _ <- evaluate (dataframeDimensions df) + case getColumn "a" df of + Just col -> evaluate (show col) + Nothing -> pure "" + case result of + Left _ -> pure () + Right _ -> + assertFailure + "schema Int over unparseable cell should raise under NoSafeRead" + +-- Schema + MaybeRead: unparseable cells become nulls. +testSchemaMaybeReadNullsFailures :: Test +testSchemaMaybeReadNullsFailures = TestLabel "typed_schema_mayberead" $ + TestCase $ + withCsv "schema_maybe" "a\n1\nx\n2\n" $ \path -> do + let schema = Schema (M.fromList [("a", SType (P.Proxy @Int))]) + opts = + defaultReadOptions + { typeSpec = + SpecifyTypes + (M.toList (elements schema)) + (typeSpec defaultReadOptions) + , safeRead = MaybeRead + } + df <- D.fastReadCsvWithOpts opts path + expectColumn + "schema Int + MaybeRead nulls the bad cell" + "a" + df + (DI.fromList @(Maybe Int) [Just 1, Nothing, Just 2]) + +-- Preserved: NoSafeRead treats only empty cells as null, so "NA" still +-- demotes a numeric column to Text. +testNoSafeReadNADemotesToText :: Test +testNoSafeReadNADemotesToText = TestLabel "typed_nosaferead_na" $ + TestCase $ + withCsv "na_demotes" "a\n1\nNA\n" $ \path -> do + df <- D.fastReadCsv path + expectColumn + "NA is data under NoSafeRead" + "a" + df + (DI.fromList @T.Text ["1", "NA"]) + +-- Preserved: MaybeRead honours the canonical missing-token set. +testMaybeReadMissingTokens :: Test +testMaybeReadMissingTokens = TestLabel "typed_mayberead_missing" $ + TestCase $ + withCsv "maybe_missing" "a\n1\nNA\n2\n" $ \path -> do + let opts = defaultReadOptions{safeRead = MaybeRead} + df <- D.fastReadCsvWithOpts opts path + expectColumn + "NA is null under MaybeRead" + "a" + df + (DI.fromList @(Maybe Int) [Just 1, Nothing, Just 2]) + +-- Preserved: custom missing indicators extend the canonical set. +testMaybeReadCustomMissingToken :: Test +testMaybeReadCustomMissingToken = TestLabel "typed_mayberead_custom" $ + TestCase $ + withCsv "maybe_custom" "a\n1\nMISSING\n2\n" $ \path -> do + let opts = + defaultReadOptions + { safeRead = MaybeRead + , missingIndicators = + "MISSING" : missingIndicators defaultReadOptions + } + df <- D.fastReadCsvWithOpts opts path + expectColumn + "custom token is null under MaybeRead" + "a" + df + (DI.fromList @(Maybe Int) [Just 1, Nothing, Just 2]) + +-- Preserved: quoted numerics parse as numbers (quote stripping happens +-- before the typed parse). +testQuotedNumbersParse :: Test +testQuotedNumbersParse = TestLabel "typed_quoted_numbers" $ + TestCase $ + withCsv "quoted_nums" "a,b\n\"42\",\"1.5\"\n\"7\",\"2.5\"\n" $ \path -> do + df <- D.fastReadCsv path + expectColumn "quoted ints" "a" df (DI.fromList @Int [42, 7]) + expectColumn "quoted doubles" "b" df (DI.fromList @Double [1.5, 2.5]) + +-- Preserved: Bool and Date inference still land on typed columns. +testBoolAndDateInference :: Test +testBoolAndDateInference = TestLabel "typed_bool_date" $ + TestCase $ + withCsv "bool_date" "b,d\nTrue,2020-01-01\nFalse,2021-12-31\n" $ \path -> do + df <- D.fastReadCsv path + expectColumn "bool column" "b" df (DI.fromList @Bool [True, False]) + expectColumn + "date column" + "d" + df + (DI.fromList [read @Day "2020-01-01", read "2021-12-31"]) + +-- Preserved: a null inside a date column yields Maybe Day. +testDateWithNulls :: Test +testDateWithNulls = TestLabel "typed_date_nulls" $ + TestCase $ + withCsv "date_nulls" "d,x\n2020-01-01,a\n,b\n2021-12-31,c\n" $ \path -> do + df <- D.fastReadCsv path + expectColumn + "nullable date column" + "d" + df + ( DI.fromList + [ Just (read @Day "2020-01-01") + , Nothing + , Just (read "2021-12-31") + ] + ) + +-- The scalar tail (last <64 bytes after the SIMD chunks) must honour the +-- quote state carried out of the SIMD region: a quoted field that opens +-- before the 64-byte boundary and closes inside the tail may contain +-- separators and newlines. +testQuoteSpansChunkBoundary :: Test +testQuoteSpansChunkBoundary = TestLabel "typed_quote_spans_boundary" $ + TestCase $ do + let pad = T.replicate 70 "x" + body = "a,b\n" <> pad <> ",\"with,comma\nand newline\"\n" + withCsv "boundary_quote" body $ \path -> do + df <- D.fastReadCsv path + assertEqual "one data row" 1 (fst (dataframeDimensions df)) + expectColumn + "quoted field spanning the SIMD/tail boundary" + "b" + df + (DI.fromList @T.Text ["with,comma\nand newline"]) + +-- An unclosed quote that opens inside the scalar tail must still raise. +testUnclosedQuoteInTail :: Test +testUnclosedQuoteInTail = TestLabel "typed_unclosed_in_tail" $ + TestCase $ do + let pad = T.replicate 70 "x" + body = "a\n" <> pad <> "\n\"dangling" + withCsv "tail_unclosed" body $ \path -> do + result <- try @D.CsvParseError (D.fastReadCsv path) + case result of + Left D.CsvUnclosedQuote -> pure () + other -> + assertFailure ("expected CsvUnclosedQuote, got " <> show other) + +-- S2 pin: inference must land on exactly the same columns (types and +-- values) as an explicit schema declaring the true types. +testInferredMatchesSchema :: Test +testInferredMatchesSchema = TestLabel "typed_inferred_matches_schema" $ + TestCase $ do + let rows = + T.unlines + [ T.intercalate + "," + [ T.pack (show i) + , T.pack (show (fromIntegral i + 0.25 :: Double)) + , "name" <> T.pack (show i) + , if even i then "True" else "false" + , "2024-01-0" <> T.pack (show (1 + i `mod` 9)) + ] + | i <- [1 .. 150 :: Int] + ] + body = "i,d,t,b,dt\n" <> rows + schema = + Schema + ( M.fromList + [ ("i", SType (P.Proxy @Int)) + , ("d", SType (P.Proxy @Double)) + , ("t", SType (P.Proxy @T.Text)) + , ("b", SType (P.Proxy @Bool)) + , ("dt", SType (P.Proxy @Day)) + ] + ) + withCsv "inferred_vs_schema" body $ \path -> do + inferred <- D.fastReadCsv path + explicit <- D.fastReadCsvWithSchema schema path + assertEqual "inferred == explicit schema" explicit inferred + +-- Ingest pin: an inferred string column freezes to a shared-buffer +-- 'PackedText' (no per-row boxed Text header), values byte-identical. +testTextFreezesPacked :: Test +testTextFreezesPacked = TestLabel "typed_text_freezes_packed" $ + TestCase $ + withCsv "freezes_packed" "a,b\n1,x\n2,y\n" $ \path -> do + df <- D.fastReadCsv path + case getColumn "b" df of + Nothing -> assertFailure "missing text column" + Just col -> do + assertBool "text column is PackedText" (DI.isPackedText col) + assertEqual "values byte-identical" (DI.fromList @T.Text ["x", "y"]) col + +-- The in-memory entry point must agree with the file-based one. +testFromBytesMatchesFile :: Test +testFromBytesMatchesFile = TestLabel "typed_from_bytes" $ + TestCase $ do + let body = "a,b,c\n1,x,2.5\n2,y,3.5\n,z,\n" + withCsv "from_bytes" body $ \path -> do + fromFile <- D.fastReadCsv path + fromBytes <- D.fastReadCsvFromBytes (TE.encodeUtf8 body) + assertEqual "file and bytes parses agree" fromFile fromBytes + +tests :: [Test] +tests = + [ testIntOverflowDemotesToDouble + , testPaddedDoubleStaysDouble + , testSchemaTextBypassesInference + , testSchemaStrictFailureThrows + , testSchemaMaybeReadNullsFailures + , testNoSafeReadNADemotesToText + , testMaybeReadMissingTokens + , testMaybeReadCustomMissingToken + , testQuotedNumbersParse + , testBoolAndDateInference + , testDateWithNulls + , testQuoteSpansChunkBoundary + , testUnclosedQuoteInTail + , testInferredMatchesSchema + , testTextFreezesPacked + , testFromBytesMatchesFile + ] diff --git a/tests/Properties/Csv.hs b/tests/Properties/Csv.hs --- a/tests/Properties/Csv.hs +++ b/tests/Properties/Csv.hs @@ -19,13 +19,16 @@ import qualified Data.List as L import qualified Data.Map 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.Vector as V +import DataFrame.IO.CSV (defaultReadOptions) + import Data.Type.Equality (testEquality, (:~:) (Refl)) import DataFrame.IO.CSV.Fast (CsvParseError (..)) import qualified DataFrame.IO.CSV.Fast as D -import DataFrame.Internal.Column (Column (..)) +import DataFrame.Internal.Column (Column (..), materializePacked) import DataFrame.Internal.DataFrame ( DataFrame (..), columnIndices, @@ -114,12 +117,13 @@ columnAsText name df = do idx <- M.lookup name (columnIndices df) col <- columns df V.!? idx - case col of + case materializePacked col of BoxedColumn _ (vec :: V.Vector a) -> case testEquality (typeRep @a) (typeRep @T.Text) of Just Refl -> Just (V.toList vec) Nothing -> Nothing UnboxedColumn{} -> Nothing + PackedText{} -> Nothing {- | Run a property in @IO@ against a generated CSV text, cleaning up the temp file afterwards no matter what. @@ -238,6 +242,31 @@ assert (cs == 1) assert (rs >= 0) +{- | WS-E2 determinism: a chunk-parallel read returns exactly the same +DataFrame as the sequential read, for any chunk count. The generated grid +mixes Text cells (quotes, embedded separators and newlines) with a +nullable numeric column so chunk-level type promotion is exercised too. +-} +prop_parallel_equals_sequential :: ColumnNames -> [[Cell]] -> Property +prop_parallel_equals_sequential (ColumnNames header) rows = + forAll (chooseInt (2, 7)) $ \nChunks -> + forAll (vectorOf (length rows) numCell) $ \numCol -> monadicIO $ do + let cellRows = + zipWith (\r n -> map unCell r <> [n]) rows numCol + csv = encodeCsv ',' (header <> ["num"]) cellRows + bytes = TE.encodeUtf8 csv + seqDf <- run (D.readSeparatedFromBytes 0x2C defaultReadOptions bytes) + parDf <- + run (D.readSeparatedFromBytesChunks nChunks 0x2C defaultReadOptions bytes) + assert (seqDf == parDf) + where + numCell = + frequency + [ (4, T.pack . show <$> chooseInt (-1000, 1000)) + , (2, T.pack . show <$> (choose (-10, 10) :: Gen Double)) + , (1, pure "") + ] + tests :: [Property] tests = [ property prop_bom_invariant @@ -245,4 +274,5 @@ , property prop_roundtrip_ascii , property prop_unclosed_quote_throws , property prop_boundary_sizes + , property prop_parallel_equals_sequential ]