dataframe-lazy 2.4.1.0 → 2.4.2.0
raw patch · 6 files changed
+379/−45 lines, 6 filesdep +HUnitdep +dataframe-lazydep ~directory
Dependencies added: HUnit, dataframe-lazy
Dependency ranges changed: directory
Files
- dataframe-lazy.cabal +20/−3
- src/DataFrame/Lazy/Internal/DataFrame.hs +45/−4
- src/DataFrame/Lazy/Internal/Executor.hs +109/−32
- src/DataFrame/Lazy/Internal/LogicalPlan.hs +4/−1
- src/DataFrame/Lazy/Internal/Optimizer.hs +9/−5
- test/StreamingCsv.hs +192/−0
dataframe-lazy.cabal view
@@ -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
src/DataFrame/Lazy/Internal/DataFrame.hs view
@@ -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 }
src/DataFrame/Lazy/Internal/Executor.hs view
@@ -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
src/DataFrame/Lazy/Internal/LogicalPlan.hs view
@@ -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 ++ " <reader>" show (CsvSourceStreaming path sep _) = "CsvSourceStreaming " ++ show path ++ " " ++ show sep ++ " <reader>"+ show (CsvSourceStreamingBytes path sep _) =+ "CsvSourceStreamingBytes " ++ show path ++ " " ++ show sep ++ " <reader>" show (ParquetSource path) = "ParquetSource " ++ show path -- | Sort direction used in Sort nodes and the public API.
src/DataFrame/Lazy/Internal/Optimizer.hs view
@@ -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
+ test/StreamingCsv.hs view
@@ -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