dataframe-lazy (empty) → 1.0.0.0
raw patch · 11 files changed
+2267/−0 lines, 11 filesdep +Globdep +asyncdep +attoparsec
Dependencies added: Glob, async, attoparsec, base, bytestring, containers, dataframe-core, dataframe-csv, dataframe-operations, dataframe-parquet, dataframe-parsing, deepseq, directory, filepath, stm, temporary, text, vector
Files
- LICENSE +20/−0
- dataframe-lazy.cabal +59/−0
- src/DataFrame/Lazy.hs +3/−0
- src/DataFrame/Lazy/IO/Binary.hs +413/−0
- src/DataFrame/Lazy/IO/CSV.hs +468/−0
- src/DataFrame/Lazy/Internal/DataFrame.hs +148/−0
- src/DataFrame/Lazy/Internal/Executor.hs +655/−0
- src/DataFrame/Lazy/Internal/LogicalPlan.hs +49/−0
- src/DataFrame/Lazy/Internal/Optimizer.hs +209/−0
- src/DataFrame/Lazy/Internal/PhysicalPlan.hs +36/−0
- src/DataFrame/Typed/Lazy.hs +207/−0
+ LICENSE view
@@ -0,0 +1,20 @@+Copyright (c) 2026 Michael Chavinda++Permission is hereby granted, free of charge, to any person obtaining+a copy of this software and associated documentation files (the+"Software"), to deal in the Software without restriction, including+without limitation the rights to use, copy, modify, merge, publish,+distribute, sublicense, and/or sell copies of the Software, and to+permit persons to whom the Software is furnished to do so, subject to+the following conditions:++The above copyright notice and this permission notice shall be included+in all copies or substantial portions of the Software.++THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.+IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,+TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE+SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ dataframe-lazy.cabal view
@@ -0,0 +1,59 @@+cabal-version: 2.4+name: dataframe-lazy+version: 1.0.0.0++synopsis: Lazy query engine for the dataframe ecosystem.+description:+ The lazy/streaming query engine: relational-algebra plans, optimizer,+ pull-based executor, and column-oriented spill format. Includes the+ typed lazy wrapper. Currently depends on @dataframe-csv@ and+ @dataframe-parquet@ since the executor calls those readers directly.++bug-reports: https://github.com/mchav/dataframe/issues+license: MIT+license-file: LICENSE+author: Michael Chavinda+maintainer: mschavinda@gmail.com+copyright: (c) 2024-2025 Michael Chavinda+category: Data+tested-with: GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2++common warnings+ ghc-options:+ -Wincomplete-patterns+ -Wincomplete-uni-patterns+ -Wunused-imports+ -Wunused-local-binds++library+ import: warnings+ exposed-modules:+ DataFrame.Lazy+ DataFrame.Lazy.IO.Binary+ DataFrame.Lazy.IO.CSV+ 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.0,+ dataframe-csv ^>= 1.0,+ dataframe-operations ^>= 1.0,+ dataframe-parquet ^>= 1.0,+ dataframe-parsing ^>= 1.0,+ deepseq >= 1 && < 2,+ directory >= 1.3.0.0 && < 2,+ filepath >= 1.4 && < 2,+ Glob >= 0.10 && < 1,+ stm >= 2.5 && < 3,+ temporary >= 1.3 && < 2,+ text >= 2.0 && < 3,+ vector ^>= 0.13+ hs-source-dirs: src+ default-language: Haskell2010
+ src/DataFrame/Lazy.hs view
@@ -0,0 +1,3 @@+module DataFrame.Lazy (module DataFrame.Lazy.Internal.DataFrame) where++import DataFrame.Lazy.Internal.DataFrame
+ src/DataFrame/Lazy/IO/Binary.hs view
@@ -0,0 +1,413 @@+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# 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.+-}+module DataFrame.Lazy.IO.Binary (+ spillToDisk,+ readSpilled,+ withSpilled,+) where++import Control.Exception (SomeException, bracket, try)+import Control.Monad (foldM, void, when)+import qualified Data.ByteString as BS+import qualified Data.ByteString.Builder as BSB+import qualified Data.ByteString.Internal as BSI+import qualified Data.ByteString.Unsafe as BSU+import qualified Data.List as L+import qualified Data.Map.Strict as M+import qualified Data.Text as T+import qualified Data.Text.Encoding as TE+import qualified Data.Vector as V+import qualified Data.Vector.Storable as VS+import qualified Data.Vector.Unboxed as VU++import Data.Bits (setBit, shiftL, testBit, (.|.))+import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))+import Data.Word (Word16, Word32, Word64, Word8)+import qualified DataFrame.Internal.Binary as Binary+import DataFrame.Internal.Column (+ Column (..),+ bitmapTestBit,+ buildBitmapFromValid,+ )+import DataFrame.Internal.DataFrame (DataFrame (..))+import Foreign (ForeignPtr, castForeignPtr, plusForeignPtr, sizeOf)+import System.Directory (getTemporaryDirectory, removeFile)+import System.IO (IOMode (..), hClose, openTempFile, withFile)+import Type.Reflection (typeRep)++-- ---------------------------------------------------------------------------+-- Type tags+-- ---------------------------------------------------------------------------++tagInt, tagDouble, tagText, tagMaybeInt, tagMaybeDouble, tagMaybeText :: Word8+tagInt = 0+tagDouble = 1+tagText = 2+tagMaybeInt = 3+tagMaybeDouble = 4+tagMaybeText = 5++-- ---------------------------------------------------------------------------+-- Write+-- ---------------------------------------------------------------------------++-- | Serialise a 'DataFrame' to a DFBN binary file.+spillToDisk :: FilePath -> DataFrame -> IO ()+spillToDisk path df =+ withFile path WriteMode $ \h -> BSB.hPutBuilder h (buildDataFrame df)++buildDataFrame :: DataFrame -> BSB.Builder+buildDataFrame df =+ BSB.byteString "DFBN"+ <> BSB.word32LE ncols+ <> foldMap (uncurry buildColumnSchema) (zip names cols)+ <> BSB.word64LE nrows+ <> foldMap (buildColumnData nrowsInt) cols+ where+ names =+ fmap+ fst+ (L.sortBy (\a b -> compare (snd a) (snd b)) (M.toList (columnIndices df)))+ ncols = fromIntegral (length names) :: Word32+ cols = V.toList (columns df)+ nrowsInt = fst (dataframeDimensions df)+ nrows = fromIntegral nrowsInt :: Word64++buildColumnSchema :: T.Text -> Column -> BSB.Builder+buildColumnSchema name col =+ BSB.word16LE nameLen+ <> BSB.byteString nameBytes+ <> BSB.word8 (columnTypeTag col)+ where+ nameBytes = TE.encodeUtf8 name+ nameLen = fromIntegral (BS.length nameBytes) :: Word16++columnTypeTag :: Column -> Word8+columnTypeTag (UnboxedColumn Nothing (_ :: VU.Vector a)) =+ case testEquality (typeRep @a) (typeRep @Int) of+ Just Refl -> tagInt+ Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+ Just Refl -> tagDouble+ Nothing -> error "spillToDisk: unsupported UnboxedColumn element type"+columnTypeTag (UnboxedColumn (Just _) (_ :: VU.Vector a)) =+ case testEquality (typeRep @a) (typeRep @Int) of+ Just Refl -> tagMaybeInt+ Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+ Just Refl -> tagMaybeDouble+ Nothing -> error "spillToDisk: unsupported nullable UnboxedColumn element type"+columnTypeTag (BoxedColumn Nothing _) = tagText+columnTypeTag (BoxedColumn (Just _) _) = tagMaybeText++buildColumnData :: Int -> Column -> BSB.Builder+buildColumnData _ (UnboxedColumn Nothing (v :: VU.Vector a)) =+ case testEquality (typeRep @a) (typeRep @Int) of+ Just Refl -> buildIntVector v+ Nothing ->+ case testEquality (typeRep @a) (typeRep @Double) of+ Just Refl -> buildDoubleVector v+ Nothing -> error "spillToDisk: unsupported UnboxedColumn element type"+buildColumnData _ (UnboxedColumn (Just bm) (v :: VU.Vector a)) =+ case testEquality (typeRep @a) (typeRep @Int) of+ Just Refl ->+ buildNullBitmap (V.generate (VU.length v) (bitmapTestBit bm))+ <> buildIntVector v+ Nothing ->+ case testEquality (typeRep @a) (typeRep @Double) of+ Just Refl ->+ buildNullBitmap (V.generate (VU.length v) (bitmapTestBit bm))+ <> buildDoubleVector v+ Nothing -> error "spillToDisk: unsupported nullable UnboxedColumn element type"+buildColumnData _ (BoxedColumn Nothing (v :: V.Vector a)) =+ case testEquality (typeRep @a) (typeRep @T.Text) of+ Just Refl -> buildTextVector v+ Nothing -> error "spillToDisk: unsupported BoxedColumn element type"+buildColumnData _ (BoxedColumn (Just bm) (v :: V.Vector a)) =+ let isValidVec = V.generate (V.length v) (bitmapTestBit bm)+ showText x = case testEquality (typeRep @a) (typeRep @T.Text) of+ Just Refl -> x+ Nothing -> T.pack (show x)+ texts = V.imap (\i x -> if bitmapTestBit bm i then showText x else T.empty) v+ in buildNullBitmap isValidVec <> buildTextVector texts++{- | 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.+-}+buildIntVector :: VU.Vector Int -> BSB.Builder+buildIntVector v =+ let sv = VU.convert v :: VS.Vector Int+ (fp, n) = VS.unsafeToForeignPtr0 sv+ 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).+buildDoubleVector :: VU.Vector Double -> BSB.Builder+buildDoubleVector v =+ let sv = VU.convert v :: VS.Vector Double+ (fp, n) = VS.unsafeToForeignPtr0 sv+ bs = BSI.fromForeignPtr (castForeignPtr fp) 0 (n * sizeOf (0 :: Double))+ in BSB.byteString bs++-- | Write a Text vector: (num_rows+1) Word32 offsets followed by UTF-8 payload.+buildTextVector :: V.Vector T.Text -> BSB.Builder+buildTextVector v =+ foldMap BSB.word32LE offsets <> foldMap BSB.byteString encoded+ where+ encoded = V.toList (V.map TE.encodeUtf8 v)+ offsets = scanl (\acc bs -> acc + fromIntegral (BS.length bs)) (0 :: Word32) encoded++-- | Build a null-validity bitmap: 1 bit per row, packed LSB-first into bytes.+buildNullBitmap :: V.Vector Bool -> BSB.Builder+buildNullBitmap valids = foldMap (BSB.word8 . mkByte) [0 .. numBytes - 1]+ where+ n = V.length valids+ numBytes = (n + 7) `div` 8+ mkByte byteIdx =+ foldr+ ( \bit acc ->+ let row = byteIdx * 8 + bit+ in if row < n && (valids V.! row) then setBit acc bit else acc+ )+ (0 :: Word8)+ [0 .. 7]++-- ---------------------------------------------------------------------------+-- Read+-- ---------------------------------------------------------------------------++-- | @(new_offset, value)@+type ParseResult a = Either String (Int, a)++-- | Deserialise a DFBN binary file into a 'DataFrame'.+readSpilled :: FilePath -> IO DataFrame+readSpilled path = do+ bs <- BS.readFile path+ case parseDataFrame bs 0 of+ Left err -> fail ("readSpilled: " <> err)+ Right (_, df) -> return df++parseDataFrame :: BS.ByteString -> Int -> ParseResult DataFrame+parseDataFrame bs off0 = do+ (off1, magic) <- readBytes bs off0 4+ when (magic /= "DFBN") $ Left "bad magic bytes"+ (off2, ncols) <- readWord32LE bs off1+ let ncolsInt = fromIntegral ncols :: Int+ (off3, schema) <- readN ncolsInt (readColumnSchema bs) off2+ (off4, nrows64) <- readWord64LE bs off3+ let nrows = fromIntegral nrows64 :: Int+ (off5, cols) <-+ foldM+ ( \(o, acc) (_, tag) -> do+ (o', col) <- readColumnData bs o nrows tag+ return (o', acc ++ [col])+ )+ (off4, [])+ schema+ let names = fmap fst schema+ return+ ( off5+ , DataFrame+ { columns = V.fromList cols+ , columnIndices = M.fromList (zip names [0 ..])+ , dataframeDimensions = (nrows, ncolsInt)+ , derivingExpressions = M.empty+ }+ )++readColumnSchema :: BS.ByteString -> Int -> ParseResult (T.Text, Word8)+readColumnSchema bs off = do+ (off1, nameLen) <- readWord16LE bs off+ let nameLenInt = fromIntegral nameLen :: Int+ (off2, nameBytes) <- readBytes bs off1 nameLenInt+ (off3, tag) <- readWord8 bs off2+ return (off3, (TE.decodeUtf8 nameBytes, tag))++readColumnData :: BS.ByteString -> Int -> Int -> Word8 -> ParseResult Column+readColumnData bs off nrows tag+ | tag == tagInt = do+ (off', v) <- readIntColumn bs off nrows+ return (off', UnboxedColumn Nothing v)+ | tag == tagDouble = do+ (off', v) <- readDoubleColumn bs off nrows+ return (off', UnboxedColumn Nothing v)+ | tag == tagText = do+ (off', v) <- readTextColumn bs off nrows+ return (off', BoxedColumn Nothing v)+ | tag == tagMaybeInt = do+ (off1, bitmap) <- readNullBitmap bs off nrows+ (off2, v) <- readIntColumn bs off1 nrows+ let bm = buildBitmapFromValid (VU.fromList (map (\b -> if b then 1 else 0) bitmap))+ return (off2, UnboxedColumn (Just bm) v)+ | tag == tagMaybeDouble = do+ (off1, bitmap) <- readNullBitmap bs off nrows+ (off2, v) <- readDoubleColumn bs off1 nrows+ let bm = buildBitmapFromValid (VU.fromList (map (\b -> if b then 1 else 0) bitmap))+ return (off2, UnboxedColumn (Just bm) v)+ | tag == tagMaybeText = do+ (off1, bitmap) <- readNullBitmap bs off nrows+ (off2, v) <- readTextColumn bs off1 nrows+ let bm = buildBitmapFromValid (VU.fromList (map (\b -> if b then 1 else 0) bitmap))+ 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).+-}+readIntColumn :: BS.ByteString -> Int -> Int -> ParseResult (VU.Vector Int)+readIntColumn bs off nrows+ | off + nrows * 8 > BS.length bs = Left "unexpected end of input"+ | otherwise =+ let (fp, bsOff, _) = BSI.toForeignPtr bs+ fp' = castForeignPtr (plusForeignPtr fp (bsOff + off)) :: ForeignPtr Int+ 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).+-}+readDoubleColumn ::+ BS.ByteString -> Int -> Int -> ParseResult (VU.Vector Double)+readDoubleColumn bs off nrows+ | off + nrows * 8 > BS.length bs = Left "unexpected end of input"+ | otherwise =+ let (fp, bsOff, _) = BSI.toForeignPtr bs+ fp' = castForeignPtr (plusForeignPtr fp (bsOff + off)) :: ForeignPtr Double+ sv = VS.unsafeFromForeignPtr0 fp' nrows :: VS.Vector Double+ in Right (off + nrows * 8, VU.convert sv)++readTextColumn :: BS.ByteString -> Int -> Int -> ParseResult (V.Vector T.Text)+readTextColumn bs off nrows = do+ offsets <- readWord32Array bs off (nrows + 1)+ let payloadStart = off + (nrows + 1) * 4+ totalPayload = fromIntegral (last offsets) :: Int+ when (payloadStart + totalPayload > BS.length bs) $+ Left "unexpected end of input"+ let sizes =+ zipWith+ (\a b -> fromIntegral b - fromIntegral a :: Int)+ offsets+ (drop 1 offsets)+ texts =+ zipWith+ ( \o sz ->+ TE.decodeUtf8+ (BS.take sz (BS.drop (payloadStart + fromIntegral o) bs))+ )+ offsets+ sizes+ return (payloadStart + totalPayload, V.fromList texts)++-- | Read @nrows@ null-bitmap bits (ceil(nrows\/8) bytes).+readNullBitmap :: BS.ByteString -> Int -> Int -> ParseResult [Bool]+readNullBitmap bs off nrows+ | off + numBytes > BS.length bs = Left "unexpected end of input"+ | otherwise =+ Right+ ( off + numBytes+ , take+ nrows+ [ testBit (BSU.unsafeIndex bs (off + row `div` 8)) (row `mod` 8)+ | row <- [0 ..]+ ]+ )+ where+ numBytes = (nrows + 7) `div` 8++readWord8 :: BS.ByteString -> Int -> ParseResult Word8+readWord8 bs off+ | off >= BS.length bs = Left "unexpected end of input"+ | otherwise = Right (off + 1, BSU.unsafeIndex bs off)++readWord16LE :: BS.ByteString -> Int -> ParseResult Word16+readWord16LE bs off+ | off + 2 > BS.length bs = Left "unexpected end of input"+ | otherwise =+ let b0 = fromIntegral (BSU.unsafeIndex bs off) :: Word16+ b1 = fromIntegral (BSU.unsafeIndex bs (off + 1)) :: Word16+ in Right (off + 2, b0 .|. (b1 `shiftL` 8))++readWord32LE :: BS.ByteString -> Int -> ParseResult Word32+readWord32LE bs off+ | off + 4 > BS.length bs = Left "unexpected end of input"+ | otherwise = Right (off + 4, Binary.littleEndianWord32 (BS.drop off bs))++readWord64LE :: BS.ByteString -> Int -> ParseResult Word64+readWord64LE bs off+ | off + 8 > BS.length bs = Left "unexpected end of input"+ | otherwise = Right (off + 8, Binary.littleEndianWord64 (BS.drop off bs))++-- | Read @n@ consecutive Word32LE values starting at offset @off@.+readWord32Array :: BS.ByteString -> Int -> Int -> Either String [Word32]+readWord32Array bs off n+ | off + n * 4 > BS.length bs = Left "unexpected end of input"+ | otherwise =+ Right+ [ let i = off + k * 4+ in Binary.littleEndianWord32 (BS.drop i bs)+ | k <- [0 .. n - 1]+ ]++-- | Read @n@ bytes from @bs@ at @off@.+readBytes :: BS.ByteString -> Int -> Int -> ParseResult BS.ByteString+readBytes bs off n+ | off + n > BS.length bs = Left "unexpected end of input"+ | otherwise = Right (off + n, BS.take n (BS.drop off bs))++-- | Apply @f@ @n@ times sequentially, threading the offset.+readN :: Int -> (Int -> ParseResult a) -> Int -> ParseResult [a]+readN 0 _ off = Right (off, [])+readN n f off = do+ (off', x) <- f off+ (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.+-}+withSpilled :: DataFrame -> (FilePath -> IO a) -> IO a+withSpilled df action = do+ tmpDir <- getTemporaryDirectory+ bracket+ ( do+ (path, h) <- openTempFile tmpDir "dataframe_spill.dfbn"+ hClose h+ spillToDisk path df+ return path+ )+ (\path -> void (try (removeFile path) :: IO (Either SomeException ())))+ action
+ src/DataFrame/Lazy/IO/CSV.hs view
@@ -0,0 +1,468 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module DataFrame.Lazy.IO.CSV where++import qualified Data.ByteString as BS+import qualified Data.Map as M+import qualified Data.Proxy as P+import qualified Data.Text as T+import qualified Data.Text.Encoding as TextEncoding+import qualified Data.Text.IO as TIO+import qualified Data.Vector as V+import qualified Data.Vector.Mutable as VM+import qualified Data.Vector.Unboxed.Mutable as VUM++import Control.Monad (forM_, unless, when, zipWithM_)+import Data.Attoparsec.Text (IResult (..), parseWith)+import Data.Char (intToDigit)+import Data.IORef+import Data.Maybe (fromMaybe, isJust)+import Data.Type.Equality (TestEquality (testEquality))+import Data.Word (Word8)+import DataFrame.IO.Internal.MutableColumn (freezeColumn', writeColumn)+import DataFrame.Internal.Column (+ Column (..),+ MutableColumn (..),+ columnLength,+ ensureOptional,+ freezeColumnEither,+ )+import DataFrame.Internal.DataFrame (DataFrame (..))+import DataFrame.Internal.Parsing+import DataFrame.Internal.Schema (Schema, SchemaType (..), elements)+import DataFrame.Operations.Typing (SafeReadMode (..), effectiveSafeRead)+import System.IO+import Type.Reflection+import Prelude hiding (takeWhile)++-- | Record for CSV read options.+data ReadOptions = ReadOptions+ { hasHeader :: Bool+ , inferTypes :: Bool+ , safeRead :: SafeReadMode+ {- ^ Default 'SafeReadMode' for columns without an entry in+ 'safeReadOverrides'.+ -}+ , safeReadOverrides :: [(T.Text, SafeReadMode)]+ -- ^ Per-column 'SafeReadMode' overrides; takes precedence over 'safeRead'.+ , rowRange :: !(Maybe (Int, Int)) -- (start, length)+ , seekPos :: !(Maybe Integer)+ , totalRows :: !(Maybe Int)+ , leftOver :: !T.Text+ , 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.+-}+defaultOptions :: ReadOptions+defaultOptions =+ ReadOptions+ { hasHeader = True+ , inferTypes = True+ , safeRead = NoSafeRead+ , safeReadOverrides = []+ , rowRange = Nothing+ , seekPos = Nothing+ , totalRows = Nothing+ , leftOver = ""+ , 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.+-}+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.+-}+readTsv :: FilePath -> IO DataFrame+readTsv path = fst <$> readSeparated '\t' defaultOptions path++-- | Reads a character separated file into a dataframe using mutable vectors.+readSeparated ::+ Char -> ReadOptions -> FilePath -> IO (DataFrame, (Integer, T.Text, Int))+readSeparated c opts path = do+ totalRows' <- case totalRows opts of+ Nothing ->+ countRows c path >>= \total -> if hasHeader opts then return (total - 1) else return total+ Just n -> if hasHeader opts then return (n - 1) else return n+ let (_, len') = case rowRange opts of+ Nothing -> (0, totalRows')+ Just (start, len'') -> (start, min len'' (totalRows' - rowsRead opts))+ withFile path ReadMode $ \handle -> do+ firstRow <- fmap T.strip . parseSep c <$> TIO.hGetLine handle+ let columnNames =+ 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 <-+ V.mapM+ (freezeColumn columnNamesV mutableCols nulls' opts)+ (V.generate numColumns id)+ pos <- hTell handle++ return+ ( DataFrame+ { columns = cols+ , columnIndices = M.fromList (zip columnNames [0 ..])+ , dataframeDimensions = (maybe 0 columnLength (cols V.!? 0), V.length cols)+ , derivingExpressions = M.empty+ }+ , (pos, unconsumed, r + 1)+ )+{-# INLINE readSeparated #-}++getInitialDataVectors :: Int -> VM.IOVector MutableColumn -> [T.Text] -> IO ()+getInitialDataVectors n mCol xs = do+ forM_ (zip [0 ..] xs) $ \(i, x) -> do+ col <- case inferValueType x of+ "Int" ->+ MUnboxedColumn+ <$> ( (VUM.unsafeNew n :: IO (VUM.IOVector Int)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readInt x) >> return c+ )+ "Double" ->+ MUnboxedColumn+ <$> ( (VUM.unsafeNew n :: IO (VUM.IOVector Double)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readDouble x) >> return c+ )+ _ ->+ MBoxedColumn+ <$> ( (VM.unsafeNew n :: IO (VM.IOVector T.Text)) >>= \c -> VM.unsafeWrite c 0 x >> return c+ )+ VM.unsafeWrite mCol i col+{-# INLINE getInitialDataVectors #-}++-- | Reads rows from the handle and stores values in mutable vectors.+fillColumns ::+ Int ->+ Char ->+ VM.IOVector MutableColumn ->+ VM.IOVector [(Int, T.Text)] ->+ T.Text ->+ Handle ->+ IO (T.Text, Int)+fillColumns n c mutableCols nullIndices unused handle = do+ input <- newIORef unused+ rowsRead' <- newIORef (0 :: Int)+ forM_ [1 .. (n - 1)] $ \i -> do+ atEOF <- hIsEOF handle+ input' <- readIORef input+ unless (atEOF && input' == mempty) $ do+ parseWith (TIO.hGetChunk handle) (parseRow c) input' >>= \case+ Fail _unconsumed ctx er -> do+ erpos <- hTell handle+ fail $+ "Failed to parse CSV file around "+ <> show erpos+ <> " byte; due: "+ <> show er+ <> "; context: "+ <> show ctx+ Partial _ -> do+ fail "Partial handler is called"+ Done (unconsumed :: T.Text) (row :: [T.Text]) -> do+ writeIORef input unconsumed+ modifyIORef rowsRead' (+ 1)+ zipWithM_ (writeValue mutableCols nullIndices i) [0 ..] row+ l <- readIORef input+ r <- readIORef rowsRead'+ pure (l, r)+{-# INLINE fillColumns #-}++-- | Writes a value into the appropriate column, resizing the vector if necessary.+writeValue ::+ VM.IOVector MutableColumn ->+ VM.IOVector [(Int, T.Text)] ->+ Int ->+ Int ->+ T.Text ->+ IO ()+writeValue mutableCols nullIndices count colIndex value = do+ col <- VM.unsafeRead mutableCols colIndex+ res <- writeColumn count value col+ let modify val = VM.unsafeModify nullIndices ((count, val) :) colIndex+ either modify (const (return ())) res+{-# INLINE writeValue #-}++-- | Freezes a mutable vector into an immutable one, trimming it to the actual row count.+freezeColumn ::+ V.Vector T.Text ->+ VM.IOVector MutableColumn ->+ V.Vector [(Int, T.Text)] ->+ ReadOptions ->+ Int ->+ IO Column+freezeColumn colNames mutableCols nulls opts colIndex = do+ col <- VM.unsafeRead mutableCols colIndex+ let colNulls = nulls V.! colIndex+ mode =+ effectiveSafeRead+ (safeRead opts)+ (safeReadOverrides opts)+ (colNames V.! colIndex)+ case mode of+ EitherRead -> freezeColumnEither colNulls col+ MaybeRead -> do+ frozen <- freezeColumn' colNulls col+ return $! ensureOptional frozen+ NoSafeRead -> freezeColumn' colNulls col+{-# INLINE freezeColumn #-}++-- ---------------------------------------------------------------------------+-- 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.+-}+openCsvStream ::+ Char ->+ Schema ->+ FilePath ->+ IO (Handle, [(Int, T.Text, SchemaType)])+openCsvStream sep schema path = do+ handle <- openFile path ReadMode+ hSetBuffering handle (BlockBuffering (Just (8 * 1024 * 1024)))+ headerLine <- TIO.hGetLine handle+ let headerCols = fmap (T.filter (/= '"') . T.strip) (parseSep sep headerLine)+ let schemaMap = elements schema+ let colSpec =+ [ (idx, name, stype)+ | (idx, name) <- zip [0 ..] headerCols+ , Just stype <- [M.lookup name schemaMap]+ ]+ when (null colSpec) $+ hClose handle+ >> fail+ ("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).+-}+readBatch ::+ Char ->+ [(Int, T.Text, SchemaType)] ->+ Int ->+ BS.ByteString ->+ Handle ->+ IO (Maybe (DataFrame, BS.ByteString))+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+ then return (rowIdx, remaining)+ else case findUnquotedNewline remaining of+ Nothing -> do+ chunk <- BS.hGet handle chunkSize+ if BS.null chunk+ then return (rowIdx, remaining) -- EOF+ else writeIORef bufRef (remaining <> chunk) >> loop rowIdx+ Just nlIdx -> do+ let line = BS.take nlIdx remaining+ rest' = BS.drop (nlIdx + 1) remaining+ line' =+ if not (BS.null line) && BS.last line == 0x0D+ then BS.init line+ else line+ writeIORef bufRef rest'+ forM_ (zip [0 ..] colSpec) $ \(ci, (fi, _, _)) -> do+ let fieldBs = getNthFieldBs sepByte fi line'+ col <- VM.unsafeRead mCols ci+ res <- writeColumnBs rowIdx fieldBs col+ case res of+ Left nv -> VM.unsafeModify nullsArr ((rowIdx, nv) :) ci+ Right _ -> return ()+ loop (rowIdx + 1)+ (completeRows, newLeftover) <- loop 0+ if completeRows == 0+ then return Nothing+ else do+ forM_ [0 .. numCols - 1] $ \ci -> do+ col <- VM.unsafeRead mCols ci+ VM.unsafeWrite mCols ci (sliceCol completeRows col)+ nullsVec <- V.unsafeFreeze nullsArr+ cols <- V.generateM numCols $ \ci -> do+ col <- VM.unsafeRead mCols ci+ freezeColumn' (nullsVec V.! ci) col+ let colNames = [name | (_, name, _) <- colSpec]+ return $+ Just+ ( DataFrame+ { columns = cols+ , columnIndices = M.fromList (zip colNames [0 ..])+ , dataframeDimensions = (completeRows, numCols)+ , derivingExpressions = M.empty+ }+ , newLeftover+ )++{- | Write a 'ByteString' field value directly into a mutable column,+parsing numerics without an intermediate 'T.Text' allocation.+-}+writeColumnBs ::+ Int -> BS.ByteString -> MutableColumn -> IO (Either T.Text Bool)+writeColumnBs i bs (MBoxedColumn (col :: VM.IOVector a)) =+ case testEquality (typeRep @a) (typeRep @T.Text) of+ Just Refl ->+ let val = TextEncoding.decodeUtf8Lenient bs+ in VM.unsafeWrite col i val >> return (Right True)+ Nothing -> return (Left (TextEncoding.decodeUtf8Lenient bs))+writeColumnBs i bs (MUnboxedColumn (col :: VUM.IOVector a)) =+ case testEquality (typeRep @a) (typeRep @Double) of+ Just Refl -> case readByteStringDouble bs of+ Just v -> VUM.unsafeWrite col i v >> return (Right True)+ Nothing -> VUM.unsafeWrite col i 0 >> return (Left (TextEncoding.decodeUtf8Lenient bs))+ Nothing -> case testEquality (typeRep @a) (typeRep @Int) of+ Just Refl -> case readByteStringInt bs of+ Just v -> VUM.unsafeWrite col i v >> return (Right True)+ Nothing -> VUM.unsafeWrite col i 0 >> return (Left (TextEncoding.decodeUtf8Lenient bs))+ 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.+-}+getNthFieldBs :: Word8 -> Int -> BS.ByteString -> BS.ByteString+getNthFieldBs sep targetIdx bs+ | 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+ Just i ->+ if k == 0+ 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+ | pos >= len =+ if idx == targetIdx then extract start pos else BS.empty+ | otherwise =+ let c = BS.index bs pos+ in if c == quoteChar+ then go idx start (not inQ) (pos + 1)+ else+ if c == sep && not inQ+ then+ if idx == targetIdx+ then extract start pos+ else go (idx + 1) (pos + 1) False (pos + 1)+ else go idx start inQ (pos + 1)++ extract s e =+ let fieldVal = BS.take (e - s) (BS.drop s bs)+ in if BS.length fieldVal >= 2+ && BS.head fieldVal == quoteChar+ && BS.last fieldVal == quoteChar+ then BS.init (BS.tail fieldVal)+ else fieldVal+{-# INLINE getNthFieldBs #-}++-- | Allocate a fresh 'MutableColumn' for @n@ slots based on a 'SchemaType'.+makeCol :: Int -> SchemaType -> IO MutableColumn+makeCol n (SType (_ :: P.Proxy a)) =+ case testEquality (typeRep @a) (typeRep @Int) of+ Just Refl -> MUnboxedColumn <$> (VUM.unsafeNew n :: IO (VUM.IOVector Int))+ Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+ Just Refl -> MUnboxedColumn <$> (VUM.unsafeNew n :: IO (VUM.IOVector Double))+ Nothing -> MBoxedColumn <$> (VM.unsafeNew n :: IO (VM.IOVector T.Text))++-- | Slice a 'MutableColumn' to @n@ elements (no-copy view).+sliceCol :: Int -> MutableColumn -> MutableColumn+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.+-}+findUnquotedNewline :: BS.ByteString -> Maybe Int+findUnquotedNewline bs =+ 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+ slowScan !pos !inQ+ | pos >= len = Nothing+ | otherwise =+ let c = BS.index bs pos+ in if c == 0x22+ then slowScan (pos + 1) (not inQ)+ else+ if c == 0x0A && not inQ+ then Just pos+ else slowScan (pos + 1) inQ+{-# INLINE findUnquotedNewline #-}
+ src/DataFrame/Lazy/Internal/DataFrame.hs view
@@ -0,0 +1,148 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE NumericUnderscores #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}++module DataFrame.Lazy.Internal.DataFrame where++import qualified Data.Text as T+import DataFrame.IO.CSV (CsvReader, readCsvWithSchema)+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 (..),+ LogicalPlan (..),+ SortOrder (..),+ )+import qualified DataFrame.Lazy.Internal.Optimizer as Opt+import DataFrame.Operations.Join (JoinType)++{- | A lazy query that has not been executed yet.++The query is represented as a 'LogicalPlan' tree; execution is deferred+until 'runDataFrame' is called.+-}+data LazyDataFrame = LazyDataFrame+ { plan :: LogicalPlan+ , batchSize :: Int+ }++instance Show LazyDataFrame where+ show ldf =+ "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, returning a fully-materialised 'D.DataFrame'.+The CSV reader (default: attoparsec) is set per scan via 'scanCsv' /+'scanCsvWith'.+-}+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+attoparsec reader. For the SIMD reader use 'scanCsvWith'.+-}+scanCsv :: Schema -> T.Text -> LazyDataFrame+scanCsv = scanCsvWith readCsvWithSchema++{- | Like 'scanCsv' but with an explicit CSV reader (e.g. the SIMD reader+@fastReadCsvWithSchema@ from @dataframe-fastcsv@).+-}+scanCsvWith :: CsvReader -> Schema -> T.Text -> LazyDataFrame+scanCsvWith reader schema path =+ LazyDataFrame+ { plan = Scan (CsvSource (T.unpack path) ',' reader) schema+ , batchSize = 1_000_000+ }++-- | Scan a character-separated file with the default attoparsec reader.+scanSeparated :: Char -> Schema -> T.Text -> LazyDataFrame+scanSeparated = scanSeparatedWith readCsvWithSchema++-- | Like 'scanSeparated' but with an explicit CSV reader.+scanSeparatedWith ::+ CsvReader -> Char -> Schema -> T.Text -> LazyDataFrame+scanSeparatedWith reader sep schema path =+ LazyDataFrame+ { plan = Scan (CsvSource (T.unpack path) sep reader) schema+ , batchSize = 1_000_000+ }++-- | Scan a Parquet file, directory of files, or glob pattern.+scanParquet :: Schema -> T.Text -> LazyDataFrame+scanParquet schema path =+ LazyDataFrame+ { plan = Scan (ParquetSource (T.unpack path)) schema+ , batchSize = 1_000_000+ }++-- | Add a computed column (or overwrite an existing one).+derive ::+ (C.Columnable a) => T.Text -> E.Expr a -> LazyDataFrame -> LazyDataFrame+derive name expr ldf =+ ldf{plan = Derive name (E.UExpr expr) (plan ldf)}++-- | Retain only the listed columns.+select :: [T.Text] -> LazyDataFrame -> LazyDataFrame+select cols ldf = ldf{plan = Project cols (plan ldf)}++-- | Keep rows that satisfy the predicate.+filter :: E.Expr Bool -> LazyDataFrame -> LazyDataFrame+filter cond ldf = ldf{plan = Filter cond (plan ldf)}++-- | Join two lazy queries on the given key columns.+join ::+ JoinType ->+ -- | Left join key column name+ T.Text ->+ -- | Right join key column name+ T.Text ->+ -- | Left sub-query+ LazyDataFrame ->+ -- | Right sub-query+ LazyDataFrame ->+ LazyDataFrame+join jt leftKey rightKey left right =+ LazyDataFrame+ { plan = Join jt leftKey rightKey (plan left) (plan right)+ , 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@).+-}+groupBy ::+ -- | Group-by key columns+ [T.Text] ->+ -- | @[(outputName, aggregateExpr)]@+ [(T.Text, E.UExpr)] ->+ LazyDataFrame ->+ LazyDataFrame+groupBy keys aggs ldf = ldf{plan = Aggregate keys aggs (plan ldf)}++-- | Sort the result by the given @(column, direction)@ pairs.+sortBy :: [(T.Text, SortOrder)] -> LazyDataFrame -> LazyDataFrame+sortBy cols ldf = ldf{plan = Sort cols (plan ldf)}++-- | Retain at most @n@ rows.+take :: Int -> LazyDataFrame -> LazyDataFrame+take n ldf = ldf{plan = Limit n (plan ldf)}
+ src/DataFrame/Lazy/Internal/Executor.hs view
@@ -0,0 +1,655 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# 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.+-}+module DataFrame.Lazy.Internal.Executor (+ CsvReader,+ execute,+ foldBatches,+) where++import Control.Concurrent (forkIO, getNumCapabilities)+import Control.Concurrent.Async (mapConcurrently)+import Control.Concurrent.STM (atomically)+import Control.Concurrent.STM.TBQueue (newTBQueueIO, readTBQueue, writeTBQueue)+import Control.Exception (evaluate)+import Control.Monad (filterM, forM, forM_, when)+import qualified Data.ByteString as BS+import Data.IORef+import qualified Data.Map as M+import qualified Data.Maybe+import qualified Data.Set as S+import qualified Data.Text as T+import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))+import qualified Data.Vector.Unboxed as VU+import Data.Word (Word8)+import DataFrame.IO.CSV (CsvReader)+import qualified DataFrame.IO.Parquet as Parquet+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+import qualified DataFrame.Operations.Aggregation as Agg+import qualified DataFrame.Operations.Core as Core+import qualified DataFrame.Operations.Join as Join+import DataFrame.Operations.Merge ()+import qualified DataFrame.Operations.Permutation as Perm+import qualified DataFrame.Operations.Subset as Sub+import qualified DataFrame.Operations.Transformations as Trans+import System.Directory (doesDirectoryExist, removeFile)+import System.FilePath ((</>))+import System.FilePath.Glob (glob)+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.+-}+newtype Stream = Stream {pullBatch :: IO (Maybe D.DataFrame)}++-- | Drain all batches from a stream and concatenate them into one DataFrame.+collectStream :: Stream -> IO D.DataFrame+collectStream stream = go D.empty+ where+ go acc = do+ mb <- pullBatch stream+ case mb of+ Nothing -> return acc+ Just df -> go (acc <> df)++-- ---------------------------------------------------------------------------+-- Top-level entry point+-- ---------------------------------------------------------------------------++{- | Execute a physical plan, returning the complete result as a single+'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.+-}+foldBatches ::+ (b -> D.DataFrame -> IO b) -> b -> PhysicalPlan -> IO b+foldBatches f seed plan = do+ stream <- buildStream plan+ let loop !acc = do+ mb <- pullBatch stream+ case mb of+ Nothing -> return acc+ Just batch -> do+ !acc' <- f acc batch+ loop acc'+ loop seed++-- ---------------------------------------------------------------------------+-- Per-operator stream builders+-- ---------------------------------------------------------------------------++buildStream :: PhysicalPlan -> IO Stream+-- Scan -----------------------------------------------------------------------+buildStream (PhysicalScan (CsvSource path sep reader) cfg) =+ executeCsvScan path sep reader cfg+buildStream (PhysicalScan (ParquetSource path) cfg) =+ executeParquetScan path cfg+buildStream (PhysicalSpill child path) = do+ df <- execute child+ Bin.spillToDisk path df+ df' <- Bin.readSpilled path+ ref <- newIORef (Just df')+ return . Stream $+ ( do+ mb <- readIORef ref+ writeIORef ref Nothing+ return mb+ )+-- Filter ---------------------------------------------------------------------+buildStream (PhysicalFilter p child) = do+ childStream <- buildStream child+ return . Stream $+ ( do+ mb <- pullBatch childStream+ return $ fmap (Sub.filterWhere p) mb+ )+-- Project --------------------------------------------------------------------+buildStream (PhysicalProject cols child) = do+ childStream <- buildStream child+ return . Stream $+ ( do+ mb <- pullBatch childStream+ return $ fmap (Sub.select cols) mb+ )+-- Derive ---------------------------------------------------------------------+buildStream (PhysicalDerive name uexpr child) = do+ childStream <- buildStream child+ return . Stream $+ ( do+ mb <- pullBatch childStream+ return $ fmap (Trans.deriveMany [(name, uexpr)]) mb+ )+-- Limit ----------------------------------------------------------------------+buildStream (PhysicalLimit n child) = do+ childStream <- buildStream child+ countRef <- newIORef (0 :: Int)+ return . Stream $+ ( do+ remaining <- readIORef countRef+ if remaining >= n+ then return Nothing+ else do+ mb <- pullBatch childStream+ case mb of+ Nothing -> return Nothing+ Just df -> do+ let toTake = min (Core.nRows df) (n - remaining)+ modifyIORef' countRef (+ toTake)+ return $ Just (Sub.take toTake df)+ )+-- Sort (blocking) ------------------------------------------------------------+buildStream (PhysicalSort cols child) = do+ df <- execute child+ let sortOrds = fmap toPermSortOrder cols+ let sorted = Perm.sortBy sortOrds df+ ref <- newIORef (Just sorted)+ return . Stream $+ ( do+ mb <- readIORef ref+ writeIORef ref Nothing+ return mb+ )+-- HashAggregate --------------------------------------------------------------+buildStream (PhysicalHashAggregate keys aggs child) = do+ 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+ partials <-+ mapConcurrently+ (\_ -> workerLoop childStream keys partialAggs mergeAggs)+ [1 .. workers]+ mFinal <-+ let nonEmpty = Data.Maybe.catMaybes partials+ in case nonEmpty of+ [] -> return Nothing+ [single] -> return (Just (finalizer single))+ (a : rest) -> do+ !merged <- mergePartials keys mergeAggs a rest+ return (Just (finalizer merged))+ ref <- newIORef mFinal+ return . Stream $ do+ mb <- readIORef ref+ writeIORef ref Nothing+ return mb+ else do+ -- Fallback: materialise entire child (for CollectAgg etc.)+ df <- collectStream childStream+ let result = Agg.aggregate aggs (Agg.groupBy keys df)+ ref <- newIORef (Just result)+ return . Stream $ do+ mb <- readIORef ref+ writeIORef ref Nothing+ return mb+-- SourceDF (split pre-loaded DataFrame into batches) -------------------------+buildStream (PhysicalSourceDF bs df) = do+ let total = Core.nRows df+ posRef <- newIORef (0 :: Int)+ return . Stream $ do+ i <- readIORef posRef+ if i >= total+ then return Nothing+ else do+ let n = min bs (total - i)+ 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) =+ case jt of+ Join.INNER -> streamingHashJoin assembleInnerBatch+ Join.LEFT -> streamingHashJoin assembleLeftBatch+ _ -> do+ -- Blocking fallback for RIGHT / FULL_OUTER+ leftDf <- execute leftPlan+ rightDf <- execute rightPlan+ let result = performJoin jt leftKey rightKey leftDf rightDf+ ref <- newIORef (Just result)+ return . Stream $ do+ mb <- readIORef ref+ writeIORef ref Nothing+ return mb+ where+ streamingHashJoin assembleFn = do+ -- Materialise build (right) side once and build the compact index.+ rightDf <- execute rightPlan+ let rightDf' =+ if leftKey == rightKey+ then rightDf+ else Core.rename rightKey leftKey rightDf+ joinKey = leftKey+ 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+ case mBatch of+ Nothing -> return Nothing+ Just probeBatch -> do+ let probeHashes = Join.buildHashColumn [joinKey] probeBatch+ (probeIxs, buildIxs) = Join.hashProbeKernel ci probeHashes+ return . Just $ assembleFn csSet probeBatch rightDf' probeIxs buildIxs++ 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)+ (VU.replicate batchN False)+ (VU.map (,True) probeIxs)+ unmatchedIxs = VU.findIndices not matched+ allProbeIxs = probeIxs VU.++ unmatchedIxs+ allBuildIxs = buildIxs VU.++ VU.replicate (VU.length unmatchedIxs) (-1)+ 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+ let result = performJoin jt leftKey rightKey leftDf rightDf+ ref <- newIORef (Just result)+ return . Stream $+ ( do+ mb <- readIORef ref+ writeIORef ref Nothing+ return mb+ )++-- ---------------------------------------------------------------------------+-- Streaming aggregation helpers+-- ---------------------------------------------------------------------------++{- | True when an aggregate expression can be computed incrementally+(i.e., partial results can be merged without materialising all rows).+-}++{- | One worker's loop: pull batches off the shared child stream until+exhausted, building up a per-worker accumulator.+-}+workerLoop ::+ Stream ->+ [T.Text] ->+ [E.NamedExpr] ->+ [E.NamedExpr] ->+ IO (Maybe D.DataFrame)+workerLoop childStream keys partialAggs mergeAggs = loop Nothing+ where+ loop !acc = do+ mb <- pullBatch childStream+ case mb of+ Nothing -> return acc+ Just batch -> do+ !partial <-+ evaluate . D.forceDataFrame $+ Agg.aggregate partialAggs (Agg.groupBy keys batch)+ !next <- case acc of+ Nothing -> return (Just partial)+ Just a -> do+ !merged <-+ evaluate . D.forceDataFrame $+ Agg.aggregate mergeAggs (Agg.groupBy keys (a <> partial))+ return (Just merged)+ loop next++-- | Merge a head accumulator with the rest of the workers' partials.+mergePartials ::+ [T.Text] ->+ [E.NamedExpr] ->+ D.DataFrame ->+ [D.DataFrame] ->+ IO D.DataFrame+mergePartials keys mergeAggs = go+ where+ go !acc [] = return acc+ go !acc (p : ps) = do+ !merged <-+ evaluate . D.forceDataFrame $+ 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)) _)) =+ case testEquality (typeRep @a) (typeRep @b) of+ Just Refl -> True -- self-merging: min, max, sum+ 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+ Nothing ->+ case testEquality (typeRep @a) (typeRep @b) of+ Just Refl -> True -- seeded self-merging+ 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)+-}+buildAggPlan ::+ [(T.Text, E.UExpr)] ->+ ( [(T.Text, E.UExpr)]+ , [(T.Text, E.UExpr)]+ , D.DataFrame -> D.DataFrame+ )+buildAggPlan aggs = foldl combine ([], [], id) (map processAgg aggs)+ where+ 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)+ 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 ->+ ( [(name, ue)]+ , [(name, E.UExpr (E.Agg (E.FoldAgg n Nothing f) (E.Col @a name)))]+ , id+ )+ Nothing ->+ -- a /= b but a = Int: merge by sum (backward compat)+ case testEquality (typeRep @a) (typeRep @Int) of+ Just Refl ->+ ( [(name, ue)]+ ,+ [+ ( name+ , E.UExpr+ (E.Agg (E.FoldAgg "sum" Nothing ((+) :: Int -> Int -> Int)) (E.Col @Int name))+ )+ ]+ , 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 ->+ ( [(name, ue)]+ ,+ [+ ( name+ , E.UExpr+ (E.Agg (E.FoldAgg "sum" Nothing ((+) :: Int -> Int -> Int)) (E.Col @Int name))+ )+ ]+ , id+ )+ Nothing ->+ 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)))]+ , 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+ n+ seed+ (step :: acc -> b -> acc)+ (merge :: acc -> acc -> acc)+ (fin :: acc -> a)+ )+ (inner :: E.Expr b)+ ) ->+ let partialExpr =+ E.UExpr+ ( E.Agg+ (E.MergeAgg n seed step merge (id :: acc -> acc))+ inner+ )+ mergeExpr =+ E.UExpr+ ( E.Agg+ (E.FoldAgg ("merge_" <> n) Nothing merge)+ (E.Col @acc name)+ )+ finalize df =+ let accCol = D.unsafeGetColumn name df+ finalCol =+ either+ (error "buildAggPlan: MergeAgg finalize failed")+ id+ (C.mapColumn @acc @a fin accCol)+ in D.insertColumn name finalCol df+ in ( [(name, partialExpr)]+ , [(name, mergeExpr)]+ , finalize+ )+ _ -> ([(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'.+-}+executeParquetScan :: FilePath -> ScanConfig -> IO Stream+executeParquetScan path cfg+ | Parquet.isHFUri path = executeHFParquetScan path cfg+ | otherwise = do+ isDir <- doesDirectoryExist path+ let pat = if isDir then path </> "*" else path+ matches <- glob pat+ files <- filterM (fmap not . doesDirectoryExist) matches+ when (null files) $+ error ("executeParquetScan: no parquet files found for " ++ path)+ let opts =+ Parquet.defaultParquetReadOptions+ { Parquet.selectedColumns = Just (M.keys (elements (scanSchema cfg)))+ , Parquet.predicate = scanPushdownPredicate cfg+ }+ ref <- newIORef files+ return . Stream $ do+ fs <- readIORef ref+ case fs of+ [] -> return Nothing+ (f : rest) -> do+ writeIORef ref rest+ Just <$> Parquet.readParquetWithOpts opts f++{- | HuggingFace Parquet scan. Files are resolved once (API call or direct URL)+then downloaded one at a time as the stream is pulled — so only one file's worth+of data is in memory at a time, regardless of dataset size.+-}++-- TODO: mchavinda - this should be a more general online file scanner.+executeHFParquetScan :: FilePath -> ScanConfig -> IO Stream+executeHFParquetScan path cfg = do+ ref <- case Parquet.parseHFUri path of+ Left err -> error err+ Right r -> pure r+ mToken <- Parquet.getHFToken+ hfFiles <-+ if Parquet.hasGlob (Parquet.hfGlob ref)+ then Parquet.resolveHFUrls mToken ref+ else do+ let url = Parquet.directHFUrl ref+ filename = last $ T.splitOn "/" (Parquet.hfGlob ref)+ pure [Parquet.HFParquetFile url "" "" filename]+ when (null hfFiles) $+ error ("executeParquetScan: no HF parquet files found for " ++ path)+ let opts =+ Parquet.defaultParquetReadOptions+ { Parquet.selectedColumns = Just (M.keys (elements (scanSchema cfg)))+ , Parquet.predicate = scanPushdownPredicate cfg+ }+ filesRef <- newIORef hfFiles+ return . Stream $ do+ fs <- readIORef filesRef+ case fs of+ [] -> return Nothing+ (f : rest) -> do+ writeIORef filesRef rest+ -- Download a single file, read it, then return the batch.+ [localPath] <- Parquet.downloadHFFiles mToken [f]+ Just <$> Parquet.readParquetWithOpts opts localPath++-- ---------------------------------------------------------------------------+-- 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.+-}+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 ->+ forM_ (sliceIntoBatches batchSz df) $ \b ->+ atomically (writeTBQueue queue (Just b))+ atomically (writeTBQueue queue Nothing)+ return . Stream $+ ( 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+ Nothing -> df+ Just p -> Sub.filterWhere p df+ in return (Just df')+ )++-- | Slice a 'DataFrame' into row-bounded batches of at most @n@ rows.+sliceIntoBatches :: Int -> D.DataFrame -> [D.DataFrame]+sliceIntoBatches n df =+ let total = Core.nRows df+ 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.+-}+splitCsvAtNewlines :: Int -> FilePath -> IO [FilePath]+splitCsvAtNewlines n path = do+ bs <- BS.readFile path+ let (header, rest) = BS.break (== nl) bs+ body = BS.drop 1 rest+ bodyLen = BS.length body+ rawOffsets = [(bodyLen * i) `div` n | i <- [0 .. n]]+ snapped = 0 : map (snap body) (init (drop 1 rawOffsets)) ++ [bodyLen]+ ranges = zip snapped (drop 1 snapped)+ slices =+ [ BS.take (hi - lo) (BS.drop lo body)+ | (lo, hi) <- ranges+ , hi > lo+ ]+ forM slices $ \chunk -> do+ p <- emptySystemTempFile "lazy_csv_chunk_.csv"+ BS.writeFile p (header <> BS.singleton nl <> chunk)+ return p+ where+ nl :: Word8+ nl = 0x0A+ snap body off =+ case BS.elemIndex nl (BS.drop off body) of+ Just i -> off + i + 1+ Nothing -> BS.length body++-- ---------------------------------------------------------------------------+-- Join helper+-- ---------------------------------------------------------------------------++{- | Route join to the existing Operations.Join implementation.+When the left and right key names differ, rename the right key before joining.+-}+performJoin ::+ Join.JoinType -> T.Text -> T.Text -> D.DataFrame -> D.DataFrame -> D.DataFrame+performJoin jt leftKey rightKey leftDf rightDf =+ if leftKey == rightKey+ then Join.join jt [leftKey] rightDf leftDf+ else+ let rightRenamed = Core.rename rightKey leftKey rightDf+ in Join.join jt [leftKey] rightRenamed leftDf++-- ---------------------------------------------------------------------------+-- Sort order conversion+-- ---------------------------------------------------------------------------++-- | Convert plan-level sort order to the Permutation module's SortOrder.+toPermSortOrder :: (T.Text, SortOrder) -> Perm.SortOrder+toPermSortOrder (col, Ascending) = Perm.Asc (E.Col @T.Text col)+toPermSortOrder (col, Descending) = Perm.Desc (E.Col @T.Text col)
+ src/DataFrame/Lazy/Internal/LogicalPlan.hs view
@@ -0,0 +1,49 @@+{-# LANGUAGE GADTs #-}++module DataFrame.Lazy.Internal.LogicalPlan where++import qualified Data.Text as T+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)++-- | Data source for a scan node.+data DataSource+ = -- | path, separator, CSV reader (e.g. attoparsec or SIMD)+ CsvSource FilePath Char CsvReader+ | ParquetSource FilePath++instance Show DataSource where+ show (CsvSource path sep _) =+ "CsvSource " ++ show path ++ " " ++ show sep ++ " <reader>"+ show (ParquetSource path) = "ParquetSource " ++ show path++-- | Sort direction used in Sort nodes and the public API.+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.+-}+data LogicalPlan+ = -- | Read columns described by the schema from a source.+ Scan DataSource Schema+ | -- | Retain only the listed columns.+ Project [T.Text] LogicalPlan+ | -- | Keep rows matching the predicate.+ Filter (E.Expr Bool) LogicalPlan+ | -- | Add or overwrite a column via an expression.+ Derive T.Text E.UExpr 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+ | -- | Sort by a list of (column, direction) pairs.+ Sort [(T.Text, SortOrder)] LogicalPlan+ | -- | Retain at most N rows.+ Limit Int LogicalPlan+ | -- | Lift an already-loaded DataFrame into the lazy plan.+ SourceDF D.DataFrame+ deriving (Show)
+ src/DataFrame/Lazy/Internal/Optimizer.hs view
@@ -0,0 +1,209 @@+{-# LANGUAGE OverloadedStrings #-}++module DataFrame.Lazy.Internal.Optimizer (optimize) where++import qualified Data.Map as M+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++After rule application @toPhysical@ selects concrete operators.+-}+optimize :: Int -> LogicalPlan -> PhysicalPlan+optimize batchSz =+ toPhysical batchSz+ . eliminateDeadColumns+ . 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)) =+ fuseFilters (Filter (andExpr p1 p2) (fuseFilters child))+fuseFilters (Filter p child) = Filter p (fuseFilters child)+fuseFilters (Project cols child) = Project cols (fuseFilters child)+fuseFilters (Derive name expr child) = Derive name expr (fuseFilters child)+fuseFilters (Join jt l r left right) =+ Join jt l r (fuseFilters left) (fuseFilters right)+fuseFilters (Aggregate keys aggs child) =+ Aggregate keys aggs (fuseFilters child)+fuseFilters (Sort cols child) = Sort cols (fuseFilters child)+fuseFilters (Limit n child) = Limit n (fuseFilters child)+fuseFilters leaf = leaf++-- | Logical AND of two @Bool@ expressions.+andExpr :: E.Expr Bool -> E.Expr Bool -> E.Expr Bool+andExpr =+ E.Binary+ ( E.MkBinaryOp+ { E.binaryFn = (&&)+ , E.binaryName = "and"+ , E.binarySymbol = Just "&&"+ , E.binaryCommutative = True+ , E.binaryPrecedence = 3+ }+ )++-- ---------------------------------------------------------------------------+-- 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@.+-}+pushPredicates :: LogicalPlan -> LogicalPlan+pushPredicates (Filter p (Derive name expr child))+ | name `notElem` E.getColumns p =+ Derive name expr (pushPredicates (Filter p child))+ | otherwise =+ Filter p (Derive name expr (pushPredicates child))+pushPredicates (Filter p (Project cols child))+ | all (`elem` cols) (E.getColumns p) =+ Project cols (pushPredicates (Filter p child))+ | otherwise =+ Filter p (Project cols (pushPredicates child))+pushPredicates (Filter p child) = Filter p (pushPredicates child)+pushPredicates (Project cols child) = Project cols (pushPredicates child)+pushPredicates (Derive name expr child) = Derive name expr (pushPredicates child)+pushPredicates (Join jt l r left right) =+ Join jt l r (pushPredicates left) (pushPredicates right)+pushPredicates (Aggregate keys aggs child) =+ Aggregate keys aggs (pushPredicates child)+pushPredicates (Sort cols child) = Sort cols (pushPredicates child)+pushPredicates (Limit n child) = Limit n (pushPredicates child)+pushPredicates leaf = leaf++-- ---------------------------------------------------------------------------+-- 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).+-}+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)+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 liftMaybe2 S.union lRef rRef+referencedCols (Aggregate keys aggs child) =+ let aggCols = S.fromList (keys <> concatMap (uExprCols . 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++liftMaybe2 :: (a -> b -> c) -> Maybe a -> Maybe b -> Maybe c+liftMaybe2 f (Just a) (Just b) = Just (f a b)+liftMaybe2 _ _ _ = Nothing++uExprCols :: E.UExpr -> [T.Text]+uExprCols (E.UExpr expr) = E.getColumns expr++-- | Drop @Derive@ nodes whose output column is never consumed downstream.+eliminateDeadColumns :: LogicalPlan -> LogicalPlan+eliminateDeadColumns plan = go (referencedCols plan) plan+ where+ go needed (Derive name expr child) =+ case needed of+ Nothing -> Derive name expr (go needed child)+ Just cols+ | name `S.notMember` cols -> go needed child+ | otherwise ->+ Derive+ name+ expr+ (go (Just (S.union cols (S.fromList (uExprCols 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) =+ Project cols (go (Just (S.fromList cols)) child)+ go needed (Join jt l r left right) =+ 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+ 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)+ go needed (Limit n child) = Limit n (go needed child)+ go needed (Scan ds schema) =+ case needed of+ Nothing -> Scan ds schema+ Just cols ->+ 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.++Join strategy: always HashJoin (the executor can fall back to SortMerge+at runtime once statistics are available).+-}+toPhysical :: Int -> LogicalPlan -> PhysicalPlan+-- Special case: Filter directly on a Scan → push into ScanConfig.+toPhysical batchSz (Filter p (Scan (CsvSource path sep reader) schema)) =+ PhysicalScan+ (CsvSource path sep reader)+ (ScanConfig batchSz sep schema (Just p))+toPhysical batchSz (Scan (CsvSource path sep reader) schema) =+ PhysicalScan+ (CsvSource path sep reader)+ (ScanConfig batchSz sep schema Nothing)+toPhysical batchSz (Filter p (Scan (ParquetSource path) schema)) =+ PhysicalScan+ (ParquetSource path)+ (ScanConfig batchSz ',' schema (Just p))+toPhysical batchSz (Scan (ParquetSource path) schema) =+ PhysicalScan+ (ParquetSource path)+ (ScanConfig batchSz ',' schema Nothing)+toPhysical batchSz (Project cols child) =+ PhysicalProject cols (toPhysical batchSz child)+toPhysical batchSz (Filter p child) =+ PhysicalFilter p (toPhysical batchSz child)+toPhysical batchSz (Derive name expr child) =+ PhysicalDerive name expr (toPhysical batchSz child)+toPhysical batchSz (Join jt l r left right) =+ PhysicalHashJoin+ jt+ l+ r+ (toPhysical batchSz left)+ (toPhysical batchSz right)+toPhysical batchSz (Aggregate keys aggs child) =+ PhysicalHashAggregate keys aggs (toPhysical batchSz child)+toPhysical batchSz (Sort cols child) =+ PhysicalSort cols (toPhysical batchSz child)+toPhysical batchSz (Limit n child) =+ PhysicalLimit n (toPhysical batchSz child)+toPhysical batchSz (SourceDF df) = PhysicalSourceDF batchSz df
+ src/DataFrame/Lazy/Internal/PhysicalPlan.hs view
@@ -0,0 +1,36 @@+module DataFrame.Lazy.Internal.PhysicalPlan where++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)++-- | Scan-level configuration: batch size, separator, optional pushdowns.+data ScanConfig = ScanConfig+ { scanBatchSize :: !Int+ , scanSeparator :: !Char+ , scanSchema :: !Schema+ , scanPushdownPredicate :: !(Maybe (E.Expr Bool))+ }+ deriving (Show)++{- | Physical plan: every node carries enough information for the executor+to allocate resources and choose algorithms without further analysis.+-}+data PhysicalPlan+ = PhysicalScan DataSource ScanConfig+ | PhysicalProject [T.Text] PhysicalPlan+ | PhysicalFilter (E.Expr Bool) PhysicalPlan+ | PhysicalDerive T.Text E.UExpr 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+ | PhysicalSort [(T.Text, SortOrder)] PhysicalPlan+ | PhysicalLimit Int PhysicalPlan+ | -- | Materialize child to a binary file on disk (used for build sides).+ PhysicalSpill PhysicalPlan FilePath+ | -- | Emit an already-loaded DataFrame as a stream of batches of size @n@.+ PhysicalSourceDF Int D.DataFrame+ deriving (Show)
+ src/DataFrame/Typed/Lazy.hs view
@@ -0,0 +1,207 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- |+Module : DataFrame.Typed.Lazy+Copyright : (c) 2025+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.++@+{\-\# LANGUAGE DataKinds, TypeApplications, TypeOperators \#-\}+import qualified DataFrame.Typed.Lazy as TL+import DataFrame.Typed (Column)++type Schema = '[Column \"id\" Int, Column \"name\" Text, Column \"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 '[Column \"id\" Int, Column \"name\" Text]+ print df+@+-}+module DataFrame.Typed.Lazy (+ -- * Core type+ TypedLazyDataFrame,++ -- * Data sources+ scanCsv,+ scanSeparated,+ scanParquet,+ fromDataFrame,+ fromTypedDataFrame,++ -- * Schema-preserving operations+ filter,+ take,++ -- * Schema-modifying operations+ derive,+ select,++ -- * Aggregation+ groupBy,+ aggregate,++ -- * Joins+ join,++ -- * Sort+ sortBy,++ -- * Execution+ run,++ -- * Re-exports for pipeline construction+ module DataFrame.Typed.Expr,+ module DataFrame.Typed.Types,+ SortOrder (..),+) where++import Data.Kind (Type)+import Data.Proxy (Proxy (..))+import qualified Data.Text as T+import GHC.TypeLits (KnownSymbol, Symbol, symbolVal)+import Prelude hiding (filter, take)++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.Typed.Expr+import DataFrame.Typed.Freeze (unsafeFreeze)+import DataFrame.Typed.Schema+import DataFrame.Typed.Types++-- | A lazy query with compile-time schema tracking.+newtype TypedLazyDataFrame (cols :: [Type]) = TLD {_unTLD :: LazyDataFrame}++instance Show (TypedLazyDataFrame cols) where+ show (TLD ldf) = "TypedLazyDataFrame { " ++ show ldf ++ " }"++-- | Scan a CSV file with a given schema.+scanCsv ::+ Schema ->+ T.Text ->+ TypedLazyDataFrame cols+scanCsv schema path = TLD (L.scanCsv schema path)++-- | Scan a character-separated file with a given schema.+scanSeparated ::+ Char ->+ Schema ->+ T.Text ->+ TypedLazyDataFrame cols+scanSeparated sep schema path = TLD (L.scanSeparated sep schema path)++-- | Scan a Parquet file, directory, or glob pattern with a given schema.+scanParquet ::+ Schema ->+ T.Text ->+ TypedLazyDataFrame cols+scanParquet schema path = TLD (L.scanParquet schema path)++-- | Lift an already-loaded eager 'TypedDataFrame' into a lazy plan.+fromDataFrame :: TypedDataFrame cols -> TypedLazyDataFrame cols+fromDataFrame (TDF df) = TLD (L.fromDataFrame df)++-- | Synonym for 'fromDataFrame'.+fromTypedDataFrame :: TypedDataFrame cols -> TypedLazyDataFrame cols+fromTypedDataFrame = fromDataFrame++-- | Keep rows that satisfy the predicate.+filter :: TExpr cols Bool -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols+filter (TExpr expr) (TLD ldf) = TLD (L.filter expr ldf)++-- | Retain at most @n@ rows.+take :: Int -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols+take n (TLD ldf) = TLD (L.take n ldf)++-- | Add a computed column.+derive ::+ forall name a cols.+ (KnownSymbol name, C.Columnable a, AssertAbsent name cols) =>+ TExpr cols a ->+ TypedLazyDataFrame cols ->+ TypedLazyDataFrame (Snoc cols (Column name a))+derive (TExpr expr) (TLD ldf) =+ TLD (L.derive (T.pack (symbolVal (Proxy @name))) expr ldf)++-- | Retain only the listed columns.+select ::+ forall (names :: [Symbol]) cols.+ (AllKnownSymbol names, AssertAllPresent names cols) =>+ TypedLazyDataFrame cols ->+ TypedLazyDataFrame (SubsetSchema names cols)+select (TLD ldf) = TLD (L.select (DataFrame.Typed.Schema.symbolVals @names) ldf)++-- | A typed lazy grouped query.+newtype TypedLazyGrouped (keys :: [Symbol]) (cols :: [Type]) = TLG+ { _unTLG :: ([T.Text], LazyDataFrame)+ }++-- | Group by key columns.+groupBy ::+ forall (keys :: [Symbol]) cols.+ (AllKnownSymbol keys, AssertAllPresent keys cols) =>+ TypedLazyDataFrame cols ->+ TypedLazyGrouped keys cols+groupBy (TLD ldf) = TLG (DataFrame.Typed.Schema.symbolVals @keys, ldf)++-- | Aggregate a grouped lazy query.+aggregate ::+ forall keys cols aggs.+ TAgg keys cols aggs ->+ TypedLazyGrouped keys cols ->+ TypedLazyDataFrame (Append (GroupKeyColumns keys cols) (Reverse aggs))+aggregate tagg (TLG (keys, ldf)) =+ TLD (L.groupBy keys (aggToNamedExprs tagg) ldf)++-- | Join two lazy queries on a shared key column.+join ::+ JoinType ->+ T.Text ->+ T.Text ->+ TypedLazyDataFrame left ->+ TypedLazyDataFrame right ->+ TypedLazyDataFrame left -- TODO: compute join result schema+join jt leftKey rightKey (TLD left) (TLD right) =+ TLD (L.join jt leftKey rightKey left right)++-- | Sort the result by column name and direction.+sortBy ::+ [(T.Text, SortOrder)] ->+ TypedLazyDataFrame cols ->+ TypedLazyDataFrame cols+sortBy cols (TLD ldf) = TLD (L.sortBy cols ldf)++-- | Execute the lazy query and return a typed DataFrame.+run ::+ forall cols.+ (KnownSchema cols) =>+ TypedLazyDataFrame cols ->+ IO (TypedDataFrame cols)+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 TAggNil = []+aggToNamedExprs (TAggCons name (TExpr expr) rest) =+ (name, E.UExpr expr) : aggToNamedExprs rest