diff --git a/dataframe-lazy.cabal b/dataframe-lazy.cabal --- a/dataframe-lazy.cabal +++ b/dataframe-lazy.cabal @@ -1,6 +1,6 @@ -cabal-version: 2.4 +cabal-version: 3.4 name: dataframe-lazy -version: 1.1.0.2 +version: 2.0.0.0 synopsis: Lazy query engine for the dataframe ecosystem. description: The lazy/streaming query engine: relational-algebra plans, optimizer, @@ -31,28 +31,34 @@ DataFrame.Lazy DataFrame.Lazy.IO.Binary DataFrame.Lazy.IO.CSV + DataFrame.Typed.Lazy + -- Relational-algebra plan tree, optimizer, and pull-based executor: + -- sealed. SortOrder(..)/LazyDataFrame re-exported via DataFrame.Lazy. + other-modules: DataFrame.Lazy.Internal.DataFrame DataFrame.Lazy.Internal.Executor DataFrame.Lazy.Internal.LogicalPlan DataFrame.Lazy.Internal.Optimizer DataFrame.Lazy.Internal.PhysicalPlan - DataFrame.Typed.Lazy build-depends: base >= 4 && < 5, async >= 2.2 && < 3, - attoparsec >= 0.12 && < 0.15, - bytestring >= 0.11 && < 0.13, - containers >= 0.6.7 && < 0.9, - dataframe-core ^>= 1.1, - dataframe-csv ^>= 1.0.2, - dataframe-operations ^>= 1.1.1, - dataframe-parquet ^>= 1.1, - dataframe-parsing ^>= 1.0.2, + attoparsec >= 0.12 && < 0.16, + bytestring >= 0.11 && < 0.14, + containers >= 0.6.7 && < 0.10, + dataframe-core >= 2.0 && < 2.1, + dataframe-core:internal >= 2.0 && < 2.1, + dataframe-csv >= 2.0 && < 2.1, + dataframe-csv:internal >= 2.0 && < 2.1, + dataframe-operations >= 2.0 && < 2.1, + dataframe-parquet >= 1.2 && < 1.3, + dataframe-parsing >= 2.0 && < 2.1, + dataframe-parsing:internal >= 2.0 && < 2.1, directory >= 1.3.0.0 && < 2, filepath >= 1.4 && < 2, Glob >= 0.10 && < 1, stm >= 2.5 && < 3, temporary >= 1.3 && < 2, text >= 2.1 && < 3, - vector ^>= 0.13 + vector >= 0.13 && < 0.15 hs-source-dirs: src default-language: Haskell2010 diff --git a/src/DataFrame/Lazy.hs b/src/DataFrame/Lazy.hs --- a/src/DataFrame/Lazy.hs +++ b/src/DataFrame/Lazy.hs @@ -1,3 +1,7 @@ -module DataFrame.Lazy (module DataFrame.Lazy.Internal.DataFrame) where +module DataFrame.Lazy ( + module DataFrame.Lazy.Internal.DataFrame, + SortOrder (..), +) where import DataFrame.Lazy.Internal.DataFrame +import DataFrame.Lazy.Internal.LogicalPlan (SortOrder (..)) 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 @@ -36,8 +36,8 @@ ) import DataFrame.Internal.DataFrame (DataFrame (..)) import DataFrame.Internal.Parsing -import DataFrame.Internal.Schema (Schema, SchemaType (..), elements) import DataFrame.Operations.Typing (SafeReadMode (..), effectiveSafeRead) +import DataFrame.Schema (Schema, SchemaType (..), elements) import System.IO import Type.Reflection import Prelude hiding (takeWhile) @@ -59,11 +59,9 @@ , rowsRead :: !Int } -{- | By default we assume the file has a header and we infer types on read. -'safeRead' starts as 'NoSafeRead' — set it to 'MaybeRead' to wrap columns as -@Maybe a@, or 'EitherRead' to wrap as @Either Text a@ preserving the raw text -of any rows that fail to parse. Use 'safeReadOverrides' to pick a different -mode for specific columns. +{- | 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. -} defaultOptions :: ReadOptions defaultOptions = @@ -110,30 +108,24 @@ if hasHeader opts then fmap (T.filter (/= '\"')) firstRow else fmap (T.singleton . intToDigit) [0 .. (length firstRow - 1)] - -- If there was no header rewind the file cursor. unless (hasHeader opts) $ hSeek handle AbsoluteSeek 0 currPos <- hTell handle when (isJust $ seekPos opts) $ hSeek handle AbsoluteSeek (fromMaybe currPos (seekPos opts)) - -- Initialize mutable vectors for each column let numColumns = length columnNames let numRows = len' - -- Use this row to infer the types of the rest of the column. (dataRow, remainder) <- readSingleLine c (leftOver opts) handle - -- This array will track the indices of all null values for each column. nullIndices <- VM.unsafeNew numColumns VM.set nullIndices [] mutableCols <- VM.unsafeNew numColumns getInitialDataVectors numRows mutableCols dataRow - -- Read rows into the mutable vectors (unconsumed, r) <- fillColumns numRows c mutableCols nullIndices remainder handle - -- Freeze the mutable vectors into immutable ones nulls' <- V.unsafeFreeze nullIndices let !columnNamesV = V.fromList columnNames cols <- @@ -252,11 +244,9 @@ -- Streaming scan API -- --------------------------------------------------------------------------- -{- | Open a CSV/separated file for streaming, returning an open handle -(positioned just after the header line) and the column specification -for the schema columns that appear in the file header. - -The caller is responsible for closing the handle when done. +{- | 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. -} openCsvStream :: Char -> @@ -280,12 +270,9 @@ ("openCsvStream: none of the schema columns appear in the header of " <> path) return (handle, colSpec) -{- | Read up to @batchSz@ rows from the open handle, returning a batch -'DataFrame' and the unconsumed leftover text. Returns 'Nothing' when -the handle is at EOF and there is no leftover input. - -The caller must pass the leftover returned by the previous call (use @""@ -for the first call). +{- | 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). -} readBatch :: Char -> @@ -297,17 +284,13 @@ readBatch sep colSpec batchSz leftover handle = do let sepByte = fromIntegral (fromEnum sep) :: Word8 numCols = length colSpec - -- Read in 8 MB chunks; only the partial-line tail is copied on refill. chunkSize = 8 * 1024 * 1024 nullsArr <- VM.unsafeNew numCols VM.set nullsArr [] mCols <- VM.unsafeNew numCols forM_ (zip [0 ..] colSpec) $ \(ci, (_, _, st)) -> VM.unsafeWrite mCols ci =<< makeCol batchSz st - -- buf holds unprocessed bytes; refilled on demand when no newline is found. bufRef <- newIORef leftover - -- Row-by-row scan. When the buffer has no unquoted newline, fetch another chunk. - -- The copy on refill is only the partial-line tail (≤ one row ≈ few hundred bytes). let loop !rowIdx = do remaining <- readIORef bufRef if rowIdx >= batchSz @@ -316,7 +299,7 @@ Nothing -> do chunk <- BS.hGet handle chunkSize if BS.null chunk - then return (rowIdx, remaining) -- EOF + then return (rowIdx, remaining) else writeIORef bufRef (remaining <> chunk) >> loop rowIdx Just nlIdx -> do let line = BS.take nlIdx remaining @@ -389,7 +372,6 @@ | not (BS.any (== 0x22) bs) = skipFast targetIdx bs | otherwise = go 0 0 False 0 where - -- Fast path: skip fields using elemIndex (memchr); avoids pair allocation. skipFast k s = case BS.elemIndex sep s of Nothing -> if k == 0 then s else BS.empty @@ -398,7 +380,6 @@ then BS.take i s else skipFast (k - 1) (BS.drop (i + 1) s) - -- Slow path: quote-aware scan. quoteChar = 0x22 :: Word8 len = BS.length bs go !idx !start !inQ !pos @@ -448,10 +429,7 @@ case BS.elemIndex 0x0A bs of Nothing -> Nothing Just nlPos - -- No quote before the newline → safe to use this position. - -- Check with elemIndex to avoid allocating a ByteString slice. | maybe True (>= nlPos) (BS.elemIndex 0x22 bs) -> Just nlPos - -- Quote present → may be a newline inside a quoted field; scan carefully. | otherwise -> slowScan 0 False where len = BS.length bs 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 @@ -12,7 +12,6 @@ import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E -import DataFrame.Internal.Schema (Schema) import DataFrame.Lazy.Internal.Executor (execute) import DataFrame.Lazy.Internal.LogicalPlan ( DataSource (..), @@ -21,11 +20,10 @@ ) import qualified DataFrame.Lazy.Internal.Optimizer as Opt import DataFrame.Operations.Join (JoinType) - -{- | A lazy query that has not been executed yet. +import DataFrame.Schema (Schema) -The query is represented as a 'LogicalPlan' tree; execution is deferred -until 'runDataFrame' is called. +{- | A lazy query that has not been executed yet: a 'LogicalPlan' tree whose +execution is deferred until 'runDataFrame' is called. -} data LazyDataFrame = LazyDataFrame { plan :: LogicalPlan @@ -42,9 +40,7 @@ -- --------------------------------------------------------------------------- {- | Execute the lazy query: optimise the logical plan, then stream-execute -the resulting physical plan, returning a fully-materialised 'D.DataFrame'. -The CSV reader (default: attoparsec) is set per scan via 'scanCsv' / -'scanCsvWith'. +the resulting physical plan into a fully-materialised 'D.DataFrame'. -} runDataFrame :: LazyDataFrame -> IO D.DataFrame runDataFrame ldf = execute (Opt.optimize (batchSize ldf) (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,29 +8,9 @@ {-# LANGUAGE TupleSections #-} {-# LANGUAGE TypeApplications #-} -{- | Pull-based (iterator) execution engine. - -Each operator returns a 'Stream' — an IO action that produces the next -'DataFrame' batch on each call and returns 'Nothing' when exhausted. -Blocking operators (Sort, HashJoin) materialise their input before producing -output. HashAggregate uses streaming partial aggregation when all aggregate -expressions support it. - -== Streaming vs. materialize routing - -When the source(s) of an operator are BOUNDED (finite local CSV/Parquet -files — the db-benchmark / join-pipeline case), 'HashJoin' and -'HashAggregate' materialise their input and call the whole-frame optimized -eager op (parallel 'Join.join' with salt + ParRadixSort + parallel probe; -'Agg.aggregate' . 'Agg.groupBy' with parallel partitioned + low-card-direct -grouping). This inherits the eager Rounds 5-10 gains and produces results -byte-identical to the eager path. - -When a source is UNBOUNDED (an online/streaming source), the per-batch -streaming partial-aggregation / streaming-probe paths are preserved so the -query runs in constant memory. Boundedness is read off the physical plan -leaves by 'isBounded'. All current scan leaves are local files, so the -streaming paths are a latent capability with no in-tree producer. +{- | 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. -} module DataFrame.Lazy.Internal.Executor ( CsvReader, @@ -43,7 +23,7 @@ import Control.Concurrent.STM (atomically) import Control.Concurrent.STM.TBQueue (newTBQueueIO, readTBQueue, writeTBQueue) import Control.Exception (evaluate) -import Control.Monad (filterM, forM, forM_, when, unless) +import Control.Monad (filterM, forM, forM_, unless, when) import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as C8 import Data.IORef @@ -62,7 +42,6 @@ import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E -import DataFrame.Internal.Schema (elements) import qualified DataFrame.Lazy.IO.Binary as Bin import DataFrame.Lazy.Internal.LogicalPlan (DataSource (..), SortOrder (..)) import DataFrame.Lazy.Internal.PhysicalPlan @@ -73,6 +52,7 @@ import qualified DataFrame.Operations.Permutation as Perm import qualified DataFrame.Operations.Subset as Sub import qualified DataFrame.Operations.Transformations as Trans +import DataFrame.Schema (elements) import System.Directory (doesDirectoryExist, removeFile) import System.FilePath (()) import System.FilePath.Glob (glob) @@ -102,12 +82,8 @@ return mb {- | Drain all batches from a stream and concatenate them into one DataFrame. - -Batches are buffered, then each column is concatenated across all batches in a -single multi-way pass ('C.concatManyColumns', backed by 'VU.concat' / -'VB.concat'). A left-fold of @acc <> batch@ would copy the growing -accumulator on every step (≈ O(rows × batches)); this is O(rows) and keeps the -bounded materialize path on par with a single eager whole-frame read. +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. -} collectStream :: Stream -> IO D.DataFrame collectStream stream = go [] @@ -132,14 +108,9 @@ -- Boundedness analysis -- --------------------------------------------------------------------------- -{- | A plan is BOUNDED when every scan leaf reads a finite local source -(local CSV or local Parquet files): the whole result can be materialised, so -blocking/bounded operators route through the whole-frame fast eager ops. - -A plan is UNBOUNDED when any leaf is an online/streaming source, in which case -the streaming partial-aggregation / streaming-probe paths are kept so the query -runs in constant memory. No in-tree scan produces an unbounded source today, -so this always holds; the routing is retained for future streaming sources. +{- | 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. -} isBounded :: PhysicalPlan -> Bool isBounded (PhysicalScan (CsvSource{}) _) = True @@ -187,7 +158,6 @@ -- --------------------------------------------------------------------------- buildStream :: PhysicalPlan -> IO Stream --- Scan ----------------------------------------------------------------------- buildStream (PhysicalScan (CsvSource path sep reader) cfg) = executeCsvScan path sep reader cfg buildStream (PhysicalScan (CsvSourceStreaming path _sep reader) cfg) = @@ -199,7 +169,6 @@ Bin.spillToDisk path df df' <- Bin.readSpilled path materialized df' --- Filter --------------------------------------------------------------------- buildStream (PhysicalFilter p child) = do childStream <- buildStream child return . Stream $ @@ -207,7 +176,6 @@ mb <- pullBatch childStream return $ fmap (Sub.filterWhere p) mb ) --- Project -------------------------------------------------------------------- buildStream (PhysicalProject cols child) = do childStream <- buildStream child return . Stream $ @@ -215,7 +183,6 @@ mb <- pullBatch childStream return $ fmap (Sub.select cols) mb ) --- Derive --------------------------------------------------------------------- buildStream (PhysicalDerive name uexpr child) = do childStream <- buildStream child return . Stream $ @@ -223,7 +190,6 @@ mb <- pullBatch childStream return $ fmap (Trans.deriveMany [(name, uexpr)]) mb ) --- Limit ---------------------------------------------------------------------- buildStream (PhysicalLimit n child) = do childStream <- buildStream child countRef <- newIORef (0 :: Int) @@ -241,16 +207,11 @@ modifyIORef' countRef (+ toTake) return $ Just (Sub.take toTake df) ) --- Sort (blocking) ------------------------------------------------------------ buildStream (PhysicalSort cols child) = do df <- execute child let sortOrds = fmap (toPermSortOrder df) cols materialized (Perm.sortBy sortOrds df) --- HashAggregate -------------------------------------------------------------- buildStream (PhysicalHashAggregate keys aggs child) - -- BOUNDED child: materialise then run the whole-frame parallel grouping - -- (parallel partitioned + low-card-direct + vectorized scatter). Same - -- result the eager path produces; small-batch partial-agg overhead avoided. | isBounded child = do df <- execute child let result = Agg.aggregate aggs (Agg.groupBy keys df) @@ -259,12 +220,6 @@ childStream <- buildStream child if all (isStreamableAgg . snd) aggs then do - -- Parallel streaming partial aggregation: - -- * N workers, each pulls batches from the child stream and - -- maintains its own local accumulator. - -- * Once the stream is drained, the N partials are merged - -- sequentially using the same merge expression. - -- * O(|groups| × N) memory in flight, then O(|groups|). let (partialAggs, mergeAggs, finalizer) = buildAggPlan aggs nCaps <- getNumCapabilities let workers = max 1 nCaps @@ -286,10 +241,8 @@ writeIORef ref Nothing return mb else do - -- Fallback: materialise entire child (for CollectAgg etc.) df <- collectStream childStream materialized (Agg.aggregate aggs (Agg.groupBy keys df)) --- SourceDF (split pre-loaded DataFrame into batches) ------------------------- buildStream (PhysicalSourceDF bs df) = do let total = Core.nRows df posRef <- newIORef (0 :: Int) @@ -302,12 +255,7 @@ batch = Sub.range (i, i + n) df writeIORef posRef (i + n) return (Just batch) --- HashJoin — streaming probe (INNER/LEFT) or blocking fallback ---------------- buildStream (PhysicalHashJoin jt leftKey rightKey leftPlan rightPlan) - -- BOUNDED probe side: materialise both sides and call the whole-frame eager - -- join (parallel probe + ParRadixSort + salt). Inherits Rounds 5/8/10 gains - -- and restores parity with eager (the streaming LEFT path interleaves - -- unmatched rows differently). Streaming kept only for an unbounded probe. | isBounded leftPlan = do leftDf <- execute leftPlan rightDf <- execute rightPlan @@ -317,13 +265,11 @@ Join.INNER -> streamingHashJoin assembleInnerBatch Join.LEFT -> streamingHashJoin assembleLeftBatch _ -> do - -- Blocking fallback for RIGHT / FULL_OUTER leftDf <- execute leftPlan rightDf <- execute rightPlan materialized (performJoin jt leftKey rightKey leftDf rightDf) where streamingHashJoin assembleFn = do - -- Materialise build (right) side once and build the compact index. rightDf <- execute rightPlan let rightDf' = if leftKey == rightKey @@ -333,7 +279,6 @@ csSet = S.fromList [joinKey] rightHashes = Join.buildHashColumn [joinKey] rightDf' ci = Join.buildCompactIndex rightHashes - -- Stream probe (left) side batch by batch. leftStream <- buildStream leftPlan return . Stream $ do mBatch <- pullBatch leftStream @@ -346,7 +291,6 @@ assembleLeftBatch csSet probeBatch rightDf' probeIxs buildIxs = let batchN = Core.nRows probeBatch - -- Mark which probe rows were matched (may have duplicates — that's fine). matched = VU.accumulate (\_ b -> b) @@ -358,8 +302,6 @@ in Join.assembleLeft csSet probeBatch rightDf' allProbeIxs allBuildIxs assembleInnerBatch = Join.assembleInner - --- SortMergeJoin (blocking on both sides) ------------------------------------- buildStream (PhysicalSortMergeJoin jt leftKey rightKey leftPlan rightPlan) = do leftDf <- execute leftPlan rightDf <- execute rightPlan @@ -417,25 +359,21 @@ isStreamableAgg (E.UExpr (E.Agg (E.CollectAgg _ _) _)) = False isStreamableAgg (E.UExpr (E.Agg (E.FoldAgg _ Nothing (_ :: a -> b -> a)) _)) = case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> True -- self-merging: min, max, sum + Just Refl -> True Nothing -> False isStreamableAgg (E.UExpr (E.Agg (E.FoldAgg _ (Just _) (_ :: a -> b -> a)) _)) = case testEquality (typeRep @a) (typeRep @Int) of - Just Refl -> True -- seeded Int fold (old-style count): merge by sum + Just Refl -> True Nothing -> case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> True -- seeded self-merging + Just Refl -> True Nothing -> False isStreamableAgg (E.UExpr (E.Agg (E.MergeAgg{}) _)) = True isStreamableAgg _ = False -{- | Build the partial, merge, and finalizer plan for a list of streamable -aggregate expressions. - -* @partialAggs@ — applied per batch, producing one row per group -* @mergeAggs@ — applied when combining two partial-result DataFrames -* @finalizer@ — post-process after all batches (needed for 'MergeAgg' - where the accumulator type differs from the output type) +{- | 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). -} buildAggPlan :: [(T.Text, E.UExpr)] -> @@ -451,7 +389,6 @@ (T.Text, E.UExpr) -> ([(T.Text, E.UExpr)], [(T.Text, E.UExpr)], D.DataFrame -> D.DataFrame) processAgg (name, ue) = case ue of - -- Seedless FoldAgg: min, max, sum (self-merging when a = b) E.UExpr (E.Agg (E.FoldAgg n Nothing (f :: a -> b -> a)) (_ :: E.Expr b)) -> case testEquality (typeRep @a) (typeRep @b) of Just Refl -> @@ -460,7 +397,6 @@ , id ) Nothing -> - -- a /= b but a = Int: merge by sum (backward compat) case testEquality (typeRep @a) (typeRep @Int) of Just Refl -> ( [(name, ue)] @@ -474,7 +410,6 @@ , id ) Nothing -> ([(name, ue)], [(name, ue)], id) - -- Seeded FoldAgg: old-style count (a = Int) E.UExpr (E.Agg (E.FoldAgg n (Just _) (f :: a -> b -> a)) (_ :: E.Expr b)) -> case testEquality (typeRep @a) (typeRep @Int) of Just Refl -> @@ -496,10 +431,6 @@ , id ) Nothing -> ([(name, ue)], [(name, ue)], id) - -- MergeAgg: count, mean, etc. - -- Partial step: accumulate into acc type (using id as finalizer). - -- Merge step: apply merge function to two acc-typed partial results. - -- Finalizer: apply fin to convert acc column to output type. E.UExpr ( E.Agg ( E.MergeAgg @@ -571,31 +502,20 @@ -- CSV scan implementation -- --------------------------------------------------------------------------- -{- | CSV scan, SIMD-parallel. - -The file is read once into memory, split at newline boundaries into N -ByteString slices (N = RTS capabilities), and each slice is parsed in -parallel with the SIMD reader from "DataFrame.IO.CSV.Fast" via the -in-memory entry point — no temp-file roundtrip. The resulting per-chunk -DataFrames are sliced into batches and a dedicated thread feeds them -into a bounded queue. Pushdown predicates are applied per batch by the -consumer. +{- | SIMD-parallel CSV scan: the file is split at newline boundaries into one +slice per capability, parsed concurrently, sliced into batches, and fed through +a bounded queue. Pushdown predicates are applied per batch by the consumer. -} executeCsvScan :: FilePath -> Char -> CsvReader -> ScanConfig -> IO Stream executeCsvScan path _sep reader cfg = do nCaps <- getNumCapabilities chunkPaths <- splitCsvAtNewlines (max 1 nCaps) path - -- Each chunk parses in parallel via the reader carried on the - -- 'CsvSource' plan node. Parsing and queue-feeding stay disjoint to - -- avoid 14 producers all hammering a shared TBQueue (STM contention - -- dominates throughput). let schema = scanSchema cfg batchSz = scanBatchSize cfg chunkDfs <- mapConcurrently (reader schema) chunkPaths mapM_ removeFile chunkPaths - -- Bounded queue with a single writer, N concurrent readers. queue <- newTBQueueIO (fromIntegral (max 4 (2 * nCaps))) _ <- forkIO $ do forM_ chunkDfs $ \df -> @@ -606,7 +526,6 @@ ( do mb <- atomically (readTBQueue queue) case mb of - -- Re-insert the sentinel so repeated pulls after EOF stay Nothing. Nothing -> atomically (writeTBQueue queue Nothing) >> return Nothing Just df -> let df' = case scanPushdownPredicate cfg of @@ -663,12 +582,9 @@ 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 -carrying the original header followed by an aligned-at-newlines slice -of the body. Returns the temp file paths; the caller is responsible -for removing them after use. The path-based 'fastReadCsvWithSchema' -mmap's each file, so we get OS-paged reads instead of a single -monolithic 'BS.readFile' of the whole input. +{- | 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. -} splitCsvAtNewlines :: Int -> FilePath -> IO [FilePath] splitCsvAtNewlines n path = do @@ -700,16 +616,9 @@ -- Join helper -- --------------------------------------------------------------------------- -{- | 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.) +{- | 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. -} performJoin :: Join.JoinType -> T.Text -> T.Text -> D.DataFrame -> D.DataFrame -> D.DataFrame @@ -727,15 +636,9 @@ -- Sort order conversion -- --------------------------------------------------------------------------- -{- | Convert a plan-level @(column, direction)@ into a Permutation 'SortOrder'. - -The lazy plan only carries the column name, but 'Perm.sortBy' recovers the -'Ord' dictionary from the type parameter of @E.Col \@a@ (it returns @EQ@ for -every row when that type does not match the column's stored element type). -We therefore read the materialised column's element type and emit @E.Col@ at -exactly that type so the comparator dispatches correctly. Unknown / non-'Ord' -element types fall back to comparing as 'T.Text' (a no-op, matching the prior -behaviour) rather than failing. +{- | 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'. -} toPermSortOrder :: D.DataFrame -> (T.Text, SortOrder) -> Perm.SortOrder toPermSortOrder df (col, dir) = @@ -748,8 +651,6 @@ Ascending -> Perm.Asc (E.Col @a col) Descending -> Perm.Desc (E.Col @a col) - -- Match the column's element type against the orderable types the - -- comparator can dispatch on, emitting E.Col at the matching type. dispatch :: C.Column -> Perm.SortOrder dispatch column = case column of C.PackedText{} -> mk @T.Text @@ -774,8 +675,6 @@ tryT @b @Char $ tryT @b @T.Text (mk @T.Text) - -- If @b@ equals the candidate orderable type @a@, build the SortOrder at - -- @a@; otherwise defer to the fallback. tryT :: forall b a. (Typeable b, C.Columnable a, Ord a) => 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 @@ -6,8 +6,8 @@ import DataFrame.IO.CSV (CsvReader) import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E -import DataFrame.Internal.Schema (Schema) import DataFrame.Operations.Join (JoinType) +import DataFrame.Schema (Schema) -- | Data source for a scan node. data DataSource 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 @@ -7,18 +7,13 @@ import qualified Data.Set as S import qualified Data.Text as T import qualified DataFrame.Internal.Expression as E -import DataFrame.Internal.Schema (Schema (..), elements) import DataFrame.Lazy.Internal.LogicalPlan import DataFrame.Lazy.Internal.PhysicalPlan - -{- | Optimise a logical plan and lower it to a physical plan. - -Rules applied bottom-up (in order): - 1. Filter fusion — merge consecutive Filter nodes into a conjunction - 2. Predicate pushdown — move Filter past Derive/Project toward Scan - 3. Dead column elim — drop Derive nodes whose output is never referenced +import DataFrame.Schema (Schema (..), elements) -After rule application @toPhysical@ selects concrete operators. +{- | 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. -} optimize :: Int -> LogicalPlan -> PhysicalPlan optimize batchSz = @@ -63,11 +58,9 @@ -- Rule 2: Predicate pushdown -- --------------------------------------------------------------------------- -{- | Push Filter nodes as close to the Scan as possible. - -* Past a @Derive@ when the predicate doesn't reference the derived column. -* Past a @Project@ when all predicate columns are in the projected set. -* Into @ScanConfig.scanPushdownPredicate@ when the child is a @Scan@. +{- | 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. -} pushPredicates :: LogicalPlan -> LogicalPlan pushPredicates (Filter p (Derive name expr child)) @@ -95,10 +88,9 @@ -- Rule 3: Dead column elimination -- --------------------------------------------------------------------------- -{- | Collect every column name that is explicitly referenced somewhere in the -plan (in filter predicates, sort keys, aggregate keys, projection lists, -join keys, and derived expressions). Returns Nothing when "all columns -are needed" (i.e. no Project restricts the output). +{- | 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). -} referencedCols :: LogicalPlan -> Maybe (S.Set T.Text) referencedCols (Scan _ schema) = Just (S.fromList (M.keys (elements schema))) @@ -165,13 +157,11 @@ -- Logical → Physical lowering -- --------------------------------------------------------------------------- -{- | Lower the (already-optimised) logical plan to a physical plan. - -Join strategy: always HashJoin (the executor can fall back to SortMerge -at runtime once statistics are available). +{- | Lower the (already-optimised) logical plan to a physical plan. Joins always +lower to HashJoin; the executor may fall back to SortMerge at runtime. -} toPhysical :: Int -> LogicalPlan -> PhysicalPlan --- Special case: Filter directly on a Scan → push into ScanConfig. +-- A Filter directly on a Scan is folded into the scan's pushdown predicate. toPhysical batchSz (Filter p (Scan (CsvSource path sep reader) schema)) = PhysicalScan (CsvSource path sep reader) 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 @@ -3,9 +3,9 @@ import qualified Data.Text as T import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E -import DataFrame.Internal.Schema (Schema) import DataFrame.Lazy.Internal.LogicalPlan (DataSource, SortOrder) import DataFrame.Operations.Join (JoinType) +import DataFrame.Schema (Schema) -- | Scan-level configuration: batch size, separator, optional pushdowns. data ScanConfig = ScanConfig 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 @@ -12,12 +12,9 @@ License : MIT Stability : experimental -Type-safe lazy query pipelines. - -This module combines the compile-time schema tracking of 'TypedDataFrame' -with the deferred execution of 'LazyDataFrame'. Queries are built as a -logical plan tree with phantom-typed schema tracking; execution is deferred -until 'run' is called. +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 \#-\} @@ -83,11 +80,11 @@ import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.Expression as E -import DataFrame.Internal.Schema (Schema) 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.Schema (Schema) import DataFrame.Typed.Expr import DataFrame.Typed.Freeze (unsafeFreeze) import DataFrame.Typed.Schema