diff --git a/dataframe-lazy.cabal b/dataframe-lazy.cabal --- a/dataframe-lazy.cabal +++ b/dataframe-lazy.cabal @@ -1,6 +1,6 @@ cabal-version: 2.4 name: dataframe-lazy -version: 1.1.0.0 +version: 1.1.0.2 synopsis: Lazy query engine for the dataframe ecosystem. description: The lazy/streaming query engine: relational-algebra plans, optimizer, 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 @@ -73,6 +73,13 @@ , batchSize = 1_000_000 } +scanCsvStreamingWith :: CsvReader -> Schema -> T.Text -> LazyDataFrame +scanCsvStreamingWith reader schema path = + LazyDataFrame + { plan = Scan (CsvSourceStreaming (T.unpack path) ',' reader) schema + , batchSize = 1_000_000 + } + -- | Scan a character-separated file with the default attoparsec reader. scanSeparated :: Char -> Schema -> T.Text -> LazyDataFrame scanSeparated = scanSeparatedWith readCsvWithSchema 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 @@ -43,8 +43,9 @@ import Control.Concurrent.STM (atomically) import Control.Concurrent.STM.TBQueue (newTBQueueIO, readTBQueue, writeTBQueue) import Control.Exception (evaluate) -import Control.Monad (filterM, forM, forM_, when) +import Control.Monad (filterM, forM, forM_, when, unless) import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as C8 import Data.IORef import Data.Int (Int16, Int32, Int64, Int8) import qualified Data.Map as M @@ -75,6 +76,7 @@ import System.Directory (doesDirectoryExist, removeFile) import System.FilePath (()) import System.FilePath.Glob (glob) +import System.IO (IOMode (ReadMode), hIsEOF, withFile) import System.IO.Temp (emptySystemTempFile) import Type.Reflection (typeRep) @@ -141,6 +143,7 @@ -} isBounded :: PhysicalPlan -> Bool isBounded (PhysicalScan (CsvSource{}) _) = True +isBounded (PhysicalScan (CsvSourceStreaming{}) _) = False isBounded (PhysicalScan (ParquetSource _) _) = True isBounded (PhysicalProject _ c) = isBounded c isBounded (PhysicalFilter _ c) = isBounded c @@ -187,6 +190,8 @@ -- Scan ----------------------------------------------------------------------- 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 (ParquetSource path) cfg) = executeParquetScan path cfg buildStream (PhysicalSpill child path) = do @@ -610,6 +615,47 @@ in return (Just df') ) +executeCsvScanStreaming :: FilePath -> CsvReader -> ScanConfig -> IO Stream +executeCsvScanStreaming path reader cfg = do + let schema = scanSchema 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 schema 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 + return . Stream $ do + mb <- atomically (readTBQueue queue) + case mb of + Nothing -> atomically (writeTBQueue queue Nothing) >> return Nothing + 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 + -- | Slice a 'DataFrame' into row-bounded batches of at most @n@ rows. sliceIntoBatches :: Int -> D.DataFrame -> [D.DataFrame] sliceIntoBatches n df = @@ -656,15 +702,26 @@ {- | Route join to the existing Operations.Join implementation. When the left and right key names differ, rename the right key before joining. + +'Join.join' retains its first 'DataFrame' argument and makes the second one +optional, so the lazy left sub-query ('leftDf') must be passed first for LEFT +and RIGHT joins to retain the side the caller means. INNER and FULL_OUTER are +symmetric in which rows survive, and 'Operations.Join' orders their output +columns with the renamed right frame first, so they keep the @rightDf leftDf@ +order. (Passing @rightDf@ first for LEFT/RIGHT was the source of a silent +left/right inversion: a left join dropped the left side's unmatched rows.) -} performJoin :: Join.JoinType -> T.Text -> T.Text -> D.DataFrame -> D.DataFrame -> D.DataFrame performJoin jt leftKey rightKey leftDf rightDf = - if leftKey == rightKey - then Join.join jt [leftKey] rightDf leftDf - else - let rightRenamed = Core.rename rightKey leftKey rightDf - in Join.join jt [leftKey] rightRenamed leftDf + case jt of + Join.LEFT -> Join.join jt [leftKey] leftDf rightRenamed + Join.RIGHT -> Join.join jt [leftKey] leftDf rightRenamed + _ -> Join.join jt [leftKey] rightRenamed leftDf + where + rightRenamed + | leftKey == rightKey = rightDf + | otherwise = Core.rename rightKey leftKey rightDf -- --------------------------------------------------------------------------- -- Sort order conversion 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 @@ -13,11 +13,14 @@ data DataSource = -- | path, separator, CSV reader (e.g. attoparsec or SIMD) CsvSource FilePath Char CsvReader + | CsvSourceStreaming FilePath Char CsvReader | ParquetSource FilePath instance Show DataSource where show (CsvSource path sep _) = "CsvSource " ++ show path ++ " " ++ show sep ++ " " + show (CsvSourceStreaming path sep _) = + "CsvSourceStreaming " ++ 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 @@ -180,6 +180,14 @@ PhysicalScan (CsvSource path sep reader) (ScanConfig batchSz sep schema Nothing) +toPhysical batchSz (Filter p (Scan (CsvSourceStreaming path sep reader) schema)) = + PhysicalScan + (CsvSourceStreaming path sep reader) + (ScanConfig batchSz sep schema (Just p)) +toPhysical batchSz (Scan (CsvSourceStreaming path sep reader) schema) = + PhysicalScan + (CsvSourceStreaming path sep reader) + (ScanConfig batchSz sep schema Nothing) toPhysical batchSz (Filter p (Scan (ParquetSource path) schema)) = PhysicalScan (ParquetSource path) diff --git a/src/DataFrame/Typed/Lazy.hs b/src/DataFrame/Typed/Lazy.hs --- a/src/DataFrame/Typed/Lazy.hs +++ b/src/DataFrame/Typed/Lazy.hs @@ -58,7 +58,10 @@ aggregate, -- * Joins - join, + innerJoin, + leftJoin, + rightJoin, + fullOuterJoin, -- * Sort sortBy, @@ -84,7 +87,7 @@ import DataFrame.Lazy.Internal.DataFrame (LazyDataFrame) import qualified DataFrame.Lazy.Internal.DataFrame as L import DataFrame.Lazy.Internal.LogicalPlan (SortOrder (..)) -import DataFrame.Operations.Join (JoinType) +import DataFrame.Operations.Join (JoinType (..)) import DataFrame.Typed.Expr import DataFrame.Typed.Freeze (unsafeFreeze) import DataFrame.Typed.Schema @@ -174,16 +177,73 @@ aggregate tagg (TLG (keys, ldf)) = TLD (L.groupBy keys (aggToNamedExprs tagg) ldf) --- | Join two lazy queries on a shared key column. -join :: +-- | Typed inner join on a single key column present in both schemas. +innerJoin :: + forall (key :: Symbol) left right. + ( KnownSymbol key + , AssertAllPresent '[key] left + , AssertAllPresent '[key] right + , AssertKeyTypesMatch '[key] left right + ) => + TypedLazyDataFrame left -> + TypedLazyDataFrame right -> + TypedLazyDataFrame (InnerJoinSchema '[key] left right) +innerJoin = joinOn @key INNER + +-- | Typed left join. The right table's non-key columns become @Maybe@. +leftJoin :: + forall (key :: Symbol) left right. + ( KnownSymbol key + , AssertAllPresent '[key] left + , AssertAllPresent '[key] right + , AssertKeyTypesMatch '[key] left right + ) => + TypedLazyDataFrame left -> + TypedLazyDataFrame right -> + TypedLazyDataFrame (LeftJoinSchema '[key] left right) +leftJoin = joinOn @key LEFT + +-- | Typed right join. The left table's non-key columns become @Maybe@. +rightJoin :: + forall (key :: Symbol) left right. + ( KnownSymbol key + , AssertAllPresent '[key] left + , AssertAllPresent '[key] right + , AssertKeyTypesMatch '[key] left right + ) => + TypedLazyDataFrame left -> + TypedLazyDataFrame right -> + TypedLazyDataFrame (RightJoinSchema '[key] left right) +rightJoin = joinOn @key RIGHT + +-- | Typed full outer join. Non-key columns from both tables become @Maybe@. +fullOuterJoin :: + forall (key :: Symbol) left right. + ( KnownSymbol key + , AssertAllPresent '[key] left + , AssertAllPresent '[key] right + , AssertKeyTypesMatch '[key] left right + ) => + TypedLazyDataFrame left -> + TypedLazyDataFrame right -> + TypedLazyDataFrame (FullOuterJoinSchema '[key] left right) +fullOuterJoin = joinOn @key FULL_OUTER + +{- | Runtime delegation shared by the typed joins. The lazy backend joins on a +single key whose name is the same in both schemas; the result schema is +computed by the caller's join-specific type family. +-} +joinOn :: + forall (key :: Symbol) left right out. + (KnownSymbol key) => JoinType -> - T.Text -> - T.Text -> TypedLazyDataFrame left -> TypedLazyDataFrame right -> - TypedLazyDataFrame left -- TODO: compute join result schema -join jt leftKey rightKey (TLD left) (TLD right) = - TLD (L.join jt leftKey rightKey left right) + TypedLazyDataFrame out +joinOn jt (TLD left) (TLD right) = + TLD (L.join jt keyName keyName left right) + where + keyName = T.pack (symbolVal (Proxy @key)) -- | Sort the result by column name and direction. sortBy ::