dataframe-lazy 1.1.0.0 → 1.1.0.2

raw patch · 6 files changed

+151/−16 lines, 6 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• dataframe-lazy.cabal +1/−1
• src/DataFrame/Lazy/Internal/DataFrame.hs +7/−0
• src/DataFrame/Lazy/Internal/Executor.hs +63/−6
• src/DataFrame/Lazy/Internal/LogicalPlan.hs +3/−0
• src/DataFrame/Lazy/Internal/Optimizer.hs +8/−0
• src/DataFrame/Typed/Lazy.hs +69/−9

@@ -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,

@@ -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

@@ -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

@@ -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 ++ " <reader>"+    show (CsvSourceStreaming path sep _) =+        "CsvSourceStreaming " ++ show path ++ " " ++ show sep ++ " <reader>"     show (ParquetSource path) = "ParquetSource " ++ show path  -- | Sort direction used in Sort nodes and the public API.

@@ -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)

@@ -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 ::