diff --git a/dataframe-fastcsv.cabal b/dataframe-fastcsv.cabal --- a/dataframe-fastcsv.cabal +++ b/dataframe-fastcsv.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: dataframe-fastcsv -version: 1.4.0.1 +version: 1.4.1.0 synopsis: SIMD-accelerated CSV reader for the dataframe library. description: A fast, SIMD-accelerated CSV/TSV reader using memory-mapped I/O @@ -57,9 +57,9 @@ build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.14, containers >= 0.6.7 && < 0.10, - dataframe-core >= 2.2 && < 2.3, + dataframe-core >= 2.4 && < 2.5, dataframe-csv >= 2.3 && < 2.4, - dataframe-operations >= 2.2 && < 2.3, + dataframe-operations >= 2.4 && < 2.5, dataframe-parsing >= 2.2 && < 2.3, dataframe-parsing >= 2.2 && < 2.3, mmap >= 0.5.8 && < 0.7, @@ -91,10 +91,10 @@ Properties.Csv build-depends: base >= 4 && < 5, containers >= 0.6.7 && < 0.10, - dataframe-core >= 2.2 && < 2.3, + dataframe-core >= 2.4 && < 2.5, dataframe-csv >= 2.3 && < 2.4, dataframe-fastcsv, - dataframe-operations >= 2.2 && < 2.3, + dataframe-operations >= 2.4 && < 2.5, dataframe-parsing >= 2.2 && < 2.3, directory >= 1.3.0.0 && < 2, HUnit >= 1.6 && < 1.8, diff --git a/src/DataFrame/IO/CSV/Fast/Columns.hs b/src/DataFrame/IO/CSV/Fast/Columns.hs --- a/src/DataFrame/IO/CSV/Fast/Columns.hs +++ b/src/DataFrame/IO/CSV/Fast/Columns.hs @@ -44,6 +44,7 @@ import DataFrame.IO.CSV.Fast.Passes import DataFrame.IO.CSV.Fast.Slice (extractField) import DataFrame.Internal.Column (Column, ensureOptional, fromVector) +import DataFrame.Internal.DictEncode (dictCompactColumn) import DataFrame.Operations.Typing ( ParseOptions (..), ParsingAssumption (..), @@ -218,13 +219,13 @@ 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) + pure (finishMode mode (dictCompactColumn 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) + pure (finishMode mode (dictCompactColumn c), False) {- | The original Text-materializing pipeline, kept for 'EitherRead' and exotic schema types (cold paths). diff --git a/src/DataFrame/IO/CSV/Fast/Core.hs b/src/DataFrame/IO/CSV/Fast/Core.hs --- a/src/DataFrame/IO/CSV/Fast/Core.hs +++ b/src/DataFrame/IO/CSV/Fast/Core.hs @@ -17,7 +17,7 @@ import qualified Data.Vector as Vector import qualified Data.Vector.Storable as VS -import Control.Exception (throwIO) +import Control.Exception (evaluate, throwIO) import Control.Monad (unless, when) import Data.Char (ord) import qualified Data.Text as T @@ -49,7 +49,13 @@ ) 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.IO.CSV.Fast.Slice ( + FieldCtx (..), + FieldLayout (..), + compactRows, + extractField, + stripBom, + ) import DataFrame.Internal.DataFrame (DataFrame (..), forceDataFrame) import DataFrame.Operations.Typing (effectiveSafeRead, parseWithTypes) @@ -100,8 +106,7 @@ FieldCtx { fcFile = file , fcBS = byteStringView contentLen file - , fcDelims = indices - , fcRowEnds = rowEnds + , fcLayout = LayoutFlat indices rowEnds , fcContentLen = contentLen , fcTrim = fastCsvTrimUnquoted opts } @@ -137,10 +142,17 @@ when (fastCsvOnRaggedRow opts == RaiseOnRagged) $ checkNoRaggedRows rowEnds dataRows numCol traceMarkerIO "fastcsv:rows-done" + mapM_ (\(nm, ix) -> evaluate nm >> evaluate ix) wanted + ctx' <- + evaluate $ + case compactRows contentLen indices rowEnds numCol of + Just lay -> ctx{fcLayout = lay} + Nothing -> ctx + traceMarkerIO "fastcsv:index-compacted" VS.unsafeWith file $ \filePtr -> do let env = ColumnEnv - { ceCtx = ctx + { ceCtx = ctx' , ceRows = dataRows , ceNumRow = numRow , cePtr = filePtr diff --git a/src/DataFrame/IO/CSV/Fast/Parallel.hs b/src/DataFrame/IO/CSV/Fast/Parallel.hs --- a/src/DataFrame/IO/CSV/Fast/Parallel.hs +++ b/src/DataFrame/IO/CSV/Fast/Parallel.hs @@ -22,8 +22,9 @@ import qualified Data.Vector.Storable as VS import Control.Concurrent (getNumCapabilities) -import Control.Exception (throwIO) +import Control.Exception (evaluate, throwIO) import Data.Either (partitionEithers) +import qualified Data.List as L import Debug.Trace (traceMarkerIO) import System.Mem (performMajorGC) @@ -33,6 +34,8 @@ import DataFrame.IO.CSV.Fast.Workers (pooledRun) import DataFrame.Internal.Column (Column, forceColumn) import DataFrame.Internal.ColumnBuilder (mergeColumns) +import DataFrame.Internal.DictEncode (dictCompactColumn) +import DataFrame.Operations.Typing (SafeReadMode) -- | Below this input size the fan-out overhead outweighs the parallelism. parallelThresholdBytes :: Int @@ -77,12 +80,17 @@ mapM (uncurry (buildColumn env)) wanted | otherwise = do width <- getNumCapabilities - let plans = [planColumn env name fieldIx | (name, fieldIx) <- wanted] - fields = map snd wanted + plans <- + mapM + (\(name, fieldIx) -> evaluate (planColumn env name fieldIx)) + wanted + let fields = map snd wanted ranges = chunkRanges chunks (ceNumRow env) traceMarkerIO "fastcsv:fanout" perChunk <- pooledRun width [runChunk env (zip fields plans) r | r <- ranges] traceMarkerIO "fastcsv:join-done" + let byCol = L.transpose perChunk + mapM_ (mapM_ (\cc -> cc `seq` pure ())) byCol -- 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 @@ -90,27 +98,54 @@ 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 - name - plan - fieldIx - [(rng, cols !! slot) | (rng, cols) <- zip ranges perChunk] - forceColumn col `seq` pure r - | (slot, ((name, fieldIx), plan)) <- zip [0 ..] (zip wanted plans) - ] + -- is deferred to the final forceDataFrame walk. 'evaluate' resolves + -- the plan dispatch now, so a schema task is a closure over its own + -- column's data only — never 'env'. + tasks <- + mapM + (\(wp, cps) -> evaluate (mergeTask width env ranges wp cps)) + (zip (zip wanted plans) byCol) + merged <- pooledRun width tasks traceMarkerIO "fastcsv:merge-done" pure merged where chunks = min nChunks (ceNumRow env) +mergeTask :: + Int -> + ColumnEnv -> + [(Int, Int)] -> + ((T.Text, Int), ColumnPlan) -> + [ChunkCol] -> + IO (Column, Bool) +mergeTask width env ranges ((name, fieldIx), plan) colParts = + case plan of + PlanSchema mode _ t -> mergeSchema width name mode t parts + _ -> forced (mergeColumn width env name plan fieldIx parts) + where + parts = zip ranges colParts + +-- | Splice a schema column's chunks; no 'ColumnEnv' capture by design. +mergeSchema :: + Int -> + T.Text -> + SafeReadMode -> + TargetType -> + [((Int, Int), ChunkCol)] -> + IO (Column, Bool) +mergeSchema width name mode t parts = + case partitionEithers [r | (_, CCSchema r) <- parts] of + ([], cs) -> forced $ do + c <- mergePassCols width cs + pure (finishMode mode (dictCompactColumn c), False) + (failed, _) -> throwIO (schemaError name t (minimum failed)) + +-- | Force the merged payload inside the task, not on the result walk. +forced :: IO (Column, Bool) -> IO (Column, Bool) +forced act = do + r@(col, _) <- act + forceColumn col `seq` pure r + -- | Even split of @n@ rows into @k@ @(offset, length)@ ranges. chunkRanges :: Int -> Int -> [(Int, Int)] chunkRanges k n = @@ -155,15 +190,10 @@ 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)) + PlanSchema mode _ t -> mergeSchema width name mode t parts PlanChain mode nspec steps _ -> do c <- mergeChain width env nspec steps col [(r, oc) | (r, CCChain oc) <- parts] - pure (finishMode mode c, False) + pure (finishMode mode (dictCompactColumn 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 diff --git a/src/DataFrame/IO/CSV/Fast/Slice.hs b/src/DataFrame/IO/CSV/Fast/Slice.hs --- a/src/DataFrame/IO/CSV/Fast/Slice.hs +++ b/src/DataFrame/IO/CSV/Fast/Slice.hs @@ -7,6 +7,8 @@ -} module DataFrame.IO.CSV.Fast.Slice ( FieldCtx (..), + FieldLayout (..), + compactRows, withFieldSlice, withParseSlice, sliceBS, @@ -21,28 +23,83 @@ import qualified Data.Text as Text import qualified Data.Text.Encoding as TextEncoding import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Storable.Mutable as VSM +import Control.Monad.ST (runST) import Data.Text (Text) -import Data.Word (Word8) +import Data.Word (Word16, Word8) import Foreign.C.Types (CSize) import DataFrame.IO.CSV.Fast.Index (cr, quote) +-- | How @(row, col)@ resolves to delimiter byte positions. +data FieldLayout + = {- | Flat delimiter positions + row-end ordinals: the shape the scanner + emits. Handles ragged and blank rows; 8 bytes per delimiter. + -} + LayoutFlat !(VS.Vector CSize) !(VS.Vector Int) + | {- | Uniform-stride rows: byte start per row plus 'Word16' deltas of + each field-end from its row start (@deltas ! (r*stride + col)@). + Built by 'compactRows'; ~3.5x smaller than the flat index. + -} + LayoutRows !(VS.Vector Int) !(VS.Vector Word16) !Int + -- | 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. + , fcLayout :: !FieldLayout + -- ^ Delimiter index (flat, or row-compacted by 'compactRows'). , fcContentLen :: !Int , fcTrim :: !Bool -- ^ 'DataFrame.IO.CSV.fastCsvTrimUnquoted'. } +compactRows :: + Int -> VS.Vector CSize -> VS.Vector Int -> Int -> Maybe FieldLayout +compactRows contentLen delims rowEnds stride + | totalRows == 0 || stride <= 0 = Nothing + | VS.length delims < totalRows * stride = Nothing + | otherwise = runST $ do + starts <- VSM.unsafeNew totalRows + deltas <- VSM.unsafeNew (totalRows * stride) + let go !r !base !rowStart + | r >= totalRows = pure True + | VS.unsafeIndex rowEnds r /= base + stride - 1 = pure False + | otherwise = do + let lastRaw = + fromIntegral (VS.unsafeIndex delims (base + stride - 1)) + lastPos = min contentLen lastRaw + if lastPos - rowStart > 65535 + then pure False + else do + VSM.unsafeWrite starts r rowStart + let fill !j + | j >= stride = pure () + | otherwise = do + let p = + min + contentLen + (fromIntegral (VS.unsafeIndex delims (base + j))) + VSM.unsafeWrite + deltas + (r * stride + j) + (fromIntegral (p - rowStart)) + fill (j + 1) + fill 0 + go (r + 1) (base + stride) (lastRaw + 1) + ok <- go 0 0 0 + if ok + then do + s <- VS.unsafeFreeze starts + d <- VS.unsafeFreeze deltas + pure (Just (LayoutRows s d stride)) + else pure Nothing + where + totalRows = VS.length rowEnds + {- | 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 @@ -51,36 +108,56 @@ -} {-# 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 +withFieldSlice ctx r col k = case fcLayout ctx of + LayoutRows starts deltas stride -> + if col >= stride 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 + let !rowStart = VS.unsafeIndex starts r + !fieldEndClamped = + rowStart + + fromIntegral (VS.unsafeIndex deltas (r * stride + col)) + !fieldStart = + if col == 0 + then rowStart else - fromIntegral - (VS.unsafeIndex (fcDelims ctx) (boundaryIdx - 1)) + rowStart + + fromIntegral + (VS.unsafeIndex deltas (r * stride + col - 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 + in finish fieldStart fieldEndClamped + LayoutFlat delims rowEnds -> + let 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 delims boundaryIdx) :: Int + fieldEndClamped = min fieldEndRaw (fcContentLen ctx) + fieldStart = + if boundaryIdx == 0 + then 0 + else + fromIntegral + (VS.unsafeIndex delims (boundaryIdx - 1)) + + 1 + in finish fieldStart fieldEndClamped + where + finish !fieldStart !fieldEndClamped = + let 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 diff --git a/src/DataFrame/IO/CSV/Fast/TextMerge.hs b/src/DataFrame/IO/CSV/Fast/TextMerge.hs --- a/src/DataFrame/IO/CSV/Fast/TextMerge.hs +++ b/src/DataFrame/IO/CSV/Fast/TextMerge.hs @@ -17,6 +17,7 @@ import Control.Monad (void) import Control.Monad.ST (stToIO) +import Data.Int (Int32) import DataFrame.IO.CSV.Fast.Workers (pooledRun) import DataFrame.Internal.Column (Column (..)) import DataFrame.Internal.ColumnMerge ( @@ -25,7 +26,7 @@ spliceBitmaps, tcRows, ) -import DataFrame.Internal.PackedText (mkPackedContiguous) +import DataFrame.Internal.PackedText (mkPackedContiguous, mkPackedContiguous32) {- | Merge text chunks with @width@-way parallel byte copies + offset rebase, then wrap the shared buffer as 'PackedText'. Single chunks take @@ -39,26 +40,38 @@ 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 splice writeOff = + 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 + writeOff (rOff + i) (bOff + VU.unsafeIndex co i) + fill (i + 1) + fill 1 + | (c, bOff, rOff) <- zip3 cs byteOffs rowOffs + ] + -- The final width is known before allocation: Int32 offsets whenever + -- the merged buffer stays under 2^31 bytes (the common case). + packed <- + if totalBytes <= fromIntegral (maxBound :: Int32) + then do + offsMV <- VUM.unsafeNew (totalRows + 1) :: IO (VUM.IOVector Int32) + VUM.unsafeWrite offsMV 0 0 + splice (\i v -> VUM.unsafeWrite offsMV i (fromIntegral v)) + mkPackedContiguous32 + <$> stToIO (A.unsafeFreeze marr) + <*> VU.unsafeFreeze offsMV + else do + offsMV <- VUM.unsafeNew (totalRows + 1) :: IO (VUM.IOVector Int) + VUM.unsafeWrite offsMV 0 0 + splice (VUM.unsafeWrite offsMV) + mkPackedContiguous + <$> stToIO (A.unsafeFreeze marr) + <*> VU.unsafeFreeze offsMV let !bm = spliceBitmaps [(tcBitmap c, tcRows c) | c <- cs] - pure (PackedText bm (mkPackedContiguous arr offs)) + pure (PackedText bm packed) diff --git a/src/DataFrame/IO/CSV/Fast/Workers.hs b/src/DataFrame/IO/CSV/Fast/Workers.hs --- a/src/DataFrame/IO/CSV/Fast/Workers.hs +++ b/src/DataFrame/IO/CSV/Fast/Workers.hs @@ -12,7 +12,7 @@ import Control.Concurrent (forkIO) import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar) -import Control.Exception (SomeException, throwIO, try) +import Control.Exception (ErrorCall (..), SomeException, throwIO, try) import Control.Monad (when) import Data.IORef (atomicModifyIORef', newIORef) @@ -31,6 +31,10 @@ {- | 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. + +Each action slot is cleared before the action runs, so data captured by a +completed closure is unreachable as soon as it finishes — callers rely on +this to release per-column chunk payloads during the parallel merge. -} pooledRun :: Int -> [IO a] -> IO [a] pooledRun width actions @@ -38,14 +42,18 @@ | otherwise = do next <- newIORef 0 out <- VM.unsafeNew n - let acts = V.fromListN n actions - worker = do + acts <- VM.unsafeNew n + sequence_ [VM.unsafeWrite acts i a | (i, a) <- zip [0 ..] actions] + let worker = do i <- atomicModifyIORef' next (\j -> (j + 1, j)) when (i < n) $ do - r <- acts V.! i + act <- VM.unsafeRead acts i + VM.unsafeWrite acts i consumed + r <- act VM.write out i r worker _ <- forkJoin (replicate width worker) V.toList <$> V.freeze out where n = length actions + consumed = throwIO (ErrorCall "pooledRun: slot already consumed")