dataframe-lazy 2.4.2.0 → 2.4.2.1
raw patch · 10 files changed
+97/−411 lines, 10 filesdep ~dataframe-coredep ~dataframe-operationsdep ~dataframe-parquet
Dependency ranges changed: dataframe-core, dataframe-operations, dataframe-parquet
Files
- dataframe-lazy.cabal +5/−5
- src/DataFrame/Lazy/IO/Binary.hs +6/−56
- src/DataFrame/Lazy/IO/CSV.hs +11/−40
- src/DataFrame/Lazy/Internal/DataFrame.hs +12/−97
- src/DataFrame/Lazy/Internal/Executor.hs +37/−116
- src/DataFrame/Lazy/Internal/LogicalPlan.hs +3/−5
- src/DataFrame/Lazy/Internal/Optimizer.hs +11/−37
- src/DataFrame/Lazy/Internal/PhysicalPlan.hs +3/−5
- src/DataFrame/Typed/Lazy.hs +4/−31
- test/StreamingCsv.hs +5/−19
dataframe-lazy.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.4 name: dataframe-lazy-version: 2.4.2.0+version: 2.4.2.1 synopsis: Lazy query engine for the dataframe ecosystem. description: The lazy/streaming query engine: relational-algebra plans, optimizer,@@ -45,10 +45,10 @@ attoparsec >= 0.12 && < 0.16, bytestring >= 0.11 && < 0.14, containers >= 0.6.7 && < 0.10,- dataframe-core >= 2.5 && < 2.6,+ dataframe-core >= 2.6 && < 2.7, dataframe-csv >= 2.3 && < 2.4,- dataframe-operations >= 2.5 && < 2.6,- dataframe-parquet >= 1.5.0.1 && < 1.6,+ dataframe-operations >= 2.6 && < 2.7,+ dataframe-parquet >= 1.6 && < 1.7, dataframe-parsing >= 2.2 && < 2.3, directory >= 1.3.0.0 && < 2, filepath >= 1.4 && < 2,@@ -68,7 +68,7 @@ build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.14, containers >= 0.6.7 && < 0.10,- dataframe-core >= 2.5 && < 2.6,+ dataframe-core >= 2.6 && < 2.7, dataframe-csv >= 2.3 && < 2.4, dataframe-lazy, dataframe-parsing >= 2.2 && < 2.3,
src/DataFrame/Lazy/IO/Binary.hs view
@@ -5,31 +5,7 @@ {-# LANGUAGE StrictData #-} {-# LANGUAGE TypeApplications #-} -{- | Simple column-oriented binary spill format (DFBN).--Layout (all integers little-endian):--@-[magic: 4 bytes] "DFBN"-[num_columns: 4 bytes] Word32- per column:- [name_len: 2 bytes] Word16 (byte length of UTF-8 name)- [name: name_len bytes]- [type_tag: 1 byte] Word8-[num_rows: 8 bytes] Word64--per column data block (order matches schema):- type_tag 0 (Int): num_rows × Int64 LE- type_tag 1 (Double): num_rows × Double LE (IEEE 754)- type_tag 2 (Text): (num_rows+1) × Word32 offsets ++ payload bytes (UTF-8)- type_tag 3 (Maybe Int): ceil(num_rows/8)-byte null bitmap ++ num_rows × Int64 LE- type_tag 4 (Maybe Double): ceil(num_rows/8)-byte null bitmap ++ num_rows × Double LE- type_tag 5 (Maybe Text): ceil(num_rows/8)-byte null bitmap- ++ (num_rows+1) × Word32 offsets ++ payload bytes-@--Null bitmap: bit @i@ of byte @i\/8@ is 1 when row @i@ is non-null.--}+-- | Simple column-oriented binary spill format (DFBN). module DataFrame.Lazy.IO.Binary ( spillToDisk, readSpilled,@@ -66,10 +42,6 @@ import System.IO (IOMode (..), hClose, openTempFile, withFile) import Type.Reflection (typeRep) --- ------------------------------------------------------------------------------ Type tags--- ---------------------------------------------------------------------------- tagInt, tagDouble, tagText, tagMaybeInt, tagMaybeDouble, tagMaybeText :: Word8 tagInt = 0 tagDouble = 1@@ -78,10 +50,6 @@ tagMaybeDouble = 4 tagMaybeText = 5 --- ------------------------------------------------------------------------------ Write--- ---------------------------------------------------------------------------- -- | Serialise a 'DataFrame' to a DFBN binary file. spillToDisk :: FilePath -> DataFrame -> IO () spillToDisk path df =@@ -165,9 +133,7 @@ in buildNullBitmap isValidVec <> buildTextVector texts buildColumnData n c@(PackedText _ _) = buildColumnData n (materializePacked c) -{- | Bulk-encode an Int vector as 8-byte LE values (native layout on LE platforms).-hPutBuilder flushes synchronously so the underlying ForeignPtr outlives the Builder.--}+-- | Bulk-encode an Int vector as 8-byte little-endian values. buildIntVector :: VU.Vector Int -> BSB.Builder buildIntVector v = let sv = VU.convert v :: VS.Vector Int@@ -175,7 +141,7 @@ bs = BSI.fromForeignPtr (castForeignPtr fp) 0 (n * sizeOf (0 :: Int)) in BSB.byteString bs --- | Bulk-encode a Double vector as 8-byte LE IEEE 754 values (native layout on LE platforms).+-- | Bulk-encode a Double vector as 8-byte LE IEEE 754 values (native on LE). buildDoubleVector :: VU.Vector Double -> BSB.Builder buildDoubleVector v = let sv = VU.convert v :: VS.Vector Double@@ -206,10 +172,6 @@ (0 :: Word8) [0 .. 7] --- ------------------------------------------------------------------------------ Read--- ---------------------------------------------------------------------------- -- | @(new_offset, value)@ type ParseResult a = Either String (Int, a) @@ -285,10 +247,7 @@ return (off2, BoxedColumn (Just bm) v) | otherwise = Left ("unknown type tag " <> show tag) -{- | Zero-copy Int column read: reuses the ByteString buffer's ForeignPtr.-Safe as long as 'bs' stays live during the caller's use of the resulting vector.-Only correct on little-endian platforms (aarch64/x86_64).--}+-- | Zero-copy Int column read that reuses the ByteString's buffer. readIntColumn :: BS.ByteString -> Int -> Int -> ParseResult (VU.Vector Int) readIntColumn bs off nrows | off + nrows * 8 > BS.length bs = Left "unexpected end of input"@@ -298,10 +257,7 @@ sv = VS.unsafeFromForeignPtr0 fp' nrows :: VS.Vector Int in Right (off + nrows * 8, VU.convert sv) -{- | Zero-copy Double column read: reuses the ByteString buffer's ForeignPtr.-Safe as long as 'bs' stays live during the caller's use of the resulting vector.-Only correct on little-endian platforms (aarch64/x86_64).--}+-- | Zero-copy Double column read that reuses the ByteString's buffer. readDoubleColumn :: BS.ByteString -> Int -> Int -> ParseResult (VU.Vector Double) readDoubleColumn bs off nrows@@ -398,13 +354,7 @@ (off'', xs) <- readN (n - 1) f off' return (off'', x : xs) --- ------------------------------------------------------------------------------ Bracket helper--- -----------------------------------------------------------------------------{- | Spill a DataFrame to a temporary file, run an action with the path,-then delete the file even if the action throws.--}+-- | Spill a DataFrame to a temp file, run an action, then delete the file. withSpilled :: DataFrame -> (FilePath -> IO a) -> IO a withSpilled df action = do tmpDir <- getTemporaryDirectory
src/DataFrame/Lazy/IO/CSV.hs view
@@ -47,9 +47,7 @@ { hasHeader :: Bool , inferTypes :: Bool , safeRead :: SafeReadMode- {- ^ Default 'SafeReadMode' for columns without an entry in- 'safeReadOverrides'.- -}+ -- ^ Default 'SafeReadMode' for columns without an override. , safeReadOverrides :: [(T.Text, SafeReadMode)] -- ^ Per-column 'SafeReadMode' overrides; takes precedence over 'safeRead'. , rowRange :: !(Maybe (Int, Int)) -- (start, length)@@ -59,10 +57,7 @@ , rowsRead :: !Int } -{- | Default read options: assume a header, infer types, no safe-read wrapping.-Set 'safeRead' to 'MaybeRead'/'EitherRead' to wrap columns; use-'safeReadOverrides' to pick a different mode per column.--}+-- | Default read options: assume a header, infer types, no safe-read wrapping. defaultOptions :: ReadOptions defaultOptions = ReadOptions@@ -77,17 +72,11 @@ , rowsRead = 0 } -{- | Reads a CSV file from the given path.-Note this file stores intermediate temporary files-while converting the CSV from a row to a columnar format.--}+-- | Read a CSV file from a path, staging intermediate temp files. readCsv :: FilePath -> IO DataFrame readCsv path = fst <$> readSeparated ',' defaultOptions path -{- | Reads a tab separated file from the given path.-Note this file stores intermediate temporary files-while converting the CSV from a row to a columnar format.--}+-- | Read a tab-separated file from a path, staging intermediate temp files. readTsv :: FilePath -> IO DataFrame readTsv path = fst <$> readSeparated '\t' defaultOptions path @@ -202,7 +191,7 @@ pure (l, r) {-# INLINE fillColumns #-} --- | Writes a value into the appropriate column, resizing the vector if necessary.+-- | Writes a value into the right column, resizing the vector if needed. writeValue :: VM.IOVector MutableColumn -> VM.IOVector [(Int, T.Text)] ->@@ -217,7 +206,7 @@ either modify (const (return ())) res {-# INLINE writeValue #-} --- | Freezes a mutable vector into an immutable one, trimming it to the actual row count.+-- | Freezes a mutable vector into an immutable one, trimmed to the row count. freezeColumn :: V.Vector T.Text -> VM.IOVector MutableColumn ->@@ -241,14 +230,7 @@ NoSafeRead -> freezeColumn' colNulls col {-# INLINE freezeColumn #-} --- ------------------------------------------------------------------------------ Streaming scan API--- -----------------------------------------------------------------------------{- | Open a file for streaming: returns a handle positioned after the header-and the column spec for the schema columns present in the header.-The caller must close the handle when done.--}+-- | Open a file for streaming, returning the handle and the column spec. openCsvStream :: Char -> Schema ->@@ -272,10 +254,7 @@ ("openCsvStream: none of the schema columns appear in the header of " <> path) return (handle, colSpec) -{- | Read up to @batchSz@ rows, returning a batch 'DataFrame' and the unconsumed-leftover text; 'Nothing' at EOF with no leftover. Pass the previous call's-leftover back in (use @""@ on the first call).--}+-- | Read up to @batchSz@ rows, returning a batch and the leftover text. readBatch :: Char -> [(Int, T.Text, SchemaType)] ->@@ -342,9 +321,7 @@ , newLeftover ) -{- | Write a 'ByteString' field value directly into a mutable column,-parsing numerics without an intermediate 'T.Text' allocation.--}+-- | Write a 'ByteString' field directly into a mutable column. writeColumnBs :: Int -> BS.ByteString -> MutableColumn -> IO (Either T.Text Bool) writeColumnBs i bs (MBoxedColumn (col :: VM.IOVector a)) =@@ -365,10 +342,7 @@ Nothing -> return (Left (TextEncoding.decodeUtf8Lenient bs)) {-# INLINE writeColumnBs #-} -{- | Extracts the Nth field (0-indexed), respecting double quotes and stripping them.-Fast path: uses memchr-based 'BS.break' when no quotes are present in the line.-Slow path: quote-aware character-by-character scan.--}+-- | Extract the Nth field (0-indexed), respecting and stripping quotes. getNthFieldBs :: Word8 -> Int -> BS.ByteString -> BS.ByteString getNthFieldBs sep targetIdx bs | not (BS.any (== 0x22) bs) = skipFast targetIdx bs@@ -422,10 +396,7 @@ sliceCol n (MBoxedColumn col) = MBoxedColumn (VM.take n col) sliceCol n (MUnboxedColumn col) = MUnboxedColumn (VUM.take n col) -{- | Finds the index of the next unquoted newline (0x0A).-Fast path: uses memchr (SIMD) and falls back to a quote-aware linear scan-only if a double-quote appears before the candidate newline.--}+-- | Find the index of the next unquoted newline (0x0A). findUnquotedNewline :: BS.ByteString -> Maybe Int findUnquotedNewline bs = case BS.elemIndex 0x0A bs of
src/DataFrame/Lazy/Internal/DataFrame.hs view
@@ -27,9 +27,7 @@ import DataFrame.Operations.Join (JoinType) import DataFrame.Schema (Schema) -{- | A lazy query that has not been executed yet: a 'LogicalPlan' tree whose-execution is deferred until 'runDataFrame' is called.--}+-- | A lazy query: a 'LogicalPlan' whose execution waits for 'runDataFrame'. data LazyDataFrame = LazyDataFrame { plan :: LogicalPlan , batchSize :: Int@@ -40,52 +38,19 @@ "LazyDataFrame { batchSize = " <> (show (batchSize ldf) <> (", plan = " <> (show (plan ldf) <> " }"))) --- ------------------------------------------------------------------------------ Entry point--- -----------------------------------------------------------------------------{- | Execute the lazy query: optimise the logical plan, then stream-execute-the resulting physical plan into a fully-materialised 'D.DataFrame'.--}+-- | Optimise and stream-execute the lazy query into a 'D.DataFrame'. runDataFrame :: LazyDataFrame -> IO D.DataFrame runDataFrame ldf = execute (Opt.optimize (batchSize ldf) (plan ldf)) --- ------------------------------------------------------------------------------ Builders that construct the logical plan tree--- ---------------------------------------------------------------------------- -- | Lift an already-loaded eager 'D.DataFrame' into the lazy plan. fromDataFrame :: D.DataFrame -> LazyDataFrame fromDataFrame df = LazyDataFrame{plan = SourceDF df, batchSize = 1_000_000} -{- | Scan a CSV file with the default comma separator and the in-tree-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\")--@--}+-- | Scan a CSV file with the default comma separator and strict reader. scanCsv :: Schema -> T.Text -> LazyDataFrame scanCsv = scanCsvWith readSeparated -{- | Like 'scanCsv' but with an explicit CSV reader (e.g. the SIMD reader-@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\")--@--}+-- | Like 'scanCsv' but with an explicit 'CsvReader'. scanCsvWith :: CsvReader -> Schema -> T.Text -> LazyDataFrame scanCsvWith reader schema path = LazyDataFrame@@ -93,35 +58,11 @@ , batchSize = 1_000_000 } -{- | 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")--@--}+-- | Scan a CSV in bounded-memory windows with the default strict reader. 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\")--@--}+-- | Like 'scanCsvWith', but reads the file in bounded-memory windows. scanCsvStreamingWith :: CsvReader -> Schema -> T.Text -> LazyDataFrame scanCsvStreamingWith reader schema path = LazyDataFrame@@ -129,16 +70,7 @@ , 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")--@--}+-- | Stream a CSV in bounded-memory windows decoded by a 'CsvBytesReader'. scanCsvStreamingBytesWith :: CsvBytesReader -> Schema -> T.Text -> LazyDataFrame scanCsvStreamingBytesWith reader schema path = LazyDataFrame@@ -146,25 +78,11 @@ , batchSize = 1_000_000 } -{- | Scan a character-separated file with the default strict reader.--==== __Example__-@-ghci> L.runDataFrame (L.scanSeparated ';' schema \"customers.txt\")--@--}+-- | Scan a character-separated file with the default strict reader. scanSeparated :: Char -> Schema -> T.Text -> LazyDataFrame scanSeparated = scanSeparatedWith readSeparated -{- | Like 'scanSeparated' but with an explicit CSV reader.--==== __Example__-@-ghci> L.runDataFrame (L.scanSeparatedWith Fast.fastReadCsvWithOpts ';' schema \"customers.txt\")--@--}+-- | Like 'scanSeparated' but with an explicit CSV reader. scanSeparatedWith :: CsvReader -> Char -> Schema -> T.Text -> LazyDataFrame scanSeparatedWith reader sep schema path =@@ -185,7 +103,7 @@ derive :: (C.Columnable a) => T.Text -> E.Expr a -> LazyDataFrame -> LazyDataFrame derive name expr ldf =- ldf{plan = Derive name (E.UExpr expr) (plan ldf)}+ ldf{plan = Derive name (E.SomeExpr expr) (plan ldf)} -- | Retain only the listed columns. select :: [T.Text] -> LazyDataFrame -> LazyDataFrame@@ -213,15 +131,12 @@ , batchSize = batchSize left } -{- | Group by a set of columns and compute aggregate expressions.--Each aggregate expression should use an 'Agg' node (e.g. @sumOf@, @meanOf@).--}+-- | Group by a set of columns and compute aggregate expressions. groupBy :: -- | Group-by key columns [T.Text] -> -- | @[(outputName, aggregateExpr)]@- [(T.Text, E.UExpr)] ->+ [(T.Text, E.SomeExpr)] -> LazyDataFrame -> LazyDataFrame groupBy keys aggs ldf = ldf{plan = Aggregate keys aggs (plan ldf)}
src/DataFrame/Lazy/Internal/Executor.hs view
@@ -8,10 +8,7 @@ {-# LANGUAGE TupleSections #-} {-# LANGUAGE TypeApplications #-} -{- | Pull-based (iterator) execution engine: each operator returns a 'Stream'-yielding the next 'DataFrame' batch or 'Nothing' at end. Bounded local sources-materialise into eager whole-frame ops; unbounded sources stream in constant memory.--}+-- | Pull-based execution engine where each operator returns a 'Stream'. module DataFrame.Lazy.Internal.Executor ( CsvReader, execute,@@ -66,19 +63,10 @@ import System.IO.Temp (emptySystemTempFile) import Type.Reflection (typeRep) --- ------------------------------------------------------------------------------ Stream abstraction--- -----------------------------------------------------------------------------{- | A pull-based stream: each call to the action yields the next batch or-'Nothing' when the stream is exhausted. State is captured by the closure.--}+-- | A pull-based stream whose action yields the next batch or 'Nothing'. newtype Stream = Stream {pullBatch :: IO (Maybe D.DataFrame)} -{- | Wrap an already-computed 'DataFrame' as a single-batch stream: the first-pull yields it, every subsequent pull yields 'Nothing'. Used by blocking-operators that materialise their whole result up front.--}+-- | Wrap a computed 'DataFrame' as a single-batch stream. materialized :: D.DataFrame -> IO Stream materialized df = do ref <- newIORef (Just df)@@ -87,10 +75,7 @@ writeIORef ref Nothing return mb -{- | Drain all batches from a stream and concatenate them into one DataFrame.-Columns are concatenated in a single multi-way pass (O(rows)) rather than a-left-fold of @acc <> batch@ that would recopy the accumulator on every step.--}+-- | Drain a stream and concatenate its batches into one DataFrame. collectStream :: Stream -> IO D.DataFrame collectStream stream = go [] where@@ -105,19 +90,12 @@ concatBatches [] = D.empty concatBatches [df] = df concatBatches batches@(first : _) =- D.fromNamedColumns+ D.fromColumns [ (name, C.concatManyColumns [D.unsafeGetColumn name b | b <- batches]) | name <- D.columnNames first ] --- ------------------------------------------------------------------------------ Boundedness analysis--- -----------------------------------------------------------------------------{- | True when every scan leaf is a finite local source, so the whole result can-be materialised and routed through the eager whole-frame ops. False for-online/streaming sources, which keep the constant-memory streaming paths.--}+-- | True when every scan leaf is a finite local source. isBounded :: PhysicalPlan -> Bool isBounded (PhysicalScan (CsvSource{}) _) = True isBounded (PhysicalScan (CsvSourceStreaming{}) _) = False@@ -134,19 +112,11 @@ isBounded (PhysicalHashJoin _ _ _ l r) = isBounded l && isBounded r isBounded (PhysicalSortMergeJoin _ _ _ l r) = isBounded l && isBounded r --- ------------------------------------------------------------------------------ Top-level entry point--- -----------------------------------------------------------------------------{- | Execute a physical plan, returning the complete result as a single-'DataFrame'.--}+-- | Execute a physical plan into a single complete 'DataFrame'. execute :: PhysicalPlan -> IO D.DataFrame execute plan = buildStream plan >>= collectStream -{- | Fold a function over every batch produced by a physical plan.-The fold is strict in the accumulator; each batch is discarded after folding.--}+-- | Strictly fold a function over every batch a physical plan produces. foldBatches :: (b -> D.DataFrame -> IO b) -> b -> PhysicalPlan -> IO b foldBatches f seed plan = do@@ -160,10 +130,6 @@ loop acc' loop seed --- ------------------------------------------------------------------------------ Per-operator stream builders--- ---------------------------------------------------------------------------- buildStream :: PhysicalPlan -> IO Stream buildStream (PhysicalScan (CsvSource path sep reader) cfg) = executeCsvScan path sep reader cfg@@ -316,13 +282,7 @@ rightDf <- execute rightPlan materialized (performJoin jt leftKey rightKey leftDf rightDf) --- ------------------------------------------------------------------------------ Streaming aggregation helpers--- -----------------------------------------------------------------------------{- | One worker's loop: pull batches off the shared child stream until-exhausted, building up a per-worker accumulator.--}+-- | One worker's loop: pull shared batches into a per-worker accumulator. workerLoop :: Stream -> [T.Text] ->@@ -364,30 +324,27 @@ Agg.aggregate mergeAggs (Agg.groupBy keys (acc <> p)) go merged ps -isStreamableAgg :: E.UExpr -> Bool-isStreamableAgg (E.UExpr (E.Agg (E.CollectAgg _ _) _)) = False-isStreamableAgg (E.UExpr (E.Agg (E.FoldAgg _ Nothing (_ :: a -> b -> a)) _)) =+isStreamableAgg :: E.SomeExpr -> Bool+isStreamableAgg (E.SomeExpr (E.Agg (E.CollectAgg _ _) _)) = False+isStreamableAgg (E.SomeExpr (E.Agg (E.FoldAgg _ Nothing (_ :: a -> b -> a)) _)) = case testEquality (typeRep @a) (typeRep @b) of Just Refl -> True Nothing -> False-isStreamableAgg (E.UExpr (E.Agg (E.FoldAgg _ (Just _) (_ :: a -> b -> a)) _)) =+isStreamableAgg (E.SomeExpr (E.Agg (E.FoldAgg _ (Just _) (_ :: a -> b -> a)) _)) = case testEquality (typeRep @a) (typeRep @Int) of Just Refl -> True Nothing -> case testEquality (typeRep @a) (typeRep @b) of Just Refl -> True Nothing -> False-isStreamableAgg (E.UExpr (E.Agg (E.MergeAgg{}) _)) = True+isStreamableAgg (E.SomeExpr (E.Agg (E.MergeAgg{}) _)) = True isStreamableAgg _ = False -{- | Build the (partial, merge, finalizer) plan for a list of streamable-aggregates: @partialAggs@ run per batch, @mergeAggs@ combine two partial-results, and @finalizer@ post-processes (for 'MergeAgg' acc≠output types).--}+-- | Build the (partial, merge, finalizer) plan for streamable aggregates. buildAggPlan ::- [(T.Text, E.UExpr)] ->- ( [(T.Text, E.UExpr)]- , [(T.Text, E.UExpr)]+ [(T.Text, E.SomeExpr)] ->+ ( [(T.Text, E.SomeExpr)]+ , [(T.Text, E.SomeExpr)] , D.DataFrame -> D.DataFrame ) buildAggPlan aggs = foldl combine ([], [], id) (map processAgg aggs)@@ -395,14 +352,14 @@ combine (p1, m1, f1) (p2, m2, f2) = (p1 ++ p2, m1 ++ m2, f1 . f2) processAgg ::- (T.Text, E.UExpr) ->- ([(T.Text, E.UExpr)], [(T.Text, E.UExpr)], D.DataFrame -> D.DataFrame)+ (T.Text, E.SomeExpr) ->+ ([(T.Text, E.SomeExpr)], [(T.Text, E.SomeExpr)], D.DataFrame -> D.DataFrame) processAgg (name, ue) = case ue of- E.UExpr (E.Agg (E.FoldAgg n Nothing (f :: a -> b -> a)) (_ :: E.Expr b)) ->+ E.SomeExpr (E.Agg (E.FoldAgg n Nothing (f :: a -> b -> a)) (_ :: E.Expr b)) -> case testEquality (typeRep @a) (typeRep @b) of Just Refl -> ( [(name, ue)]- , [(name, E.UExpr (E.Agg (E.FoldAgg n Nothing f) (E.Col @a name)))]+ , [(name, E.SomeExpr (E.Agg (E.FoldAgg n Nothing f) (E.Col @a name)))] , id ) Nothing ->@@ -412,21 +369,21 @@ , [ ( name- , E.UExpr+ , E.SomeExpr (E.Agg (E.FoldAgg "sum" Nothing ((+) :: Int -> Int -> Int)) (E.Col @Int name)) ) ] , id ) Nothing -> ([(name, ue)], [(name, ue)], id)- E.UExpr (E.Agg (E.FoldAgg n (Just _) (f :: a -> b -> a)) (_ :: E.Expr b)) ->+ E.SomeExpr (E.Agg (E.FoldAgg n (Just _) (f :: a -> b -> a)) (_ :: E.Expr b)) -> case testEquality (typeRep @a) (typeRep @Int) of Just Refl -> ( [(name, ue)] , [ ( name- , E.UExpr+ , E.SomeExpr (E.Agg (E.FoldAgg "sum" Nothing ((+) :: Int -> Int -> Int)) (E.Col @Int name)) ) ]@@ -436,11 +393,11 @@ case testEquality (typeRep @a) (typeRep @b) of Just Refl -> ( [(name, ue)]- , [(name, E.UExpr (E.Agg (E.FoldAgg n Nothing f) (E.Col @a name)))]+ , [(name, E.SomeExpr (E.Agg (E.FoldAgg n Nothing f) (E.Col @a name)))] , id ) Nothing -> ([(name, ue)], [(name, ue)], id)- E.UExpr+ E.SomeExpr ( E.Agg ( E.MergeAgg n@@ -452,13 +409,13 @@ (inner :: E.Expr b) ) -> let partialExpr =- E.UExpr+ E.SomeExpr ( E.Agg (E.FoldAgg ("partial_" <> n) (Just seed) step) inner ) mergeExpr =- E.UExpr+ E.SomeExpr ( E.Agg (E.FoldAgg ("merge_" <> n) Nothing merge) (E.Col @acc name)@@ -477,14 +434,7 @@ ) _ -> ([(name, ue)], [(name, ue)], id) --- ------------------------------------------------------------------------------ Parquet scan implementation--- -----------------------------------------------------------------------------{- | Scan a Parquet file, directory, or glob. Each file becomes one batch.-Column projection and predicate pushdown are forwarded to 'readParquetWithOpts'-via 'ParquetReadOptions'.--}+-- | Scan a Parquet file, directory or glob; each file becomes one batch. executeParquetScan :: FilePath -> ScanConfig -> IO Stream executeParquetScan path cfg = do isDir <- doesDirectoryExist path@@ -507,21 +457,12 @@ writeIORef ref rest Just <$> Parquet.readParquetWithOpts opts f --- ------------------------------------------------------------------------------ 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.--}+-- | Read options for a CSV scan: schema gives types, source 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.--}+-- | Parallel CSV scan splitting the file at newlines into slices. executeCsvScan :: FilePath -> Char -> CsvReader -> ScanConfig -> IO Stream executeCsvScan path sep reader cfg = do nCaps <- getNumCapabilities@@ -623,12 +564,7 @@ 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.--}+-- | Find the last record-ending LF in a CSV buffer. lastCompleteRecordNewline :: BS.ByteString -> Maybe Int lastCompleteRecordNewline bytes | not (BS.elem quote bytes) = BS.elemIndexEnd nl bytes@@ -636,9 +572,7 @@ 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.+ -- Each read stays inside its bounds-checked branch (no shared guard). go !i !inside !lastNewline | i >= len = lastNewline | otherwise = case BSU.unsafeIndex bytes i of@@ -667,10 +601,7 @@ starts = [0, n .. total - 1] in [Sub.range (s, min (s + n) total) df | s <- starts] -{- | Split a CSV file at newline boundaries into @n@ temp files, each with the-original header plus a body slice. Returns the paths (caller removes them); the-per-file reader mmap's each, giving OS-paged reads instead of one monolithic read.--}+-- | Split a CSV at newline boundaries into temp files, each with the header. splitCsvAtNewlines :: Int -> FilePath -> IO [FilePath] splitCsvAtNewlines n path = do bs <- BS.readFile path@@ -697,10 +628,7 @@ Just i -> off + i + 1 Nothing -> BS.length body -{- | Route a join to 'Operations.Join', renaming the right key when the names-differ. 'Join.join' keeps its first argument, so LEFT/RIGHT must pass 'leftDf'-first to retain the intended side; symmetric INNER/FULL_OUTER pass @rightDf@ first.--}+-- | Route a join to 'Operations.Join', renaming the right key if needed. performJoin :: Join.JoinType -> T.Text -> T.Text -> D.DataFrame -> D.DataFrame -> D.DataFrame performJoin jt leftKey rightKey leftDf rightDf =@@ -713,14 +641,7 @@ | leftKey == rightKey = rightDf | otherwise = Core.rename rightKey leftKey rightDf --- ------------------------------------------------------------------------------ Sort order conversion--- -----------------------------------------------------------------------------{- | Convert a plan-level @(column, direction)@ into a Permutation 'SortOrder',-emitting @E.Col@ at the materialised column's element type so 'Perm.sortBy'-dispatches its comparator correctly; unknown/non-'Ord' types fall back to 'T.Text'.--}+-- | Convert a plan-level (column, direction) into a Permutation 'SortOrder'. toPermSortOrder :: D.DataFrame -> (T.Text, SortOrder) -> Perm.SortOrder toPermSortOrder df (col, dir) = case M.lookup col (D.columnIndices df) of
src/DataFrame/Lazy/Internal/LogicalPlan.hs view
@@ -30,9 +30,7 @@ data SortOrder = Ascending | Descending deriving (Show, Eq, Ord) -{- | Relational-algebra tree that represents what the query computes.-No physical decisions (batch size, join strategy) are made here.--}+-- | Relational-algebra tree describing what the query computes. data LogicalPlan = -- | Read columns described by the schema from a source. Scan DataSource Schema@@ -41,11 +39,11 @@ | -- | Keep rows matching the predicate. Filter (E.Expr Bool) LogicalPlan | -- | Add or overwrite a column via an expression.- Derive T.Text E.UExpr LogicalPlan+ Derive T.Text E.SomeExpr LogicalPlan | -- | Join two sub-plans on the given key columns. Join JoinType T.Text T.Text LogicalPlan LogicalPlan | -- | Group then aggregate.- Aggregate [T.Text] [(T.Text, E.UExpr)] LogicalPlan+ Aggregate [T.Text] [(T.Text, E.SomeExpr)] LogicalPlan | -- | Sort by a list of (column, direction) pairs. Sort [(T.Text, SortOrder)] LogicalPlan | -- | Retain at most N rows.
src/DataFrame/Lazy/Internal/Optimizer.hs view
@@ -11,10 +11,7 @@ import DataFrame.Lazy.Internal.PhysicalPlan import DataFrame.Schema (Schema (..), elements) -{- | Optimise a logical plan and lower it to a physical plan: fuse filters, push-predicates toward the scan, drop dead derived columns, then let @toPhysical@-select concrete operators.--}+-- | Optimise a logical plan and lower it to a physical plan. optimize :: Int -> LogicalPlan -> PhysicalPlan optimize batchSz = toPhysical batchSz@@ -22,10 +19,6 @@ . pushPredicates . fuseFilters --- ------------------------------------------------------------------------------ Rule 1: Filter fusion--- ---------------------------------------------------------------------------- -- | Merge @Filter p1 (Filter p2 child)@ into @Filter (p1 && p2) child@. fuseFilters :: LogicalPlan -> LogicalPlan fuseFilters (Filter p1 (Filter p2 child)) =@@ -54,14 +47,7 @@ } ) --- ------------------------------------------------------------------------------ Rule 2: Predicate pushdown--- -----------------------------------------------------------------------------{- | Push Filter nodes as close to the Scan as possible: past a @Derive@ it-doesn't reference, past a @Project@ that keeps all its columns, and into the-@Scan@'s pushdown predicate.--}+-- | Push Filter nodes as close to the Scan as possible. pushPredicates :: LogicalPlan -> LogicalPlan pushPredicates (Filter p (Derive name expr child)) | name `notElem` E.getColumns p =@@ -84,36 +70,29 @@ pushPredicates (Limit n child) = Limit n (pushPredicates child) pushPredicates leaf = leaf --- ------------------------------------------------------------------------------ Rule 3: Dead column elimination--- -----------------------------------------------------------------------------{- | Collect every column name referenced anywhere in the plan (filters, sorts,-aggregate/join keys, projections, derived expressions). 'Nothing' means all-columns are needed (no Project restricts the output).--}+-- | Collect every column name referenced anywhere in the plan. referencedCols :: LogicalPlan -> Maybe (S.Set T.Text) referencedCols (Scan _ schema) = Just (S.fromList (M.keys (elements schema))) referencedCols (Project cols _) = Just (S.fromList cols) referencedCols (Filter p child) = fmap (S.union (S.fromList (E.getColumns p))) (referencedCols child) referencedCols (Derive _ expr child) =- fmap (S.union (S.fromList (uExprCols expr))) (referencedCols child)+ fmap (S.union (S.fromList (someExprCols expr))) (referencedCols child) referencedCols (Join _ l r left right) = let keySet = S.fromList [l, r] lRef = fmap (S.union keySet) (referencedCols left) rRef = fmap (S.union keySet) (referencedCols right) in S.union <$> lRef <*> rRef referencedCols (Aggregate keys aggs child) =- let aggCols = S.fromList (keys <> concatMap (uExprCols . snd) aggs)+ let aggCols = S.fromList (keys <> concatMap (someExprCols . snd) aggs) in fmap (S.union aggCols) (referencedCols child) referencedCols (Sort cols child) = fmap (S.union (S.fromList (fmap fst cols))) (referencedCols child) referencedCols (Limit _ child) = referencedCols child referencedCols (SourceDF _) = Nothing -uExprCols :: E.UExpr -> [T.Text]-uExprCols (E.UExpr expr) = E.getColumns expr+someExprCols :: E.SomeExpr -> [T.Text]+someExprCols (E.SomeExpr expr) = E.getColumns expr -- | Drop @Derive@ nodes whose output column is never consumed downstream. eliminateDeadColumns :: LogicalPlan -> LogicalPlan@@ -128,7 +107,7 @@ Derive name expr- (go (Just (S.union cols (S.fromList (uExprCols expr)))) child)+ (go (Just (S.union cols (S.fromList (someExprCols expr)))) child) go needed (Filter p child) = Filter p (go (fmap (S.union (S.fromList (E.getColumns p))) needed) child) go _needed (Project cols child) =@@ -137,7 +116,8 @@ let keySet = fmap (S.union (S.fromList [l, r])) needed in Join jt l r (go keySet left) (go keySet right) go needed (Aggregate keys aggs child) =- let aggCols = fmap (S.union (S.fromList (keys <> concatMap (uExprCols . snd) aggs))) needed+ let aggCols =+ fmap (S.union (S.fromList (keys <> concatMap (someExprCols . snd) aggs))) needed in Aggregate keys aggs (go aggCols child) go needed (Sort cols child) = Sort cols (go (fmap (S.union (S.fromList (fmap fst cols))) needed) child)@@ -149,13 +129,7 @@ Scan ds (Schema (M.filterWithKey (\k _ -> k `S.member` cols) (elements schema))) go _ (SourceDF df) = SourceDF df --- ------------------------------------------------------------------------------ Logical → Physical lowering--- -----------------------------------------------------------------------------{- | Lower the (already-optimised) logical plan to a physical plan. Joins always-lower to HashJoin; the executor may fall back to SortMerge at runtime.--}+-- | Lower the optimised logical plan to a physical plan. toPhysical :: Int -> LogicalPlan -> PhysicalPlan -- A Filter directly on a Scan is folded into the scan's pushdown predicate. toPhysical batchSz (Filter p (Scan (CsvSource path sep reader) schema)) =
src/DataFrame/Lazy/Internal/PhysicalPlan.hs view
@@ -16,17 +16,15 @@ } deriving (Show) -{- | Physical plan: every node carries enough information for the executor-to allocate resources and choose algorithms without further analysis.--}+-- | Physical plan: nodes carry what the executor needs to pick algorithms. data PhysicalPlan = PhysicalScan DataSource ScanConfig | PhysicalProject [T.Text] PhysicalPlan | PhysicalFilter (E.Expr Bool) PhysicalPlan- | PhysicalDerive T.Text E.UExpr PhysicalPlan+ | PhysicalDerive T.Text E.SomeExpr PhysicalPlan | PhysicalHashJoin JoinType T.Text T.Text PhysicalPlan PhysicalPlan | PhysicalSortMergeJoin JoinType T.Text T.Text PhysicalPlan PhysicalPlan- | PhysicalHashAggregate [T.Text] [(T.Text, E.UExpr)] PhysicalPlan+ | PhysicalHashAggregate [T.Text] [(T.Text, E.SomeExpr)] PhysicalPlan | PhysicalSort [(T.Text, SortOrder)] PhysicalPlan | PhysicalLimit Int PhysicalPlan | -- | Materialize child to a binary file on disk (used for build sides).
src/DataFrame/Typed/Lazy.hs view
@@ -6,31 +6,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- |-Module : DataFrame.Typed.Lazy-Copyright : (c) 2024 - 2026 Michael Chavinda-License : MIT-Stability : experimental--Type-safe lazy query pipelines: compile-time schema tracking ('TypedDataFrame')-with the deferred execution of 'LazyDataFrame'. Queries build a phantom-typed-logical plan; execution is deferred until 'run'.--@-{\-\# LANGUAGE DataKinds, TypeApplications, TypeOperators \#-\}-import qualified DataFrame.Typed.Lazy as TL-import DataFrame.Typed (Column)--type Schema = '[ '(\"id\", Int), '(\"name\", Text), '(\"score\", Double)]--main = do- let query = TL.scanCsv \@Schema \"data.csv\"- & TL.filter (TL.col \@\"score\" TL..>. TL.lit 0.5)- & TL.select \@'[\"id\", \"name\"]- df <- TL.run query -- TypedDataFrame '[ '(\"id\", Int), '(\"name\", Text)]- print df-@--}+-- | Type-safe lazy query pipelines with compile-time schema tracking. module DataFrame.Typed.Lazy ( -- * Core type TypedLazyDataFrame,@@ -226,10 +202,7 @@ 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.--}+-- | Runtime delegation shared by the typed joins. joinOn :: forall (key :: Symbol) left right out. (KnownSymbol key) =>@@ -258,7 +231,7 @@ run (TLD ldf) = unsafeFreeze <$> L.runDataFrame ldf -- | Convert TAgg to untyped named expressions for the lazy groupBy.-aggToNamedExprs :: TAgg keys cols aggs -> [(T.Text, E.UExpr)]+aggToNamedExprs :: TAgg keys cols aggs -> [(T.Text, E.SomeExpr)] aggToNamedExprs TAggNil = [] aggToNamedExprs (TAggCons name (TExpr expr) rest) =- (name, E.UExpr expr) : aggToNamedExprs rest+ (name, E.SomeExpr expr) : aggToNamedExprs rest
test/StreamingCsv.hs view
@@ -32,11 +32,7 @@ ) 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.--}+-- | Regression: a path-only reader sees one complete 64 MiB-or-EOF window. streamingPathReaderFillsWindow :: Test streamingPathReaderFillsWindow = TestLabel "streaming_path_reader_fills_window" $ TestCase $ do@@ -55,10 +51,7 @@ 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.--}+-- | A record newline exactly at the 64 MiB read boundary must not be lost. streamingWindowBoundaryIsLossless :: Test streamingWindowBoundaryIsLossless = TestLabel "streaming_window_boundary_is_lossless" $ TestCase $ do@@ -90,10 +83,7 @@ ] 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.--}+-- | A quoted newline spanning a 64 MiB boundary stays in one window. streamingQuotedNewlineSpansWindow :: Test streamingQuotedNewlineSpansWindow = TestLabel "streaming_quoted_newline_spans_window" $ TestCase $ do@@ -116,9 +106,7 @@ [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.--}+-- | The default streaming path strips a BOM and matches an eager read. streamingDefaultMatchesEager :: Test streamingDefaultMatchesEager = TestLabel "streaming_default_matches_eager" $ TestCase $ do@@ -136,9 +124,7 @@ 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.--}+-- | A parser failure on the producer thread must fail the consumer. streamingReaderFailurePropagates :: Test streamingReaderFailurePropagates = TestLabel "streaming_reader_failure_propagates" $ TestCase $ do