diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/dataframe-lazy.cabal b/dataframe-lazy.cabal new file mode 100644 --- /dev/null +++ b/dataframe-lazy.cabal @@ -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 diff --git a/src/DataFrame/Lazy.hs b/src/DataFrame/Lazy.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Lazy.hs @@ -0,0 +1,3 @@ +module DataFrame.Lazy (module DataFrame.Lazy.Internal.DataFrame) where + +import DataFrame.Lazy.Internal.DataFrame diff --git a/src/DataFrame/Lazy/IO/Binary.hs b/src/DataFrame/Lazy/IO/Binary.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Lazy/IO/Binary.hs @@ -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 diff --git a/src/DataFrame/Lazy/IO/CSV.hs b/src/DataFrame/Lazy/IO/CSV.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Lazy/IO/CSV.hs @@ -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 #-} diff --git a/src/DataFrame/Lazy/Internal/DataFrame.hs b/src/DataFrame/Lazy/Internal/DataFrame.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Lazy/Internal/DataFrame.hs @@ -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)} diff --git a/src/DataFrame/Lazy/Internal/Executor.hs b/src/DataFrame/Lazy/Internal/Executor.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Lazy/Internal/Executor.hs @@ -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) diff --git a/src/DataFrame/Lazy/Internal/LogicalPlan.hs b/src/DataFrame/Lazy/Internal/LogicalPlan.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Lazy/Internal/LogicalPlan.hs @@ -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 ++ " " + 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) diff --git a/src/DataFrame/Lazy/Internal/Optimizer.hs b/src/DataFrame/Lazy/Internal/Optimizer.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Lazy/Internal/Optimizer.hs @@ -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 diff --git a/src/DataFrame/Lazy/Internal/PhysicalPlan.hs b/src/DataFrame/Lazy/Internal/PhysicalPlan.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Lazy/Internal/PhysicalPlan.hs @@ -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) diff --git a/src/DataFrame/Typed/Lazy.hs b/src/DataFrame/Typed/Lazy.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Lazy.hs @@ -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