diff --git a/dataframe-lazy.cabal b/dataframe-lazy.cabal --- a/dataframe-lazy.cabal +++ b/dataframe-lazy.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: dataframe-lazy -version: 2.4.1.0 +version: 2.4.2.0 synopsis: Lazy query engine for the dataframe ecosystem. description: The lazy/streaming query engine: relational-algebra plans, optimizer, @@ -47,9 +47,8 @@ containers >= 0.6.7 && < 0.10, dataframe-core >= 2.5 && < 2.6, dataframe-csv >= 2.3 && < 2.4, - dataframe-csv >= 2.3 && < 2.4, dataframe-operations >= 2.5 && < 2.6, - dataframe-parquet >= 1.5 && < 1.6, + dataframe-parquet >= 1.5.0.1 && < 1.6, dataframe-parsing >= 2.2 && < 2.3, directory >= 1.3.0.0 && < 2, filepath >= 1.4 && < 2, @@ -59,4 +58,22 @@ text >= 2.1 && < 3, vector >= 0.13 && < 0.15 hs-source-dirs: src + default-language: Haskell2010 + +test-suite streaming-csv + import: warnings + type: exitcode-stdio-1.0 + main-is: StreamingCsv.hs + hs-source-dirs: test + build-depends: base >= 4 && < 5, + bytestring >= 0.11 && < 0.14, + containers >= 0.6.7 && < 0.10, + dataframe-core >= 2.5 && < 2.6, + dataframe-csv >= 2.3 && < 2.4, + dataframe-lazy, + dataframe-parsing >= 2.2 && < 2.3, + directory >= 1.3 && < 2, + HUnit >= 1.6 && < 1.8, + temporary >= 1.3 && < 2, + text >= 2.1 && < 3 default-language: Haskell2010 diff --git a/src/DataFrame/Lazy/Internal/DataFrame.hs b/src/DataFrame/Lazy/Internal/DataFrame.hs --- a/src/DataFrame/Lazy/Internal/DataFrame.hs +++ b/src/DataFrame/Lazy/Internal/DataFrame.hs @@ -8,7 +8,12 @@ module DataFrame.Lazy.Internal.DataFrame where import qualified Data.Text as T -import DataFrame.IO.CSV (CsvReader, readSeparated) +import DataFrame.IO.CSV ( + CsvBytesReader, + CsvReader, + decodeSeparatedStrict, + readSeparated, + ) import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E @@ -88,12 +93,31 @@ , batchSize = 1_000_000 } -{- | Like 'scanCsvWith', but the file is read in bounded-memory windows -instead of one pass per chunk — for files too large to hold in memory even -after the schema's projection. +{- | Scan a CSV file in bounded-memory windows with the default in-tree +strict reader. Windows are decoded directly from memory, without temporary +files. +Use this for files too large to hold in memory even after the schema's +projection. + ==== __Example__ @ +ghci> L.runDataFrame (L.scanCsvStreaming schema "huge.csv") + +@ +-} +scanCsvStreaming :: Schema -> T.Text -> LazyDataFrame +scanCsvStreaming = scanCsvStreamingBytesWith decodeSeparatedStrict + +{- | Like 'scanCsvWith', but the file is read in bounded-memory windows. + +This compatibility entry point accepts an existing path-based 'CsvReader'. +Because such a reader can only consume file paths, each window is staged in a +temporary file. New readers should use 'scanCsvStreamingBytesWith' to decode +windows directly from memory. + +==== __Example__ +@ ghci> L.runDataFrame (L.scanCsvStreamingWith Fast.fastReadCsvWithOpts schema \"huge.csv\") @ @@ -102,6 +126,23 @@ scanCsvStreamingWith reader schema path = LazyDataFrame { plan = Scan (CsvSourceStreaming (T.unpack path) ',' reader) schema + , batchSize = 1_000_000 + } + +{- | Stream a CSV file in bounded-memory windows decoded by an in-memory +'CsvBytesReader'. This avoids the temporary-file staging required by +'scanCsvStreamingWith'. + +==== __Example__ +@ +ghci> L.runDataFrame (L.scanCsvStreamingBytesWith decodeSeparatedStrict schema "huge.csv") + +@ +-} +scanCsvStreamingBytesWith :: CsvBytesReader -> Schema -> T.Text -> LazyDataFrame +scanCsvStreamingBytesWith reader schema path = + LazyDataFrame + { plan = Scan (CsvSourceStreamingBytes (T.unpack path) ',' reader) schema , batchSize = 1_000_000 } diff --git a/src/DataFrame/Lazy/Internal/Executor.hs b/src/DataFrame/Lazy/Internal/Executor.hs --- a/src/DataFrame/Lazy/Internal/Executor.hs +++ b/src/DataFrame/Lazy/Internal/Executor.hs @@ -18,14 +18,15 @@ foldBatches, ) where -import Control.Concurrent (forkIO, getNumCapabilities) +import Control.Concurrent (forkFinally, forkIO, getNumCapabilities) import Control.Concurrent.Async (mapConcurrently) import Control.Concurrent.STM (atomically) import Control.Concurrent.STM.TBQueue (newTBQueueIO, readTBQueue, writeTBQueue) -import Control.Exception (evaluate) +import Control.Exception (evaluate, finally, throwIO) import Control.Monad (filterM, forM, forM_, unless, when) import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as C8 +import qualified Data.ByteString.Unsafe as BSU import Data.IORef import Data.Int (Int16, Int32, Int64, Int8) import qualified Data.Map as M @@ -37,7 +38,12 @@ import qualified Data.Vector as VB import qualified Data.Vector.Unboxed as VU import Data.Word (Word16, Word32, Word64, Word8) -import DataFrame.IO.CSV (CsvReader, ReadOptions (..), schemaReadOptions) +import DataFrame.IO.CSV ( + CsvBytesReader, + CsvReader, + ReadOptions (..), + schemaReadOptions, + ) import qualified DataFrame.IO.Parquet as Parquet import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.DataFrame as D @@ -56,7 +62,7 @@ import System.Directory (doesDirectoryExist, removeFile) import System.FilePath (()) import System.FilePath.Glob (glob) -import System.IO (IOMode (ReadMode), hIsEOF, withFile) +import System.IO (IOMode (ReadMode), withFile) import System.IO.Temp (emptySystemTempFile) import Type.Reflection (typeRep) @@ -115,6 +121,7 @@ isBounded :: PhysicalPlan -> Bool isBounded (PhysicalScan (CsvSource{}) _) = True isBounded (PhysicalScan (CsvSourceStreaming{}) _) = False +isBounded (PhysicalScan (CsvSourceStreamingBytes{}) _) = False isBounded (PhysicalScan (ParquetSource _) _) = True isBounded (PhysicalProject _ c) = isBounded c isBounded (PhysicalFilter _ c) = isBounded c @@ -162,6 +169,8 @@ executeCsvScan path sep reader cfg buildStream (PhysicalScan (CsvSourceStreaming path sep reader) cfg) = executeCsvScanStreaming path sep reader cfg +buildStream (PhysicalScan (CsvSourceStreamingBytes path sep reader) cfg) = + executeCsvScanStreamingBytes path sep reader cfg buildStream (PhysicalScan (ParquetSource path) cfg) = executeParquetScan path cfg buildStream (PhysicalSpill child path) = do @@ -445,7 +454,7 @@ let partialExpr = E.UExpr ( E.Agg - (E.MergeAgg n seed step merge (id :: acc -> acc)) + (E.FoldAgg ("partial_" <> n) (Just seed) step) inner ) mergeExpr = @@ -543,43 +552,111 @@ executeCsvScanStreaming :: FilePath -> Char -> CsvReader -> ScanConfig -> IO Stream -executeCsvScanStreaming path sep reader cfg = do +executeCsvScanStreaming path sep reader = + executeCsvScanStreamingWith path sep (Left reader) + +executeCsvScanStreamingBytes :: + FilePath -> Char -> CsvBytesReader -> ScanConfig -> IO Stream +executeCsvScanStreamingBytes path sep reader = + executeCsvScanStreamingWith path sep (Right reader) + +executeCsvScanStreamingWith :: + FilePath -> + Char -> + Either CsvReader CsvBytesReader -> + ScanConfig -> + IO Stream +executeCsvScanStreamingWith path sep reader cfg = do let opts = scanReadOptions sep cfg batchSz = scanBatchSize cfg windowBytes = 64 * 1024 * 1024 :: Int queue <- newTBQueueIO 8 - _ <- forkIO $ - withFile path ReadMode $ \h -> do - header <- C8.hGetLine h - let feed bytes = - unless (BS.null bytes) $ do - p <- emptySystemTempFile "lazy_csv_win_.csv" - BS.writeFile p (header <> BS.singleton nl <> bytes) - df <- reader opts p - removeFile p - forM_ (sliceIntoBatches batchSz df) $ \b -> - atomically (writeTBQueue queue (Just b)) - loop leftover = do - eof <- hIsEOF h - if eof - then feed leftover >> atomically (writeTBQueue queue Nothing) - else do - chunk <- BS.hGetSome h windowBytes - let buf = leftover <> chunk - case BS.elemIndexEnd nl buf of - Nothing -> loop buf - Just i -> feed (BS.take i buf) >> loop (BS.drop (i + 1) buf) - loop BS.empty + _ <- + forkFinally + ( withFile path ReadMode $ \h -> do + rawHeader <- C8.hGetLine h + let header = stripUtf8Bom rawHeader + let feed bytes = + unless (BS.null bytes) $ do + let window = header <> BS.singleton nl <> bytes + parsed <- case reader of + Right bytesReader -> bytesReader opts window + Left pathReader -> do + p <- emptySystemTempFile "lazy_csv_win_.csv" + (BS.writeFile p window >> pathReader opts p) + `finally` removeFile p + forM_ (sliceIntoBatches batchSz parsed) $ \b -> + atomically (writeTBQueue queue (Right (Just b))) + loop leftover = do + chunk <- BS.hGet h windowBytes + if BS.null chunk + then feed leftover + else do + let buf = leftover <> chunk + case lastCompleteRecordNewline buf of + Nothing -> loop buf + Just i -> feed (BS.take i buf) >> loop (BS.drop (i + 1) buf) + loop BS.empty + ) + ( \result -> + atomically $ + writeTBQueue queue $ + case result of + Left err -> Left err + Right () -> Right Nothing + ) return . Stream $ do - mb <- atomically (readTBQueue queue) - case mb of - Nothing -> atomically (writeTBQueue queue Nothing) >> return Nothing - Just df -> + item <- atomically (readTBQueue queue) + case item of + Left err -> throwIO err + Right Nothing -> + atomically (writeTBQueue queue (Right Nothing)) >> return Nothing + Right (Just df) -> let df' = case scanPushdownPredicate cfg of Nothing -> df Just p -> Sub.filterWhere p df in return (Just df') where + nl :: Word8 + nl = 0x0A + + stripUtf8Bom bytes = + Data.Maybe.fromMaybe bytes (BS.stripPrefix "\xEF\xBB\xBF" bytes) + +{- | Find the last record-ending LF in a CSV buffer. The common unquoted case +uses bytestring's optimized reverse search; only buffers containing a quote +need the scalar RFC 4180 quote-state scan. The buffer always starts at a record +boundary, so rescanning it after an incomplete quoted record is sufficient to +handle a doubled or closing quote split across read windows. +-} +lastCompleteRecordNewline :: BS.ByteString -> Maybe Int +lastCompleteRecordNewline bytes + | not (BS.elem quote bytes) = BS.elemIndexEnd nl bytes + | otherwise = go 0 False Nothing + where + !len = BS.length bytes + + -- Every read stays inside the branch that has already bounds-checked its + -- offset: a guard binding shared across alternatives can be forced on the + -- terminating @i == len@ iteration, indexing one byte past the buffer. + go !i !inside !lastNewline + | i >= len = lastNewline + | otherwise = case BSU.unsafeIndex bytes i of + current + | current == quote -> + if inside && quoteAt (i + 1) + then go (i + 2) True lastNewline + else go (i + 1) (not inside) lastNewline + | current == nl && not inside -> go (i + 1) inside (Just i) + | otherwise -> go (i + 1) inside lastNewline + + quoteAt !i + | i >= len = False + | otherwise = BSU.unsafeIndex bytes i == quote + + quote :: Word8 + quote = 0x22 + nl :: Word8 nl = 0x0A diff --git a/src/DataFrame/Lazy/Internal/LogicalPlan.hs b/src/DataFrame/Lazy/Internal/LogicalPlan.hs --- a/src/DataFrame/Lazy/Internal/LogicalPlan.hs +++ b/src/DataFrame/Lazy/Internal/LogicalPlan.hs @@ -3,7 +3,7 @@ module DataFrame.Lazy.Internal.LogicalPlan where import qualified Data.Text as T -import DataFrame.IO.CSV (CsvReader) +import DataFrame.IO.CSV (CsvBytesReader, CsvReader) import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E import DataFrame.Operations.Join (JoinType) @@ -14,6 +14,7 @@ = -- | path, separator, CSV reader (e.g. attoparsec or SIMD) CsvSource FilePath Char CsvReader | CsvSourceStreaming FilePath Char CsvReader + | CsvSourceStreamingBytes FilePath Char CsvBytesReader | ParquetSource FilePath instance Show DataSource where @@ -21,6 +22,8 @@ "CsvSource " ++ show path ++ " " ++ show sep ++ " " show (CsvSourceStreaming path sep _) = "CsvSourceStreaming " ++ show path ++ " " ++ show sep ++ " " + show (CsvSourceStreamingBytes path sep _) = + "CsvSourceStreamingBytes " ++ show path ++ " " ++ show sep ++ " " show (ParquetSource path) = "ParquetSource " ++ show path -- | Sort direction used in Sort nodes and the public API. diff --git a/src/DataFrame/Lazy/Internal/Optimizer.hs b/src/DataFrame/Lazy/Internal/Optimizer.hs --- a/src/DataFrame/Lazy/Internal/Optimizer.hs +++ b/src/DataFrame/Lazy/Internal/Optimizer.hs @@ -103,7 +103,7 @@ let keySet = S.fromList [l, r] lRef = fmap (S.union keySet) (referencedCols left) rRef = fmap (S.union keySet) (referencedCols right) - in liftMaybe2 S.union lRef rRef + in S.union <$> lRef <*> rRef referencedCols (Aggregate keys aggs child) = let aggCols = S.fromList (keys <> concatMap (uExprCols . snd) aggs) in fmap (S.union aggCols) (referencedCols child) @@ -112,10 +112,6 @@ referencedCols (Limit _ child) = referencedCols child referencedCols (SourceDF _) = Nothing -liftMaybe2 :: (a -> b -> c) -> Maybe a -> Maybe b -> Maybe c -liftMaybe2 f (Just a) (Just b) = Just (f a b) -liftMaybe2 _ _ _ = Nothing - uExprCols :: E.UExpr -> [T.Text] uExprCols (E.UExpr expr) = E.getColumns expr @@ -177,6 +173,14 @@ toPhysical batchSz (Scan (CsvSourceStreaming path sep reader) schema) = PhysicalScan (CsvSourceStreaming path sep reader) + (ScanConfig batchSz sep schema Nothing) +toPhysical batchSz (Filter p (Scan (CsvSourceStreamingBytes path sep reader) schema)) = + PhysicalScan + (CsvSourceStreamingBytes path sep reader) + (ScanConfig batchSz sep schema (Just p)) +toPhysical batchSz (Scan (CsvSourceStreamingBytes path sep reader) schema) = + PhysicalScan + (CsvSourceStreamingBytes path sep reader) (ScanConfig batchSz sep schema Nothing) toPhysical batchSz (Filter p (Scan (ParquetSource path) schema)) = PhysicalScan diff --git a/test/StreamingCsv.hs b/test/StreamingCsv.hs new file mode 100644 --- /dev/null +++ b/test/StreamingCsv.hs @@ -0,0 +1,192 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module Main (main) where + +import Control.Exception (SomeException, bracket, try) +import Control.Monad (when) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as C8 +import Data.IORef (modifyIORef', newIORef, readIORef) +import Data.List (isInfixOf) +import qualified Data.Map.Strict as M +import qualified Data.Text as T +import qualified DataFrame.IO.CSV as Csv +import DataFrame.Internal.Column (Column (..)) +import qualified DataFrame.Internal.DataFrame as D +import qualified DataFrame.Lazy as L +import DataFrame.Schema (Schema (..), schemaType) +import System.Directory (removeFile) +import qualified System.Exit as Exit +import System.IO.Temp (emptySystemTempFile) +import Test.HUnit + +withRawCsv :: BS.ByteString -> (FilePath -> IO a) -> IO a +withRawCsv bytes = + bracket + ( do + path <- emptySystemTempFile "lazy_streaming_.csv" + BS.writeFile path bytes + pure path + ) + removeFile + +{- | Regression for the @hIsEOF >> hGetSome@ interaction: the handle's 8 KiB +buffer used to turn this 16 MiB input into more than two thousand parser and +temporary-file calls. A path-only reader remains source-compatible and should +now see one complete 64 MiB-or-EOF window. +-} +streamingPathReaderFillsWindow :: Test +streamingPathReaderFillsWindow = + TestLabel "streaming_path_reader_fills_window" $ TestCase $ do + let row = "1," <> BS.replicate 4094 0x78 <> "\n" + rowCount = 4096 + input = "id,payload\n" <> BS.concat (replicate rowCount row) + schema = Schema $ M.fromList [("id", schemaType @Int)] + withRawCsv input $ \path -> do + calls <- newIORef (0 :: Int) + let countingReader opts windowPath = do + modifyIORef' calls (+ 1) + Csv.readSeparated opts windowPath + df <- + L.runDataFrame + (L.scanCsvStreamingWith countingReader schema (T.pack path)) + assertEqual "all rows decoded" (rowCount, 1) (D.dataframeDimensions df) + readIORef calls >>= assertEqual "one parser window" 1 + +{- | A record newline exactly at the 64 MiB read boundary must not be lost or +duplicated. The byte reader records only small observations, avoiding the cost +of parsing the deliberately large payload. +-} +streamingWindowBoundaryIsLossless :: Test +streamingWindowBoundaryIsLossless = + TestLabel "streaming_window_boundary_is_lossless" $ TestCase $ do + let windowBytes = 64 * 1024 * 1024 + firstRecord = "1," <> BS.replicate (windowBytes - 3) 0x78 <> "\n" + input = BS.concat ["id,payload\n", firstRecord, "2,y\n"] + schema = Schema $ M.fromList [("id", schemaType @Int)] + withRawCsv input $ \path -> do + observations <- newIORef [] + let observingReader _ bytes = do + let (header, withNewline) = BS.break (== 0x0A) bytes + body = BS.drop 1 withNewline + observation = + ( C8.unpack header + , BS.length body + , BS.unpack (BS.take 2 body) + , if BS.null body then Nothing else Just (BS.last body) + ) + modifyIORef' observations (observation :) + pure D.empty + _ <- + L.runDataFrame + (L.scanCsvStreamingBytesWith observingReader schema (T.pack path)) + actual <- reverse <$> readIORef observations + assertEqual + "two complete records, each with the original header" + [ ("id,payload", windowBytes - 1, [0x31, 0x2C], Just 0x78) + , ("id,payload", 3, [0x32, 0x2C], Just 0x79) + ] + actual + +{- | An LF inside a quoted field is data, not a record boundary. If the quoted +record crosses a 64 MiB read boundary it must still reach the reader in one +complete window. +-} +streamingQuotedNewlineSpansWindow :: Test +streamingQuotedNewlineSpansWindow = + TestLabel "streaming_quoted_newline_spans_window" $ TestCase $ do + let windowBytes = 64 * 1024 * 1024 + firstChunk = "1,\"" <> BS.replicate (windowBytes - 4) 0x78 <> "\n" + input = BS.concat ["id,payload\n", firstChunk, "\"\n"] + schema = Schema $ M.fromList [("id", schemaType @Int)] + withRawCsv input $ \path -> do + bodyLengths <- newIORef [] + let observingReader _ bytes = do + let body = BS.drop 1 (snd (BS.break (== 0x0A) bytes)) + modifyIORef' bodyLengths (BS.length body :) + pure D.empty + _ <- + L.runDataFrame + (L.scanCsvStreamingBytesWith observingReader schema (T.pack path)) + actual <- reverse <$> readIORef bodyLengths + assertEqual + "quoted record stays in one parser window" + [windowBytes + 1] + actual + +{- | The default in-memory streaming path strips a file BOM and agrees with an +eager schema-projected read, including a final record without a newline. +-} +streamingDefaultMatchesEager :: Test +streamingDefaultMatchesEager = + TestLabel "streaming_default_matches_eager" $ TestCase $ do + let input = "\xEF\xBB\xBFid,name,value\n1,Ada,3.5\n2,Alan,4.5" + schema = + Schema $ + M.fromList + [ ("id", schemaType @Int) + , ("name", schemaType @T.Text) + , ("value", schemaType @Double) + ] + withRawCsv input $ \path -> do + let opts = Csv.schemaReadOptions schema + eager <- Csv.readSeparated opts path + streamed <- L.runDataFrame (L.scanCsvStreaming schema (T.pack path)) + assertEqual "streaming bytes reader matches eager reader" eager streamed + +{- | A parser failure happens on the producer thread; it must cross the queue +and fail the consumer rather than leaving 'runDataFrame' blocked forever. +-} +streamingReaderFailurePropagates :: Test +streamingReaderFailurePropagates = + TestLabel "streaming_reader_failure_propagates" $ TestCase $ do + let input = "id\n1\n" + schema = Schema $ M.fromList [("id", schemaType @Int)] + failingReader _ _ = ioError (userError "streaming-reader-marker") + withRawCsv input $ \path -> do + result <- + try + ( L.runDataFrame + (L.scanCsvStreamingBytesWith failingReader schema (T.pack path)) + ) + case result of + Left (err :: SomeException) -> + assertBool + ("unexpected exception: " <> show err) + ("streaming-reader-marker" `isInfixOf` show err) + Right (_ :: D.DataFrame) -> assertFailure "expected reader failure" + +-- | Schema-declared text stays in the CSV reader's packed representation. +streamingKeepsTextPacked :: Test +streamingKeepsTextPacked = + TestLabel "streaming_keeps_text_packed" $ TestCase $ do + let input = + "station,value\n" + <> "Alpha,1\nBeta,2\nAlpha,3\nBeta,4\n" + <> "Alpha,5\nBeta,6\nAlpha,7\nBeta,8\n" + schema = Schema $ M.fromList [("station", schemaType @T.Text)] + withRawCsv input $ \path -> do + df <- + L.runDataFrame + (L.scanCsvStreaming schema (T.pack path)) + case D.unsafeGetColumn "station" df of + PackedText{} -> pure () + _ -> assertFailure "expected the CSV text column to stay packed" + +tests :: Test +tests = + TestList + [ streamingPathReaderFillsWindow + , streamingWindowBoundaryIsLossless + , streamingQuotedNewlineSpansWindow + , streamingDefaultMatchesEager + , streamingReaderFailurePropagates + , streamingKeepsTextPacked + ] + +main :: IO () +main = do + result <- runTestTT tests + when (failures result > 0 || errors result > 0) Exit.exitFailure