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.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, diff --git a/src/DataFrame/Lazy/IO/Binary.hs b/src/DataFrame/Lazy/IO/Binary.hs --- a/src/DataFrame/Lazy/IO/Binary.hs +++ b/src/DataFrame/Lazy/IO/Binary.hs @@ -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 diff --git a/src/DataFrame/Lazy/IO/CSV.hs b/src/DataFrame/Lazy/IO/CSV.hs --- a/src/DataFrame/Lazy/IO/CSV.hs +++ b/src/DataFrame/Lazy/IO/CSV.hs @@ -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 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 @@ -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)} 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 @@ -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 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 @@ -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. 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 @@ -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)) = diff --git a/src/DataFrame/Lazy/Internal/PhysicalPlan.hs b/src/DataFrame/Lazy/Internal/PhysicalPlan.hs --- a/src/DataFrame/Lazy/Internal/PhysicalPlan.hs +++ b/src/DataFrame/Lazy/Internal/PhysicalPlan.hs @@ -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). 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 @@ -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 diff --git a/test/StreamingCsv.hs b/test/StreamingCsv.hs --- a/test/StreamingCsv.hs +++ b/test/StreamingCsv.hs @@ -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