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.1.0.0 +version: 2.2.0.0 synopsis: Lazy query engine for the dataframe ecosystem. description: The lazy/streaming query engine: relational-algebra plans, optimizer, @@ -47,10 +47,10 @@ containers >= 0.6.7 && < 0.10, dataframe-core >= 2.1 && < 2.2, dataframe-core >= 2.1 && < 2.2, - dataframe-csv >= 2.1 && < 2.2, - dataframe-csv >= 2.1 && < 2.2, + dataframe-csv >= 2.2 && < 2.3, + dataframe-csv >= 2.2 && < 2.3, dataframe-operations >= 2.1 && < 2.2, - dataframe-parquet >= 1.3 && < 1.4, + dataframe-parquet >= 1.4 && < 1.5, dataframe-parsing >= 2.1 && < 2.2, directory >= 1.3.0.0 && < 2, filepath >= 1.4 && < 2, 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,7 @@ module DataFrame.Lazy.Internal.DataFrame where import qualified Data.Text as T -import DataFrame.IO.CSV (CsvReader, readCsvWithSchema) +import DataFrame.IO.CSV (CsvReader, readSeparated) import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E @@ -54,13 +54,32 @@ fromDataFrame df = LazyDataFrame{plan = SourceDF df, batchSize = 1_000_000} {- | Scan a CSV file with the default comma separator and the in-tree -attoparsec reader. For the SIMD reader use 'scanCsvWith'. +strict reader. For the SIMD reader use 'scanCsvWith'. + +The 'Schema' both types and selects: only the columns it names are read, +matching 'scanParquet'. + +==== __Example__ +@ +ghci> schema = D.makeSchema [("id", D.schemaType \@Int), ("name", D.schemaType \@Text)] +ghci> L.runDataFrame (L.scanCsv schema \"customers.csv\") + +@ -} scanCsv :: Schema -> T.Text -> LazyDataFrame -scanCsv = scanCsvWith readCsvWithSchema +scanCsv = scanCsvWith readSeparated {- | Like 'scanCsv' but with an explicit CSV reader (e.g. the SIMD reader -@fastReadCsvWithSchema@ from @dataframe-fastcsv@). +@fastReadCsvWithOpts@ from @dataframe-fastcsv@). The scan derives the +reader's 'DataFrame.IO.CSV.ReadOptions' from the schema and separator, so +any 'CsvReader' projects. + +==== __Example__ +@ +ghci> import qualified DataFrame.IO.CSV.Fast as Fast +ghci> L.runDataFrame (L.scanCsvWith Fast.fastReadCsvWithOpts schema \"customers.csv\") + +@ -} scanCsvWith :: CsvReader -> Schema -> T.Text -> LazyDataFrame scanCsvWith reader schema path = @@ -69,6 +88,16 @@ , 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. + +==== __Example__ +@ +ghci> L.runDataFrame (L.scanCsvStreamingWith Fast.fastReadCsvWithOpts schema \"huge.csv\") + +@ +-} scanCsvStreamingWith :: CsvReader -> Schema -> T.Text -> LazyDataFrame scanCsvStreamingWith reader schema path = LazyDataFrame @@ -76,11 +105,25 @@ , batchSize = 1_000_000 } --- | Scan a character-separated file with the default attoparsec reader. +{- | Scan a character-separated file with the default strict reader. + +==== __Example__ +@ +ghci> L.runDataFrame (L.scanSeparated ';' schema \"customers.txt\") + +@ +-} scanSeparated :: Char -> Schema -> T.Text -> LazyDataFrame -scanSeparated = scanSeparatedWith readCsvWithSchema +scanSeparated = scanSeparatedWith readSeparated --- | Like 'scanSeparated' but with an explicit CSV reader. +{- | Like 'scanSeparated' but with an explicit CSV reader. + +==== __Example__ +@ +ghci> L.runDataFrame (L.scanSeparatedWith Fast.fastReadCsvWithOpts ';' schema \"customers.txt\") + +@ +-} scanSeparatedWith :: CsvReader -> Char -> Schema -> T.Text -> LazyDataFrame scanSeparatedWith reader sep schema path = 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 @@ -37,7 +37,7 @@ 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) +import DataFrame.IO.CSV (CsvReader, ReadOptions (..), schemaReadOptions) import qualified DataFrame.IO.Parquet as Parquet import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.DataFrame as D @@ -160,8 +160,8 @@ buildStream :: PhysicalPlan -> IO Stream buildStream (PhysicalScan (CsvSource path sep reader) cfg) = executeCsvScan path sep reader cfg -buildStream (PhysicalScan (CsvSourceStreaming path _sep reader) cfg) = - executeCsvScanStreaming path reader cfg +buildStream (PhysicalScan (CsvSourceStreaming path sep reader) cfg) = + executeCsvScanStreaming path sep reader cfg buildStream (PhysicalScan (ParquetSource path) cfg) = executeParquetScan path cfg buildStream (PhysicalSpill child path) = do @@ -502,18 +502,25 @@ -- CSV scan implementation -- --------------------------------------------------------------------------- +{- | The options a CSV scan reads with: the plan's schema supplies the column +types and the projection, the source supplies the separator. +-} +scanReadOptions :: Char -> ScanConfig -> ReadOptions +scanReadOptions sep cfg = + (schemaReadOptions (scanSchema cfg)){columnSeparator = sep} + {- | SIMD-parallel CSV scan: the file is split at newline boundaries into one slice per capability, parsed concurrently, sliced into batches, and fed through a bounded queue. Pushdown predicates are applied per batch by the consumer. -} executeCsvScan :: FilePath -> Char -> CsvReader -> ScanConfig -> IO Stream -executeCsvScan path _sep reader cfg = do +executeCsvScan path sep reader cfg = do nCaps <- getNumCapabilities chunkPaths <- splitCsvAtNewlines (max 1 nCaps) path - let schema = scanSchema cfg + let opts = scanReadOptions sep cfg batchSz = scanBatchSize cfg - chunkDfs <- mapConcurrently (reader schema) chunkPaths + chunkDfs <- mapConcurrently (reader opts) chunkPaths mapM_ removeFile chunkPaths queue <- newTBQueueIO (fromIntegral (max 4 (2 * nCaps))) @@ -534,9 +541,10 @@ in return (Just df') ) -executeCsvScanStreaming :: FilePath -> CsvReader -> ScanConfig -> IO Stream -executeCsvScanStreaming path reader cfg = do - let schema = scanSchema cfg +executeCsvScanStreaming :: + FilePath -> Char -> CsvReader -> ScanConfig -> IO Stream +executeCsvScanStreaming path sep reader cfg = do + let opts = scanReadOptions sep cfg batchSz = scanBatchSize cfg windowBytes = 64 * 1024 * 1024 :: Int queue <- newTBQueueIO 8 @@ -547,7 +555,7 @@ unless (BS.null bytes) $ do p <- emptySystemTempFile "lazy_csv_win_.csv" BS.writeFile p (header <> BS.singleton nl <> bytes) - df <- reader schema p + df <- reader opts p removeFile p forM_ (sliceIntoBatches batchSz df) $ \b -> atomically (writeTBQueue queue (Just b))