dataframe 0.6.0.0 → 0.7.0.0

raw patch · 46 files changed

+11199/−6769 lines, 46 filesdep +deepseqdep +stmdep ~randomnew-component:exe:lazy-benchPVP ok

version bump matches the API change (PVP)

Dependencies added: deepseq, stm

Dependency ranges changed: random

API changes (from Hackage documentation)

Files

• CHANGELOG.md +10/−0
• app/LazyBenchmark.hs +311/−0
• benchmark/Main.hs +164/−9
• dataframe.cabal +32/−2
• src/DataFrame/Functions.hs +32/−8
• src/DataFrame/IO/CSV.hs +191/−44
• src/DataFrame/IO/Parquet.hs +226/−116
• src/DataFrame/IO/Parquet/Binary.hs +30/−3
• src/DataFrame/IO/Parquet/Dictionary.hs +160/−44
• src/DataFrame/IO/Parquet/Encoding.hs +24/−25
• src/DataFrame/IO/Parquet/Levels.hs +79/−10
• src/DataFrame/IO/Parquet/Page.hs +25/−2
• src/DataFrame/IO/Parquet/Types.hs +8/−7
• src/DataFrame/IO/Unstable/CSV.hs +36/−19
• src/DataFrame/Internal/Column.hs +418/−5
• src/DataFrame/Internal/DataFrame.hs +11/−2
• src/DataFrame/Internal/Expression.hs +22/−1
• src/DataFrame/Internal/Interpreter.hs +478/−838
• src/DataFrame/Internal/Parsing.hs +130/−2
• src/DataFrame/Internal/Row.hs +4/−0
• src/DataFrame/Internal/Types.hs +2/−1
• src/DataFrame/Lazy/IO/Binary.hs +444/−0
• src/DataFrame/Lazy/IO/CSV.hs +231/−174
• src/DataFrame/Lazy/Internal/DataFrame.hs +110/−79
• src/DataFrame/Lazy/Internal/Executor.hs +543/−0
• src/DataFrame/Lazy/Internal/LogicalPlan.hs +44/−0
• src/DataFrame/Lazy/Internal/Optimizer.hs +209/−0
• src/DataFrame/Lazy/Internal/PhysicalPlan.hs +36/−0
• src/DataFrame/Operations/Aggregation.hs +27/−7
• src/DataFrame/Operations/Join.hs +425/−192
• src/DataFrame/Typed.hs +0/−3
• src/DataFrame/Typed/Expr.hs +16/−5
• src/DataFrame/Typed/Join.hs +2/−2
• src/DataFrame/Typed/Operations.hs +1/−12
• tests/IO/JSON.hs +112/−0
• tests/Internal/Parsing.hs +237/−0
• tests/LazyParquet.hs +131/−0
• tests/Main.hs +79/−5149
• tests/Operations/Aggregations.hs +77/−0
• tests/Operations/Apply.hs +39/−0
• tests/Operations/GroupBy.hs +2/−2
• tests/Operations/Join.hs +99/−6
• tests/Operations/Statistics.hs +61/−0
• tests/Operations/Typing.hs +5087/−0
• tests/Parquet.hs +757/−0
• tests/Properties.hs +37/−0

@@ -1,5 +1,15 @@ # Revision history for dataframe +## 0.7.0.0+* This release adds A LOT of AI code to the repo (which we'll now pause in favour of refactoring, testing, and completeness for 1.0)+* The lazy reader now has a custom binary format that it spills to. This almost halved the time it takes to run the 1 billion row challenge. The lazy evaluation now also supports Parquet.+* You can now explicitly set the schema of a subset of columns (named) as opposed to setting a list.+* Join order is changed from pipelining style to function call style.+* Bug in division being marked as commutative.+* Columnable now requires nfdata.+* Parquet reads are now faster across the board.+* Better test coverage.+ ## 0.6.0.0 * New typed API see https://dataframe.readthedocs.io/en/latest/using_dataframe_in_a_standalone_script.html * Faster joins

@@ -0,0 +1,311 @@+{-# LANGUAGE NumericUnderscores #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++{- | End-to-end smoke test and benchmark for the Lazy streaming API.++     Usage:+       cabal run lazy-bench [-- [OPTIONS]]++     Options:+       --rows N       Number of rows to generate (default: 1_000_000_000)+       --file PATH    Output CSV path          (default: /tmp/lazy_1b.csv)+       --skip-gen     Skip generation if the file already exists++     The executable generates a CSV file (streaming, constant memory) then+     runs five Lazy queries over it, printing timing and result summaries.++     For heap/GC stats run with:+       cabal run lazy-bench -- +RTS -s -RTS+-}+module Main where++import Control.Monad (foldM, forM_, when)+import qualified Data.ByteString.Builder as Builder+import qualified Data.Map as M+import qualified Data.Text as T+import Data.Time (UTCTime, diffUTCTime, getCurrentTime)+import qualified DataFrame as D+import DataFrame.Internal.Schema (Schema (..), schemaType)+import qualified DataFrame.Lazy as L+import DataFrame.Operators+import System.Directory (doesFileExist, getFileSize)+import System.Environment (getArgs)+import System.Exit (exitFailure)+import System.IO (+    BufferMode (..),+    IOMode (..),+    hFlush,+    hSetBuffering,+    stdout,+    withFile,+ )+import System.Random.Stateful++-- ---------------------------------------------------------------------------+-- Defaults+-- ---------------------------------------------------------------------------++defaultRows :: Int+defaultRows = 1_000_000_000++defaultFile :: FilePath+defaultFile = "/tmp/lazy_1b.csv"++-- Rows written per Builder flush to disk.+chunkSize :: Int+chunkSize = 500_000++-- ---------------------------------------------------------------------------+-- Argument parsing+-- ---------------------------------------------------------------------------++data Opts = Opts+    { optRows :: Int+    , optFile :: FilePath+    , optSkipGen :: Bool+    }++parseArgs :: [String] -> Either String Opts+parseArgs = go (Opts defaultRows defaultFile False)+  where+    go opts [] = Right opts+    go opts ("--rows" : n : rest) = case reads n of+        [(v, "")] -> go opts{optRows = v} rest+        _ -> Left $ "bad --rows value: " ++ n+    go opts ("--file" : p : rest) = go opts{optFile = p} rest+    go opts ("--skip-gen" : rest) = go opts{optSkipGen = True} rest+    go _ (flag : _) = Left $ "unknown flag: " ++ flag++-- ---------------------------------------------------------------------------+-- CSV generation+-- ---------------------------------------------------------------------------++{- | Write @n@ rows of schema++       id (Int), x (Double), y (Double), category (Text: A/B/C/D)++     to @path@ using a streaming Builder, keeping heap usage constant.+-}+generateCsv :: FilePath -> Int -> IO ()+generateCsv path n = do+    g <- newIOGenM =<< newStdGen+    t0 <- getCurrentTime+    withFile path WriteMode $ \h -> do+        hSetBuffering h (BlockBuffering (Just (4 * 1024 * 1024)))+        Builder.hPutBuilder h (Builder.byteString "id,x,y,category\n")+        let numChunks = n `div` chunkSize+            remainder = n `mod` chunkSize+        forM_ [0 .. numChunks - 1] $ \c -> do+            bldr <- buildChunk g (c * chunkSize) chunkSize+            Builder.hPutBuilder h bldr+            when (c `mod` 200 == 0) $ do+                let done = c * chunkSize+                    pct = (done * 100) `div` n+                putStr $ "\r  " ++ show pct ++ "% — " ++ commas done ++ " rows"+                hFlush stdout+        when (remainder > 0) $ do+            bldr <- buildChunk g (numChunks * chunkSize) remainder+            Builder.hPutBuilder h bldr+    t1 <- getCurrentTime+    putStrLn $ "\r  100% — " ++ commas n ++ " rows written in " ++ showDiff t0 t1++buildChunk :: IOGenM StdGen -> Int -> Int -> IO Builder.Builder+buildChunk g baseId count =+    foldM (\acc i -> (acc <>) <$> buildRow g baseId i) mempty [0 .. count - 1]++buildRow :: IOGenM StdGen -> Int -> Int -> IO Builder.Builder+buildRow g baseId i = do+    x <- uniformRM (0.0 :: Double, 1.0) g+    y <- uniformRM (0.0 :: Double, 1.0) g+    c <- uniformRM (0 :: Int, 3) g+    return $+        Builder.intDec (baseId + i)+            <> Builder.char7 ','+            <> buildDouble x+            <> Builder.char7 ','+            <> buildDouble y+            <> Builder.char7 ','+            <> catChar c+            <> Builder.char7 '\n'++buildDouble :: Double -> Builder.Builder+buildDouble x =+    let scaled = round (x * 1_000_000) :: Int+        whole = scaled `div` 1_000_000+        frac = scaled `mod` 1_000_000+     in Builder.intDec whole+            <> Builder.char7 '.'+            <> pad6 frac++pad6 :: Int -> Builder.Builder+pad6 n+    | n < 10 = Builder.byteString "00000" <> Builder.intDec n+    | n < 100 = Builder.byteString "0000" <> Builder.intDec n+    | n < 1000 = Builder.byteString "000" <> Builder.intDec n+    | n < 10_000 = Builder.byteString "00" <> Builder.intDec n+    | n < 100_000 = Builder.byteString "0" <> Builder.intDec n+    | otherwise = Builder.intDec n++catChar :: Int -> Builder.Builder+catChar 0 = Builder.char7 'A'+catChar 1 = Builder.char7 'B'+catChar 2 = Builder.char7 'C'+catChar _ = Builder.char7 'D'++-- ---------------------------------------------------------------------------+-- Lazy queries+-- ---------------------------------------------------------------------------++runQuery :: String -> IO D.DataFrame -> IO ()+runQuery label action = do+    putStrLn $ "\n── " ++ label+    t0 <- getCurrentTime+    df <- action+    t1 <- getCurrentTime+    let (rows, cols) = D.dimensions df+    putStrLn $ "   rows returned : " ++ commas rows+    putStrLn $ "   columns       : " ++ show cols+    putStrLn $ "   time          : " ++ showDiff t0 t1+    when (rows > 0 && rows <= 30) $ print df++-- ---------------------------------------------------------------------------+-- Main+-- ---------------------------------------------------------------------------++main :: IO ()+main = do+    hSetBuffering stdout LineBuffering+    args <- getArgs+    opts <- case parseArgs args of+        Left err -> putStrLn ("Error: " ++ err) >> exitFailure+        Right o -> return o++    let path = optFile opts+        n = optRows opts+        pathT = T.pack path++    -- -----------------------------------------------------------------------+    -- Phase 1: Generate+    -- -----------------------------------------------------------------------+    putStrLn "=== Lazy API 1B-row benchmark ==="+    putStrLn $ "    rows   : " ++ commas n+    putStrLn $ "    file   : " ++ path+    putStrLn "    tip    : run with '+RTS -s -RTS' for heap stats"+    putStrLn ""++    exists <- doesFileExist path+    if optSkipGen opts && exists+        then do+            sz <- getFileSize path+            putStrLn $ "Skipping generation — file exists (" ++ showBytes sz ++ ")"+        else do+            putStrLn "Phase 1: Generating CSV …"+            generateCsv path n+            sz <- getFileSize path+            putStrLn $ "  file size: " ++ showBytes sz++    -- -----------------------------------------------------------------------+    -- Phase 2: Lazy queries+    -- -----------------------------------------------------------------------+    putStrLn "\nPhase 2: Lazy queries"++    -- Schema for the generated CSV: id (Int), x (Double), y (Double), category (Text)+    let schema =+            Schema $+                M.fromList+                    [ ("id", schemaType @Int)+                    , ("x", schemaType @Double)+                    , ("y", schemaType @Double)+                    , ("category", schemaType @T.Text)+                    ]++    -- Q1: Preview — limit 20, no filter.+    -- Demonstrates that the executor reads only the first batch.+    runQuery "Q1 — preview first 20 rows (no filter)" $+        L.runDataFrame $+            L.limit 20 $+                L.scanCsv schema pathT++    -- Q2: Filter + limit.+    -- x > 0.999 ≈ 0.1% of rows. With a 512K-row batch the executor finds+    -- ~512 matches in the first batch and stops — reads only one batch.+    runQuery "Q2 — filter (x > 0.999), limit 20" $+        L.runDataFrame $+            L.limit 20 $+                L.filter (col @Double "x" .> lit 0.999) $+                    L.scanCsv schema pathT++    -- Q3: Filter + derive + select + limit.+    -- Shows projection pushdown: only id/x/y/category are read, z is derived.+    -- Predicate pushdown moves the filter into the scan batch loop.+    runQuery "Q3 — filter (x > 0.999), derive z = x*y, select [id,z], limit 20" $+        L.runDataFrame $+            L.limit 20 $+                L.select ["id", "z"] $+                    L.derive "z" (col @Double "x" * col @Double "y") $+                        L.filter (col @Double "x" .> lit 0.999) $+                            L.scanCsv schema pathT++    -- Q4: Filter fusion demo.+    -- Two consecutive filters are fused into one AND predicate by the optimizer.+    -- Result: rows where x > 0.5 AND y > 0.5 (≈ 25% of total).+    -- We limit to keep result size manageable.+    runQuery "Q4 — filter fusion: (x > 0.5) . (y > 0.5), limit 20" $+        L.runDataFrame $+            L.limit 20 $+                L.filter (col @Double "y" .> lit 0.5) $+                    L.filter (col @Double "x" .> lit 0.5) $+                        L.scanCsv schema pathT++    -- Q5: Full scan, heavy filter, count results.+    -- x > 0.999 across the whole file ≈ 0.1% × N rows.+    -- For 1B rows that is ~1M results — materialised into one DataFrame.+    -- This query exercises streaming across all batches.+    runQuery+        ( "Q5 — full scan, filter (x > 0.999), count (~"+            ++ approx (n `div` 1000)+            ++ " rows expected)"+        )+        $ L.runDataFrame+        $ L.select ["id", "x"]+        $ L.filter (col @Double "x" .> lit 0.999)+        $ L.scanCsv schema pathT++    putStrLn "\nDone."++-- ---------------------------------------------------------------------------+-- Formatting helpers+-- ---------------------------------------------------------------------------++showDiff :: UTCTime -> UTCTime -> String+showDiff t0 t1 = show (diffUTCTime t1 t0)++commas :: Int -> String+commas n+    | n < 1000 = show n+    | otherwise = commas (n `div` 1000) ++ "," ++ pad3 (n `mod` 1000)+  where+    pad3 x+        | x < 10 = "00" ++ show x+        | x < 100 = "0" ++ show x+        | otherwise = show x++approx :: Int -> String+approx n+    | n >= 1_000_000 = show (n `div` 1_000_000) ++ "M"+    | n >= 1_000 = show (n `div` 1_000) ++ "K"+    | otherwise = show n++showBytes :: Integer -> String+showBytes b+    | b >= 1_073_741_824 = fmt (fromIntegral b / 1_073_741_824) ++ " GiB"+    | b >= 1_048_576 = fmt (fromIntegral b / 1_048_576) ++ " MiB"+    | b >= 1_024 = fmt (fromIntegral b / 1_024) ++ " KiB"+    | otherwise = show b ++ " B"+  where+    fmt :: Double -> String+    fmt x =+        show (fromIntegral (round (x * 10) :: Int) `div` 10 :: Int)+            ++ "."+            ++ show (fromIntegral (round (x * 10) :: Int) `mod` 10 :: Int)

@@ -1,3 +1,4 @@+{-# LANGUAGE NumericUnderscores #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TypeApplications #-} @@ -6,8 +7,10 @@  import Control.Monad (void) import Criterion.Main+import DataFrame.Operations.Join import DataFrame.Operators-import System.Process+import System.Process hiding (env)+import System.Random.Stateful  haskell :: IO () haskell = do@@ -99,6 +102,101 @@ parseMeasurementsTXT :: IO () parseMeasurementsTXT = parseFile "./data/measurements.txt" +{- | Generate a pair of dataframes for a 1:1 join scenario.+  Left has keys [0..n-1], right has keys [0..n-1] shuffled.+  overlap controls what fraction of keys appear in both sides.+-}+mkOneToOne :: Int -> Double -> IO (D.DataFrame, D.DataFrame)+mkOneToOne n overlap = do+    g <- newIOGenM =<< newStdGen+    let rightSize = max 1 (round (fromIntegral n * overlap))+    -- Left: keys 0..n-1 with a payload column+    let leftKeys = [0 :: Int .. n - 1]+        leftVals = [0 :: Int .. n - 1]+        leftDf =+            D.fromNamedColumns+                [ ("key", D.fromList leftKeys)+                , ("A", D.fromList leftVals)+                ]+    -- Right: take first `rightSize` keys, add non-overlapping keys for the rest+    rightPayload <- mapM (\_ -> uniformRM (0 :: Int, 1_000_000) g) [1 .. rightSize]+    let rightKeys = [0 :: Int .. rightSize - 1]+        rightDf =+            D.fromNamedColumns+                [ ("key", D.fromList rightKeys)+                , ("B", D.fromList rightPayload)+                ]+    return (leftDf, rightDf)++{- | Generate a pair of dataframes for a many-to-many join scenario.+  Keys are drawn from [0..cardinality-1], so rows share keys.+-}+mkManyToMany :: Int -> Int -> Int -> IO (D.DataFrame, D.DataFrame)+mkManyToMany leftRows rightRows cardinality = do+    g <- newIOGenM =<< newStdGen+    leftKeys <- mapM (\_ -> uniformRM (0 :: Int, cardinality - 1) g) [1 .. leftRows]+    rightKeys <-+        mapM (\_ -> uniformRM (0 :: Int, cardinality - 1) g) [1 .. rightRows]+    leftVals <- mapM (\_ -> uniformRM (0 :: Int, 1_000_000) g) [1 .. leftRows]+    rightVals <- mapM (\_ -> uniformRM (0 :: Int, 1_000_000) g) [1 .. rightRows]+    let leftDf =+            D.fromNamedColumns+                [ ("key", D.fromList leftKeys)+                , ("A", D.fromList leftVals)+                ]+        rightDf =+            D.fromNamedColumns+                [ ("key", D.fromList rightKeys)+                , ("B", D.fromList rightVals)+                ]+    return (leftDf, rightDf)++{- | Generate a pair of dataframes for a many-to-one join+  (fact table joining a dimension table).+  Left has n rows with keys drawn from [0..dimSize-1].+  Right has exactly dimSize rows with unique keys.+-}+mkManyToOne :: Int -> Int -> IO (D.DataFrame, D.DataFrame)+mkManyToOne factRows dimSize = do+    g <- newIOGenM =<< newStdGen+    factKeys <- mapM (\_ -> uniformRM (0 :: Int, dimSize - 1) g) [1 .. factRows]+    factVals <- mapM (\_ -> uniformRM (0 :: Int, 1_000_000) g) [1 .. factRows]+    dimVals <- mapM (\_ -> uniformRM (0 :: Int, 1_000_000) g) [1 .. dimSize]+    let factDf =+            D.fromNamedColumns+                [ ("key", D.fromList factKeys)+                , ("A", D.fromList factVals)+                ]+        dimDf =+            D.fromNamedColumns+                [ ("key", D.fromList [0 :: Int .. dimSize - 1])+                , ("B", D.fromList dimVals)+                ]+    return (factDf, dimDf)++mkMultiKey :: Int -> Int -> IO (D.DataFrame, D.DataFrame)+mkMultiKey leftRows rightRows = do+    g <- newIOGenM =<< newStdGen+    lk1 <- mapM (\_ -> uniformRM (0 :: Int, 99) g) [1 .. leftRows]+    lk2 <- mapM (\_ -> uniformRM (0 :: Int, 99) g) [1 .. leftRows]+    rk1 <- mapM (\_ -> uniformRM (0 :: Int, 99) g) [1 .. rightRows]+    rk2 <- mapM (\_ -> uniformRM (0 :: Int, 99) g) [1 .. rightRows]+    lv <- mapM (\_ -> uniformRM (0 :: Int, 1_000_000) g) [1 .. leftRows]+    rv <- mapM (\_ -> uniformRM (0 :: Int, 1_000_000) g) [1 .. rightRows]+    let leftDf =+            D.fromNamedColumns+                [ ("key1", D.fromList lk1)+                , ("key2", D.fromList lk2)+                , ("A", D.fromList lv)+                ]+        rightDf =+            D.fromNamedColumns+                [ ("key1", D.fromList rk1)+                , ("key2", D.fromList rk2)+                , ("B", D.fromList rv)+                ]+    return (leftDf, rightDf)+ main :: IO () main = do     output <- readProcess "cabal" ["build", "-O2"] ""@@ -135,12 +233,69 @@                     parseFileUnstableSIMD                         "./data/effects-of-covid-19-on-trade-at-15-december-2021-provisional.csv"             ]-            -- this file was removed as it's too large you may substitute-            -- with your own if you wish to run it locally on a large file-            -- , bgroup-            --     "customers.csv (334 MB)"-            --     [ bench "Attoparsec" $ nfIO $ parseFile "./data/customers.csv"-            --     , bench "Native Haskell" $ nfIO $ parseFileUnstable "./data/customers.csv"-            --     , bench "SIMD" $ nfIO $ parseFileUnstableSIMD "./data/customers.csv"-            --     ]+        , bgroup+            "join/inner/1:1"+            [ env (mkOneToOne 1_000 1.0) $ \ ~(l, r) ->+                bench "1K rows" $ nf (innerJoin ["key"] r) l+            , env (mkOneToOne 10_000 1.0) $ \ ~(l, r) ->+                bench "10K rows" $ nf (innerJoin ["key"] r) l+            , env (mkOneToOne 100_000 1.0) $ \ ~(l, r) ->+                bench "100K rows" $ nf (innerJoin ["key"] r) l+            ]+        , bgroup+            "join/inner/1:1-partial-overlap"+            [ env (mkOneToOne 100_000 0.5) $ \ ~(l, r) ->+                bench "100K rows, 50% overlap" $ nf (innerJoin ["key"] r) l+            , env (mkOneToOne 100_000 0.1) $ \ ~(l, r) ->+                bench "100K rows, 10% overlap" $ nf (innerJoin ["key"] r) l+            ]+        , bgroup+            "join/inner/many:1"+            [ env (mkManyToOne 10_000 100) $ \ ~(fact, dim) ->+                bench "10K fact x 100 dim" $ nf (innerJoin ["key"] dim) fact+            , env (mkManyToOne 100_000 1_000) $ \ ~(fact, dim) ->+                bench "100K fact x 1K dim" $ nf (innerJoin ["key"] dim) fact+            , env (mkManyToOne 100_000 100) $ \ ~(fact, dim) ->+                bench "100K fact x 100 dim" $ nf (innerJoin ["key"] dim) fact+            ]+        , bgroup+            "join/inner/many:many"+            [ env (mkManyToMany 1_000 1_000 100) $ \ ~(l, r) ->+                bench "1Kx1K, 100 keys" $ nf (innerJoin ["key"] r) l+            , env (mkManyToMany 10_000 10_000 1_000) $ \ ~(l, r) ->+                bench "10Kx10K, 1K keys" $ nf (innerJoin ["key"] r) l+            , env (mkManyToMany 10_000 10_000 100) $ \ ~(l, r) ->+                bench "10Kx10K, 100 keys" $ nf (innerJoin ["key"] r) l+            ]+        , bgroup+            "join/left"+            [ env (mkOneToOne 10_000 1.0) $ \ ~(l, r) ->+                bench "1:1, 10K rows" $ nf (leftJoin ["key"] r) l+            , env (mkOneToOne 100_000 1.0) $ \ ~(l, r) ->+                bench "1:1, 100K rows" $ nf (leftJoin ["key"] r) l+            , env (mkOneToOne 100_000 0.5) $ \ ~(l, r) ->+                bench "1:1, 100K rows, 50%" $ nf (leftJoin ["key"] r) l+            , env (mkManyToOne 100_000 1_000) $ \ ~(fact, dim) ->+                bench "many:1, 100K x 1K" $ nf (leftJoin ["key"] dim) fact+            , env (mkManyToMany 10_000 10_000 1_000) $ \ ~(l, r) ->+                bench "many:many, 10Kx10K" $ nf (leftJoin ["key"] r) l+            ]+        , bgroup+            "join/fullOuter"+            [ env (mkOneToOne 10_000 1.0) $ \ ~(l, r) ->+                bench "1:1, 10K rows" $ nf (fullOuterJoin ["key"] r) l+            , env (mkOneToOne 100_000 1.0) $ \ ~(l, r) ->+                bench "1:1, 100K rows" $ nf (fullOuterJoin ["key"] r) l+            , env (mkOneToOne 100_000 0.5) $ \ ~(l, r) ->+                bench "1:1, 100K rows, 50%" $ nf (fullOuterJoin ["key"] r) l+            , env (mkManyToOne 100_000 1_000) $ \ ~(fact, dim) ->+                bench "many:1, 100K x 1K" $ nf (fullOuterJoin ["key"] dim) fact+            ]+        , bgroup+            "join/multiKey"+            [ env (mkMultiKey 10_000 10_000) $ \ ~(l, r) ->+                bench "inner 10Kx10K, 2 keys" $ nf (innerJoin ["key1", "key2"] r) l+            , env (mkMultiKey 10_000 10_000) $ \ ~(l, r) ->+                bench "left 10Kx10K, 2 keys" $ nf (leftJoin ["key1", "key2"] r) l+            ]         ]

@@ -1,6 +1,6 @@ cabal-version:      2.4 name:               dataframe-version:            0.6.0.0+version:            0.7.0.0  synopsis: A fast, safe, and intuitive DataFrame library. @@ -89,7 +89,12 @@                     DataFrame.IO.Parquet.Time,                     DataFrame.IO.Parquet.Types,                     DataFrame.Lazy.IO.CSV,+                    DataFrame.Lazy.IO.Binary,                     DataFrame.Lazy.Internal.DataFrame,+                    DataFrame.Lazy.Internal.LogicalPlan,+                    DataFrame.Lazy.Internal.PhysicalPlan,+                    DataFrame.Lazy.Internal.Optimizer,+                    DataFrame.Lazy.Internal.Executor,                     DataFrame.Monad,                     DataFrame.DecisionTree,                     DataFrame.Typed.Types,@@ -103,6 +108,7 @@                     DataFrame.Typed.Expr,                     DataFrame.Typed     build-depends:    base >= 4 && <5,+                      deepseq >= 1 && < 2,                       aeson >= 0.11.0.0 && < 3,                       array >= 0.5.4.0 && < 0.6,                       attoparsec >= 0.12 && < 0.15,@@ -129,6 +135,7 @@                       zstd >= 0.1.2.0 && < 0.2,                       mmap >= 0.5.8 && < 0.6,                       parallel >= 3.2.2.0 && < 5,+                      stm >= 2.5 && < 3,                       filepath >= 1.4 && < 2,                       Glob >= 0.10 && < 1, @@ -176,6 +183,21 @@     default-language: Haskell2010     ghc-options: -rtsopts -threaded -with-rtsopts=-N +executable lazy-bench+    import: warnings+    main-is: LazyBenchmark.hs+    build-depends:    base >= 4 && < 5,+                      bytestring >= 0.11 && < 0.13,+                      containers >= 0.6.7 && < 0.9,+                      dataframe >= 0.5 && < 1,+                      directory >= 1.3.0.0 && < 2,+                      random >= 1 && < 2,+                      text >= 2.0 && < 3,+                      time >= 1.12 && < 2,+    hs-source-dirs:   app+    default-language: Haskell2010+    ghc-options: -rtsopts -threaded -with-rtsopts=-N+ benchmark dataframe-benchmark     import: warnings     type: exitcode-stdio-1.0@@ -184,7 +206,8 @@     build-depends: base >= 4 && < 5,                    criterion >= 1 && < 2,                    process >= 1.6 && < 2,-                   dataframe >= 0.5 && < 1+                   dataframe >= 0.5 && < 1,+                   random >= 1 && < 2,     default-language: Haskell2010     ghc-options:       -threaded@@ -198,6 +221,8 @@     other-modules: Assertions,                    Functions,                    GenDataFrame,+                   Internal.Parsing,+                   IO.JSON,                    Operations.Aggregations,                    Operations.Apply,                    Operations.Core,@@ -213,9 +238,13 @@                    Operations.Subset,                    Operations.Statistics,                    Operations.Take,+                   Operations.Typing,+                   LazyParquet,                    Parquet,+                   Properties,                    Monad     build-depends:  base >= 4 && < 5,+                    bytestring >= 0.11 && < 0.13,                     dataframe >= 0.5 && < 1,                     directory >= 1.3.0.0 && < 2,                     HUnit ^>= 1.6,@@ -229,3 +258,4 @@                     containers >= 0.6.7 && < 0.9     hs-source-dirs: tests     default-language: Haskell2010+

@@ -18,6 +18,7 @@ import DataFrame.Internal.Column import DataFrame.Internal.DataFrame (     DataFrame (..),+    empty,     unsafeGetColumn,  ) import DataFrame.Internal.Expression hiding (normalize)@@ -46,6 +47,9 @@ import Debug.Trace (trace) import Language.Haskell.TH import qualified Language.Haskell.TH.Syntax as TH+import System.Directory (doesDirectoryExist)+import System.FilePath ((</>))+import System.FilePath.Glob (glob) import Text.Regex.TDFA import Prelude hiding (maximum, minimum) import Prelude as P@@ -138,7 +142,7 @@         (MkUnaryOp{unaryFn = Prelude.not, unaryName = "not", unarySymbol = Just "~"})  count :: (Columnable a) => Expr a -> Expr Int-count = Agg (FoldAgg "count" (Just 0) (\acc _ -> acc + 1))+count = Agg (MergeAgg "count" (0 :: Int) (\c _ -> c + 1) (+) id)  collect :: (Columnable a) => Expr a -> Expr [a] collect = Agg (FoldAgg "collect" (Just []) (flip (:)))@@ -170,11 +174,16 @@ sumMaybe :: forall a. (Columnable a, Num a) => Expr (Maybe a) -> Expr a sumMaybe = Agg (CollectAgg "sumMaybe" (P.sum . Maybe.catMaybes . V.toList)) -mean :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double-mean = Agg (CollectAgg "mean" mean')-{-# SPECIALIZE DataFrame.Functions.mean :: Expr Double -> Expr Double #-}-{-# SPECIALIZE DataFrame.Functions.mean :: Expr Int -> Expr Double #-}-{-# INLINEABLE DataFrame.Functions.mean #-}+mean :: (Columnable a, Real a) => Expr a -> Expr Double+mean =+    Agg+        ( MergeAgg+            "mean"+            (MeanAcc 0.0 0)+            (\(MeanAcc s c) x -> MeanAcc (s + realToFrac x) (c + 1))+            (\(MeanAcc s1 c1) (MeanAcc s2 c2) -> MeanAcc (s1 + s2) (c1 + c2))+            (\(MeanAcc s c) -> if c == 0 then 0 / 0 else s / fromIntegral c)+        )  meanMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double meanMaybe = Agg (CollectAgg "meanMaybe" (mean' . optionalToDoubleVector))@@ -442,8 +451,23 @@  declareColumnsFromParquetFile :: String -> DecsQ declareColumnsFromParquetFile path = do-    metadata <- liftIO (Parquet.readMetadataFromPath path)-    let df = schemaToEmptyDataFrame (schema metadata)+    isDir <- liftIO $ doesDirectoryExist path++    let pat = if isDir then path </> "*.parquet" else path++    matches <- liftIO $ glob pat++    files <- liftIO $ filterM (fmap Prelude.not . doesDirectoryExist) matches+    df <-+        liftIO $+            foldM+                ( \acc p -> do+                    (metadata, _) <- liftIO (Parquet.readMetadataFromPath path)+                    let d = schemaToEmptyDataFrame (schema metadata)+                    pure $ acc <> d+                )+                DataFrame.Internal.DataFrame.empty+                files     declareColumns df  schemaToEmptyDataFrame :: [SchemaElement] -> DataFrame

@@ -10,6 +10,8 @@  module DataFrame.IO.CSV where +import qualified Data.ByteString as BS+import qualified Data.ByteString.Char8 as C import qualified Data.ByteString.Lazy as BL import qualified Data.List as L import qualified Data.Map.Strict as M@@ -66,6 +68,7 @@     = BuilderInt !(PagedUnboxedVector Int) !(PagedUnboxedVector Word8)     | BuilderDouble !(PagedUnboxedVector Double) !(PagedUnboxedVector Word8)     | BuilderText !(PagedVector T.Text) !(PagedUnboxedVector Word8)+    | BuilderBS !(PagedVector BS.ByteString) !(PagedUnboxedVector Word8)  newPagedVector :: IO (PagedVector a) newPagedVector = do@@ -123,10 +126,21 @@     active <- readIORef activeRef     chunks <- readIORef chunksRef -    lastChunk <- V.unsafeFreeze (VM.slice 0 count active)+    writeIORef chunksRef [] -- release chunk references+    let frozenChunks = reverse chunks+        totalLen = count + sum (map V.length frozenChunks) -    return $! V.concat (reverse (lastChunk : chunks))+    mv <- VM.unsafeNew totalLen +    let copyChunk !offset chunk = do+            V.copy (VM.slice offset (V.length chunk) mv) chunk+            pure (offset + V.length chunk)++    offset <- foldM copyChunk 0 frozenChunks+    VM.copy (VM.slice offset count mv) (VM.slice 0 count active)++    V.unsafeFreeze mv+ freezePagedUnboxedVector ::     (VUM.Unbox a) => PagedUnboxedVector a -> IO (VU.Vector a) freezePagedUnboxedVector (PagedUnboxedVector chunksRef activeRef countRef) = do@@ -134,14 +148,29 @@     active <- readIORef activeRef     chunks <- readIORef chunksRef -    lastChunk <- VU.unsafeFreeze (VUM.slice 0 count active)-    return $! VU.concat (reverse (lastChunk : chunks))+    writeIORef chunksRef [] -- release chunk references+    let frozenChunks = reverse chunks+        totalLen = count + sum (map VU.length frozenChunks) +    mv <- VUM.unsafeNew totalLen++    let copyChunk !offset chunk = do+            VU.copy (VUM.slice offset (VU.length chunk) mv) chunk+            pure (offset + VU.length chunk)++    offset <- foldM copyChunk 0 frozenChunks+    VUM.copy (VUM.slice offset count mv) (VUM.slice 0 count active)++    VU.unsafeFreeze mv+ -- | STANDARD CONFIG TYPES data HeaderSpec = NoHeader | UseFirstRow | ProvideNames [T.Text]     deriving (Eq, Show) -data TypeSpec = InferFromSample Int | SpecifyTypes [SchemaType] | NoInference+data TypeSpec+    = InferFromSample Int+    | SpecifyTypes [(T.Text, SchemaType)]+    | NoInference  -- | CSV read parameters. data ReadOptions = ReadOptions@@ -175,9 +204,9 @@ shouldInferFromSample (InferFromSample _) = True shouldInferFromSample _ = False -schemaTypes :: TypeSpec -> [SchemaType]-schemaTypes (SpecifyTypes xs) = xs-schemaTypes _ = []+schemaTypeMap :: TypeSpec -> M.Map T.Text SchemaType+schemaTypeMap (SpecifyTypes xs) = M.fromList xs+schemaTypeMap _ = M.empty  typeInferenceSampleSize :: TypeSpec -> Int typeInferenceSampleSize (InferFromSample n) = n@@ -270,46 +299,42 @@                 )      (sampleRow, _) <- peekStream rowsToProcess-    builderCols <- initializeColumns (V.toList sampleRow) opts+    builderCols <- initializeColumns columnNames (V.toList sampleRow) opts+    let !builderColsV = V.fromList builderCols     processStream         (missingIndicators opts)         rowsToProcess-        builderCols+        builderColsV         (numColumns opts) -    frozenCols <- V.fromList <$> mapM freezeBuilderColumn builderCols+    frozenCols <- V.mapM (finalizeBuilderColumn opts) builderColsV     let numRows = maybe 0 columnLength (frozenCols V.!? 0) -    let df =-            DataFrame-                frozenCols-                (M.fromList (zip columnNames [0 ..]))-                (numRows, V.length frozenCols)-                M.empty -- TODO give typed column references     return $-        if shouldInferFromSample (typeSpec opts)-            then-                parseDefaults-                    (missingIndicators opts)-                    (typeInferenceSampleSize (typeSpec opts))-                    (safeRead opts)-                    (dateFormat opts)-                    df-            else-                if not (null (schemaTypes (typeSpec opts)))-                    then parseWithTypes (schemaTypes (typeSpec opts)) df-                    else df+        DataFrame+            frozenCols+            (M.fromList (zip columnNames [0 ..]))+            (numRows, V.length frozenCols)+            M.empty -- TODO give typed column references -initializeColumns :: [BL.ByteString] -> ReadOptions -> IO [BuilderColumn]-initializeColumns row opts = case typeSpec opts of-    NoInference -> zipWithM initColumn row (expandTypes [])-    InferFromSample _ -> zipWithM initColumn row (expandTypes [])-    SpecifyTypes ts -> zipWithM initColumn row (expandTypes ts)+initializeColumns ::+    [T.Text] -> [BL.ByteString] -> ReadOptions -> IO [BuilderColumn]+initializeColumns names row opts = zipWithM initColumn names (map lookupType names)   where-    expandTypes xs = xs ++ replicate (length row - length xs) (schemaType @T.Text)-    initColumn :: BL.ByteString -> SchemaType -> IO BuilderColumn-    initColumn _ t = do+    typeMap = schemaTypeMap (typeSpec opts)+    -- Return Nothing for columns that should be inferred from BS+    shouldInfer = case typeSpec opts of+        InferFromSample _ -> True+        SpecifyTypes _ -> True+        NoInference -> False+    lookupType name = M.lookup name typeMap+    initColumn :: T.Text -> Maybe SchemaType -> IO BuilderColumn+    initColumn _ Nothing | shouldInfer = do         validityRef <- newPagedUnboxedVector+        BuilderBS <$> newPagedVector <*> pure validityRef+    initColumn _ mtype = do+        validityRef <- newPagedUnboxedVector+        let t = fromMaybe (schemaType @T.Text) mtype         case t of             SType (_ :: P.Proxy a) -> case testEquality (typeRep @a) (typeRep @Int) of                 Just Refl -> BuilderInt <$> newPagedUnboxedVector <*> pure validityRef@@ -320,7 +345,7 @@ processStream ::     [T.Text] ->     CsvStream.Records (V.Vector BL.ByteString) ->-    [BuilderColumn] ->+    V.Vector BuilderColumn ->     Maybe Int ->     IO () processStream _ _ _ (Just 0) = return ()@@ -330,11 +355,12 @@ processStream missing (Cons (Left err) _) _ _ = error ("CSV Parse Error: " ++ err) processStream missing (Nil _ _) _ _ = return () -processRow :: [T.Text] -> V.Vector BL.ByteString -> [BuilderColumn] -> IO ()-processRow missing !vals !cols = V.zipWithM_ processValue vals (V.fromList cols)+processRow ::+    [T.Text] -> V.Vector BL.ByteString -> V.Vector BuilderColumn -> IO ()+processRow missing !vals !cols = V.zipWithM_ processValue vals cols   where     processValue !bs !col = do-        let bs' = BL.toStrict bs+        let !bs' = BL.toStrict bs         case col of             BuilderInt gv valid -> case readByteStringInt bs' of                 Just !i -> appendPagedUnboxedVector gv i >> appendPagedUnboxedVector valid 1@@ -347,26 +373,147 @@                 if isNullish val || val `elem` missing                     then appendPagedVector gv T.empty >> appendPagedUnboxedVector valid 0                     else appendPagedVector gv val >> appendPagedUnboxedVector valid 1+            BuilderBS gv valid -> do+                let !bs'' = C.strip bs'+                if isNullishBS bs'' || TE.decodeUtf8Lenient bs'' `elem` missing+                    then appendPagedVector gv BS.empty >> appendPagedUnboxedVector valid 0+                    else appendPagedVector gv bs'' >> appendPagedUnboxedVector valid 1  freezeBuilderColumn :: BuilderColumn -> IO Column freezeBuilderColumn (BuilderInt gv validRef) = do     vec <- freezePagedUnboxedVector gv     valid <- freezePagedUnboxedVector validRef     if VU.all (== 1) valid-        then return $ UnboxedColumn vec+        then return $! UnboxedColumn vec         else constructOptional vec valid freezeBuilderColumn (BuilderDouble gv validRef) = do     vec <- freezePagedUnboxedVector gv     valid <- freezePagedUnboxedVector validRef     if VU.all (== 1) valid-        then return $ UnboxedColumn vec+        then return $! UnboxedColumn vec         else constructOptional vec valid freezeBuilderColumn (BuilderText gv validRef) = do     vec <- freezePagedVector gv     valid <- freezePagedUnboxedVector validRef     if VU.all (== 1) valid-        then return $ BoxedColumn vec+        then return $! BoxedColumn vec         else constructOptionalBoxed vec valid+freezeBuilderColumn (BuilderBS _ _) =+    error+        "freezeBuilderColumn: BuilderBS must be finalized via finalizeBuilderColumn"++finalizeBuilderColumn :: ReadOptions -> BuilderColumn -> IO Column+finalizeBuilderColumn opts (BuilderBS gv validRef) = do+    vec <- freezePagedVector gv+    valid <- freezePagedUnboxedVector validRef+    return $! inferColumnFromBS opts vec valid+finalizeBuilderColumn _ bc = freezeBuilderColumn bc++inferColumnFromBS ::+    ReadOptions -> V.Vector BS.ByteString -> VU.Vector Word8 -> Column+inferColumnFromBS opts vec valid =+    let sampleN = let n = typeInferenceSampleSize (typeSpec opts) in if n == 0 then 100 else n+        dfmt = dateFormat opts+        asMaybeFull = V.generate (V.length vec) $ \i ->+            if valid VU.! i == 1 then Just (vec V.! i) else Nothing+        samples = V.take sampleN asMaybeFull+        assumption = makeParsingAssumptionBS dfmt samples+     in case assumption of+            IntAssumption -> handleBSInt dfmt asMaybeFull+            DoubleAssumption -> handleBSDouble asMaybeFull+            BoolAssumption -> handleBSBool asMaybeFull+            DateAssumption -> handleBSDate dfmt asMaybeFull+            TextAssumption -> handleBSText asMaybeFull+            NoAssumption -> handleBSNo dfmt asMaybeFull++makeParsingAssumptionBS ::+    String -> V.Vector (Maybe BS.ByteString) -> ParsingAssumption+makeParsingAssumptionBS dfmt asMaybe+    | V.all (== Nothing) asMaybe = NoAssumption+    | vecSameConstructor asMaybe asMaybeBool = BoolAssumption+    | vecSameConstructor asMaybe asMaybeInt+        && vecSameConstructor asMaybe asMaybeDouble =+        IntAssumption+    | vecSameConstructor asMaybe asMaybeDouble = DoubleAssumption+    | vecSameConstructor asMaybe asMaybeDate = DateAssumption+    | otherwise = TextAssumption+  where+    asMaybeBool = V.map (>>= readByteStringBool) asMaybe+    asMaybeInt = V.map (>>= readByteStringInt) asMaybe+    asMaybeDouble = V.map (>>= readByteStringDouble) asMaybe+    asMaybeDate = V.map (>>= readByteStringDate dfmt) asMaybe++handleBSBool :: V.Vector (Maybe BS.ByteString) -> Column+handleBSBool asMaybe+    | parsableAsBool =+        maybe (fromVector asMaybeBool) fromVector (sequenceA asMaybeBool)+    | otherwise = handleBSText asMaybe+  where+    asMaybeBool = V.map (>>= readByteStringBool) asMaybe+    parsableAsBool = vecSameConstructor asMaybe asMaybeBool++handleBSInt :: String -> V.Vector (Maybe BS.ByteString) -> Column+handleBSInt dfmt asMaybe+    | parsableAsInt =+        maybe (fromVector asMaybeInt) fromVector (sequenceA asMaybeInt)+    | parsableAsDouble =+        maybe (fromVector asMaybeDouble) fromVector (sequenceA asMaybeDouble)+    | otherwise = handleBSText asMaybe+  where+    asMaybeInt = V.map (>>= readByteStringInt) asMaybe+    asMaybeDouble = V.map (>>= readByteStringDouble) asMaybe+    parsableAsInt =+        vecSameConstructor asMaybe asMaybeInt+            && vecSameConstructor asMaybe asMaybeDouble+    parsableAsDouble = vecSameConstructor asMaybe asMaybeDouble++handleBSDouble :: V.Vector (Maybe BS.ByteString) -> Column+handleBSDouble asMaybe+    | parsableAsDouble =+        maybe (fromVector asMaybeDouble) fromVector (sequenceA asMaybeDouble)+    | otherwise = handleBSText asMaybe+  where+    asMaybeDouble = V.map (>>= readByteStringDouble) asMaybe+    parsableAsDouble = vecSameConstructor asMaybe asMaybeDouble++handleBSDate :: String -> V.Vector (Maybe BS.ByteString) -> Column+handleBSDate dfmt asMaybe+    | parsableAsDate =+        maybe (fromVector asMaybeDate) fromVector (sequenceA asMaybeDate)+    | otherwise = handleBSText asMaybe+  where+    asMaybeDate = V.map (>>= readByteStringDate dfmt) asMaybe+    parsableAsDate = vecSameConstructor asMaybe asMaybeDate++handleBSText :: V.Vector (Maybe BS.ByteString) -> Column+handleBSText asMaybe =+    let asMaybeText = V.map (fmap TE.decodeUtf8Lenient) asMaybe+     in maybe (fromVector asMaybeText) fromVector (sequenceA asMaybeText)++handleBSNo :: String -> V.Vector (Maybe BS.ByteString) -> Column+handleBSNo dfmt asMaybe+    | V.all (== Nothing) asMaybe =+        fromVector (V.map (const (Nothing :: Maybe T.Text)) asMaybe)+    | parsableAsBool =+        maybe (fromVector asMaybeBool) fromVector (sequenceA asMaybeBool)+    | parsableAsInt =+        maybe (fromVector asMaybeInt) fromVector (sequenceA asMaybeInt)+    | parsableAsDouble =+        maybe (fromVector asMaybeDouble) fromVector (sequenceA asMaybeDouble)+    | parsableAsDate =+        maybe (fromVector asMaybeDate) fromVector (sequenceA asMaybeDate)+    | otherwise = handleBSText asMaybe+  where+    asMaybeBool = V.map (>>= readByteStringBool) asMaybe+    asMaybeInt = V.map (>>= readByteStringInt) asMaybe+    asMaybeDouble = V.map (>>= readByteStringDouble) asMaybe+    asMaybeDate = V.map (>>= readByteStringDate dfmt) asMaybe+    parsableAsBool = vecSameConstructor asMaybe asMaybeBool+    parsableAsInt =+        vecSameConstructor asMaybe asMaybeInt+            && vecSameConstructor asMaybe asMaybeDouble+    parsableAsDouble = vecSameConstructor asMaybe asMaybeDouble+    parsableAsDate = vecSameConstructor asMaybe asMaybeDate  constructOptional ::     (VU.Unbox a, Columnable a) => VU.Vector a -> VU.Vector Word8 -> IO Column

@@ -40,6 +40,8 @@ import qualified Data.Vector.Unboxed as VU import System.FilePath ((</>)) +-- Options -----------------------------------------------------------------+ {- | Options for reading Parquet data.  These options are applied in this order:@@ -82,6 +84,8 @@         , rowRange = Nothing         } +-- Public API --------------------------------------------------------------+ {- | Read a parquet file from path and load it into a dataframe.  ==== __Example__@@ -104,9 +108,31 @@ When @selectedColumns@ is set and @predicate@ references other columns, those predicate columns are auto-included for decoding, then projected back to the requested output columns. -}++{- | Strip Parquet encoding artifact names (REPEATED wrappers and their single+  list-element children) from a raw column path, leaving user-visible names.+-}+cleanColPath :: [SNode] -> [String] -> [String]+cleanColPath nodes path = go nodes path False+  where+    go _ [] _ = []+    go ns (p : ps) skipThis =+        case L.find (\n -> sName n == p) ns of+            Nothing -> []+            Just n+                | sRep n == REPEATED && not (null (sChildren n)) ->+                    let skipChildren = length (sChildren n) == 1+                     in go (sChildren n) ps skipChildren+                | skipThis ->+                    go (sChildren n) ps False+                | null (sChildren n) ->+                    [p]+                | otherwise ->+                    p : go (sChildren n) ps False+ readParquetWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame readParquetWithOpts opts path = do-    fileMetadata <- readMetadataFromPath path+    (fileMetadata, contents) <- readMetadataFromPath path     let columnPaths = getColumnPaths (drop 1 $ schema fileMetadata)     let columnNames = map fst columnPaths     let leafNames = map (last . T.splitOn ".") columnNames@@ -135,12 +161,13 @@                         )                     ) -    colMap <- newIORef (M.empty :: M.Map T.Text DI.Column)+    let totalRows = sum (map (fromIntegral . rowGroupNumRows) (rowGroups fileMetadata)) :: Int+    colMutMap <- newIORef (M.empty :: M.Map T.Text DI.MutableColumn)+    colOffMap <- newIORef (M.empty :: M.Map T.Text Int)     lTypeMap <- newIORef (M.empty :: M.Map T.Text LogicalType) -    contents <- BSO.readFile path-     let schemaElements = schema fileMetadata+    let sNodes = parseAll (drop 1 schemaElements)     let getTypeLength :: [String] -> Maybe Int32         getTypeLength path = findTypeLength schemaElements path 0           where@@ -159,12 +186,17 @@         forM_ (zip (rowGroupColumns rowGroup) [0 ..]) $ \(colChunk, colIdx) -> do             let metadata = columnMetaData colChunk             let colPath = columnPathInSchema metadata-            let colName =-                    if null colPath+            let cleanPath = cleanColPath sNodes colPath+            let colLeafName =+                    if null cleanPath                         then T.pack $ "col_" ++ show colIdx-                        else T.pack $ last colPath+                        else T.pack $ last cleanPath+            let colFullName =+                    if null cleanPath+                        then colLeafName+                        else T.intercalate "." $ map T.pack cleanPath -            when (shouldReadColumn colName colPath) $ do+            when (shouldReadColumn colLeafName colPath) $ do                 let colDataPageOffset = columnDataPageOffset metadata                 let colDictionaryPageOffset = columnDictionaryPageOffset metadata                 let colStart =@@ -186,7 +218,11 @@                  let schemaTail = drop 1 (schema fileMetadata)                 let (maxDef, maxRep) = levelsForPath schemaTail colPath-                let lType = logicalType (schemaTail !! colIdx)+                let lType =+                        maybe+                            LOGICAL_TYPE_UNKNOWN+                            logicalType+                            (findLeafSchema schemaTail colPath)                 column <-                     processColumnPages                         (maxDef, maxRep)@@ -196,10 +232,22 @@                         maybeTypeLength                         lType -                modifyIORef colMap (M.insertWith DI.concatColumnsEither colName column)-                modifyIORef lTypeMap (M.insert colName lType)+                mutMapSnap <- readIORef colMutMap+                case M.lookup colFullName mutMapSnap of+                    Nothing -> do+                        mc <- DI.newMutableColumn totalRows column+                        DI.copyIntoMutableColumn mc 0 column+                        modifyIORef colMutMap (M.insert colFullName mc)+                        modifyIORef colOffMap (M.insert colFullName (DI.columnLength column))+                    Just mc -> do+                        off <- (M.! colFullName) <$> readIORef colOffMap+                        DI.copyIntoMutableColumn mc off column+                        modifyIORef colOffMap (M.adjust (+ DI.columnLength column) colFullName)+                modifyIORef lTypeMap (M.insert colFullName lType) -    finalColMap <- readIORef colMap+    finalMutMap <- readIORef colMutMap+    finalColMap <-+        M.traverseWithKey (\_ mc -> DI.freezeMutableColumn mc) finalMutMap     finalLTypeMap <- readIORef lTypeMap     let orderedColumns =             map@@ -237,7 +285,7 @@ readParquetFilesWithOpts opts path = do     isDir <- doesDirectoryExist path -    let pat = if isDir then path </> "*" else path+    let pat = if isDir then path </> "*.parquet" else path      matches <- glob pat @@ -252,6 +300,8 @@             dfs <- mapM (readParquetWithOpts optsWithoutRowRange) files             pure (applyRowRange opts (mconcat dfs)) +-- Options application -----------------------------------------------------+ applyRowRange :: ParquetReadOptions -> DataFrame -> DataFrame applyRowRange opts df =     maybe df (`DS.range` df) (rowRange opts)@@ -270,12 +320,15 @@         . applySelectedColumns opts         . applyPredicate opts -readMetadataFromPath :: FilePath -> IO FileMetadata+-- File and metadata parsing -----------------------------------------------++readMetadataFromPath :: FilePath -> IO (FileMetadata, BSO.ByteString) readMetadataFromPath path = do     contents <- BSO.readFile path     let (size, magicString) = contents `seq` readMetadataSizeFromFooter contents     when (magicString /= "PAR1") $ error "Invalid Parquet file"-    readMetadata contents size+    meta <- readMetadata contents size+    pure (meta, contents)  readMetadataSizeFromFooter :: BSO.ByteString -> (Int, BSO.ByteString) readMetadataSizeFromFooter contents =@@ -291,22 +344,44 @@      in         (size, magicString) +-- Schema navigation -------------------------------------------------------+ getColumnPaths :: [SchemaElement] -> [(T.Text, Int)]-getColumnPaths schema = extractLeafPaths schema 0 []+getColumnPaths schemaElements =+    let nodes = parseAll schemaElements+     in go nodes 0 [] False   where-    extractLeafPaths :: [SchemaElement] -> Int -> [T.Text] -> [(T.Text, Int)]-    extractLeafPaths [] _ _ = []-    extractLeafPaths (s : ss) idx path-        | numChildren s == 0 =-            let fullPath = T.intercalate "." (path ++ [elementName s])-             in (fullPath, idx) : extractLeafPaths ss (idx + 1) path+    go [] _ _ _ = []+    go (n : ns) idx path skipThis+        | null (sChildren n) =+            let newPath = if skipThis then path else path ++ [T.pack (sName n)]+                fullPath = T.intercalate "." newPath+             in (fullPath, idx) : go ns (idx + 1) path skipThis+        | sRep n == REPEATED =+            let skipChildren = length (sChildren n) == 1+                childLeaves = go (sChildren n) idx path skipChildren+             in childLeaves ++ go ns (idx + length childLeaves) path skipThis+        | skipThis =+            let childLeaves = go (sChildren n) idx path False+             in childLeaves ++ go ns (idx + length childLeaves) path skipThis         | otherwise =-            let newPath = if T.null (elementName s) then path else path ++ [elementName s]-                childrenCount = fromIntegral (numChildren s)-                (children, remaining) = splitAt childrenCount ss-                childResults = extractLeafPaths children idx newPath-             in childResults ++ extractLeafPaths remaining (idx + length childResults) path+            let subPath = path ++ [T.pack (sName n)]+                childLeaves = go (sChildren n) idx subPath False+             in childLeaves ++ go ns (idx + length childLeaves) path skipThis +findLeafSchema :: [SchemaElement] -> [String] -> Maybe SchemaElement+findLeafSchema elems path =+    case go (parseAll elems) path of+        Just node -> L.find (\e -> T.unpack (elementName e) == sName node) elems+        Nothing -> Nothing+  where+    go [] _ = Nothing+    go _ [] = Nothing+    go nodes [p] = L.find (\n -> sName n == p) nodes+    go nodes (p : ps) = L.find (\n -> sName n == p) nodes >>= \n -> go (sChildren n) ps++-- Page decoding -----------------------------------------------------------+ processColumnPages ::     (Int, Int) ->     [Page] ->@@ -327,59 +402,32 @@                         DictionaryPageHeader{..} ->                             let countForBools =                                     if pType == PBOOLEAN-                                        then error "is bool" Just dictionaryPageHeaderNumValues+                                        then Just dictionaryPageHeaderNumValues                                         else maybeTypeLength                              in Just (readDictVals pType (pageBytes dictPage) countForBools)                         _ -> Nothing      cols <- forM dataPages $ \page -> do+        let bs0 = pageBytes page         case pageTypeHeader (pageHeader page) of             DataPageHeader{..} -> do                 let n = fromIntegral dataPageHeaderNumValues-                let bs0 = pageBytes page-                let (defLvls, _repLvls, afterLvls) = readLevelsV1 n maxDef maxRep bs0-                let nPresent = length (filter (== maxDef) defLvls)--                case dataPageHeaderEncoding of-                    EPLAIN ->-                        case pType of-                            PBOOLEAN ->-                                let (vals, _) = readNBool nPresent afterLvls-                                 in pure (toMaybeBool maxDef defLvls vals)-                            PINT32 ->-                                let (vals, _) = readNInt32 nPresent afterLvls-                                 in pure (toMaybeInt32 maxDef defLvls vals)-                            PINT64 ->-                                let (vals, _) = readNInt64 nPresent afterLvls-                                 in pure (toMaybeInt64 maxDef defLvls vals)-                            PINT96 ->-                                let (vals, _) = readNInt96Times nPresent afterLvls-                                 in pure (toMaybeUTCTime maxDef defLvls vals)-                            PFLOAT ->-                                let (vals, _) = readNFloat nPresent afterLvls-                                 in pure (toMaybeFloat maxDef defLvls vals)-                            PDOUBLE ->-                                let (vals, _) = readNDouble nPresent afterLvls-                                 in pure (toMaybeDouble maxDef defLvls vals)-                            PBYTE_ARRAY ->-                                let (raws, _) = readNByteArrays nPresent afterLvls-                                    texts = map decodeUtf8 raws-                                 in pure (toMaybeText maxDef defLvls texts)-                            PFIXED_LEN_BYTE_ARRAY ->-                                case maybeTypeLength of-                                    Just len ->-                                        let (raws, _) = splitFixed nPresent (fromIntegral len) afterLvls-                                            texts = map decodeUtf8 raws-                                         in pure (toMaybeText maxDef defLvls texts)-                                    Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type length"-                            PARQUET_TYPE_UNKNOWN -> error "Cannot read unknown Parquet type"-                    ERLE_DICTIONARY -> decodeDictV1 dictValsM maxDef defLvls nPresent afterLvls-                    EPLAIN_DICTIONARY -> decodeDictV1 dictValsM maxDef defLvls nPresent afterLvls-                    other -> error ("Unsupported v1 encoding: " ++ show other)+                    (defLvls, repLvls, afterLvls) = readLevelsV1 n maxDef maxRep bs0+                    nPresent = length (filter (== maxDef) defLvls)+                decodePageData+                    dictValsM+                    (maxDef, maxRep)+                    pType+                    maybeTypeLength+                    dataPageHeaderEncoding+                    defLvls+                    repLvls+                    nPresent+                    afterLvls+                    "v1"             DataPageHeaderV2{..} -> do                 let n = fromIntegral dataPageHeaderV2NumValues-                let bs0 = pageBytes page-                let (defLvls, _repLvls, afterLvls) =+                    (defLvls, repLvls, afterLvls) =                         readLevelsV2                             n                             maxDef@@ -387,57 +435,119 @@                             definitionLevelByteLength                             repetitionLevelByteLength                             bs0-                let nPresent =-                        if dataPageHeaderV2NumNulls > 0-                            then fromIntegral (dataPageHeaderV2NumValues - dataPageHeaderV2NumNulls)-                            else length (filter (== maxDef) defLvls)+                    nPresent+                        | dataPageHeaderV2NumNulls > 0 =+                            fromIntegral (dataPageHeaderV2NumValues - dataPageHeaderV2NumNulls)+                        | otherwise = length (filter (== maxDef) defLvls)+                decodePageData+                    dictValsM+                    (maxDef, maxRep)+                    pType+                    maybeTypeLength+                    dataPageHeaderV2Encoding+                    defLvls+                    repLvls+                    nPresent+                    afterLvls+                    "v2" -                case dataPageHeaderV2Encoding of-                    EPLAIN ->-                        case pType of-                            PBOOLEAN ->-                                let (vals, _) = readNBool nPresent afterLvls-                                 in pure (toMaybeBool maxDef defLvls vals)-                            PINT32 ->-                                let (vals, _) = readNInt32 nPresent afterLvls-                                 in pure (toMaybeInt32 maxDef defLvls vals)-                            PINT64 ->-                                let (vals, _) = readNInt64 nPresent afterLvls-                                 in pure (toMaybeInt64 maxDef defLvls vals)-                            PINT96 ->-                                let (vals, _) = readNInt96Times nPresent afterLvls-                                 in pure (toMaybeUTCTime maxDef defLvls vals)-                            PFLOAT ->-                                let (vals, _) = readNFloat nPresent afterLvls-                                 in pure (toMaybeFloat maxDef defLvls vals)-                            PDOUBLE ->-                                let (vals, _) = readNDouble nPresent afterLvls-                                 in pure (toMaybeDouble maxDef defLvls vals)-                            PBYTE_ARRAY ->-                                let (raws, _) = readNByteArrays nPresent afterLvls-                                    texts = map decodeUtf8 raws-                                 in pure (toMaybeText maxDef defLvls texts)-                            PFIXED_LEN_BYTE_ARRAY ->-                                case maybeTypeLength of-                                    Just len ->-                                        let (raws, _) = splitFixed nPresent (fromIntegral len) afterLvls-                                            texts = map decodeUtf8 raws-                                         in pure (toMaybeText maxDef defLvls texts)-                                    Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type length"-                            PARQUET_TYPE_UNKNOWN -> error "Cannot read unknown Parquet type"-                    ERLE_DICTIONARY -> decodeDictV1 dictValsM maxDef defLvls nPresent afterLvls-                    EPLAIN_DICTIONARY -> decodeDictV1 dictValsM maxDef defLvls nPresent afterLvls-                    other -> error ("Unsupported v2 encoding: " ++ show other)             -- Cannot happen as these are filtered out by isDataPage above             DictionaryPageHeader{} -> error "processColumnPages: impossible DictionaryPageHeader"             INDEX_PAGE_HEADER -> error "processColumnPages: impossible INDEX_PAGE_HEADER"             PAGE_TYPE_HEADER_UNKNOWN -> error "processColumnPages: impossible PAGE_TYPE_HEADER_UNKNOWN"-    -- This is N^2. We should probably use mutable columns here.-    case cols of-        [] -> pure $ DI.fromList ([] :: [Maybe Int])-        (c : cs) ->-            pure $-                L.foldl' (\l r -> fromRight (error "concat failed") (DI.concatColumns l r)) c cs+    pure $ DI.concatManyColumns cols++decodePageData ::+    Maybe DictVals ->+    (Int, Int) ->+    ParquetType ->+    Maybe Int32 ->+    ParquetEncoding ->+    [Int] ->+    [Int] ->+    Int ->+    BSO.ByteString ->+    String ->+    IO DI.Column+decodePageData dictValsM (maxDef, maxRep) pType maybeTypeLength encoding defLvls repLvls nPresent afterLvls versionLabel =+    case encoding of+        EPLAIN ->+            case pType of+                PBOOLEAN ->+                    let (vals, _) = readNBool nPresent afterLvls+                     in pure $+                            if maxRep > 0+                                then stitchForRepBool maxRep maxDef repLvls defLvls vals+                                else toMaybeBool maxDef defLvls vals+                PINT32+                    | maxDef == 0+                    , maxRep == 0 ->+                        pure $ DI.fromUnboxedVector (readNInt32Vec nPresent afterLvls)+                PINT32 ->+                    let (vals, _) = readNInt32 nPresent afterLvls+                     in pure $+                            if maxRep > 0+                                then stitchForRepInt32 maxRep maxDef repLvls defLvls vals+                                else toMaybeInt32 maxDef defLvls vals+                PINT64+                    | maxDef == 0+                    , maxRep == 0 ->+                        pure $ DI.fromUnboxedVector (readNInt64Vec nPresent afterLvls)+                PINT64 ->+                    let (vals, _) = readNInt64 nPresent afterLvls+                     in pure $+                            if maxRep > 0+                                then stitchForRepInt64 maxRep maxDef repLvls defLvls vals+                                else toMaybeInt64 maxDef defLvls vals+                PINT96 ->+                    let (vals, _) = readNInt96Times nPresent afterLvls+                     in pure $+                            if maxRep > 0+                                then stitchForRepUTCTime maxRep maxDef repLvls defLvls vals+                                else toMaybeUTCTime maxDef defLvls vals+                PFLOAT+                    | maxDef == 0+                    , maxRep == 0 ->+                        pure $ DI.fromUnboxedVector (readNFloatVec nPresent afterLvls)+                PFLOAT ->+                    let (vals, _) = readNFloat nPresent afterLvls+                     in pure $+                            if maxRep > 0+                                then stitchForRepFloat maxRep maxDef repLvls defLvls vals+                                else toMaybeFloat maxDef defLvls vals+                PDOUBLE+                    | maxDef == 0+                    , maxRep == 0 ->+                        pure $ DI.fromUnboxedVector (readNDoubleVec nPresent afterLvls)+                PDOUBLE ->+                    let (vals, _) = readNDouble nPresent afterLvls+                     in pure $+                            if maxRep > 0+                                then stitchForRepDouble maxRep maxDef repLvls defLvls vals+                                else toMaybeDouble maxDef defLvls vals+                PBYTE_ARRAY ->+                    let (raws, _) = readNByteArrays nPresent afterLvls+                        texts = map decodeUtf8Lenient raws+                     in pure $+                            if maxRep > 0+                                then stitchForRepText maxRep maxDef repLvls defLvls texts+                                else toMaybeText maxDef defLvls texts+                PFIXED_LEN_BYTE_ARRAY ->+                    case maybeTypeLength of+                        Just len ->+                            let (raws, _) = splitFixed nPresent (fromIntegral len) afterLvls+                                texts = map decodeUtf8Lenient raws+                             in pure $+                                    if maxRep > 0+                                        then stitchForRepText maxRep maxDef repLvls defLvls texts+                                        else toMaybeText maxDef defLvls texts+                        Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type length"+                PARQUET_TYPE_UNKNOWN -> error "Cannot read unknown Parquet type"+        ERLE_DICTIONARY -> decodeDictV1 dictValsM maxDef maxRep repLvls defLvls nPresent afterLvls+        EPLAIN_DICTIONARY -> decodeDictV1 dictValsM maxDef maxRep repLvls defLvls nPresent afterLvls+        other -> error ("Unsupported " ++ versionLabel ++ " encoding: " ++ show other)++-- Logical type conversion -------------------------------------------------  applyLogicalType :: LogicalType -> DI.Column -> DI.Column applyLogicalType (TimestampType _ unit) col =

@@ -2,13 +2,18 @@  module DataFrame.IO.Parquet.Binary where +import Control.Exception (bracketOnError) import Control.Monad import Data.Bits import qualified Data.ByteString as BS+import qualified Data.ByteString.Unsafe as BSU import Data.Char import Data.IORef import Data.Int import Data.Word+import qualified Foreign.Marshal.Alloc as Foreign+import qualified Foreign.Ptr as Foreign+import qualified Foreign.Storable as Foreign  littleEndianWord32 :: BS.ByteString -> Word32 littleEndianWord32 bytes@@ -124,7 +129,7 @@ readString :: BS.ByteString -> IORef Int -> IO String readString buf pos = do     nameSize <- readVarIntFromBuffer @Int buf pos-    map (chr . fromIntegral) <$> replicateM nameSize (readAndAdvance pos buf)+    replicateM nameSize (chr . fromIntegral <$> readAndAdvance pos buf)  readByteStringFromBytes :: BS.ByteString -> (BS.ByteString, BS.ByteString) readByteStringFromBytes xs =@@ -136,10 +141,32 @@ readByteString :: BS.ByteString -> IORef Int -> IO BS.ByteString readByteString buf pos = do     size <- readVarIntFromBuffer @Int buf pos-    BS.pack <$> replicateM size (readAndAdvance pos buf)+    fillByteStringByWord8 size (\_ -> readAndAdvance pos buf)  readByteString' :: BS.ByteString -> Int64 -> IO BS.ByteString-readByteString' buf size = BS.pack <$> mapM (`readSingleByte` buf) [0 .. (size - 1)]+readByteString' buf size =+    fillByteStringByWord8+        (fromIntegral size)+        ((`readSingleByte` buf) . fromIntegral)++{- | Allocate a fixed-size buffer, repeat the action on each index.+Fill it into the buffer to get a ByteString.+-}+fillByteStringByWord8 :: Int -> (Int -> IO Word8) -> IO BS.ByteString+fillByteStringByWord8 size getByte = do+    bracketOnError+        (Foreign.mallocBytes size :: IO (Foreign.Ptr Word8))+        Foreign.free+        -- \^ ensures p is freed if (IO Word8) throws.+        ( \p -> do+            fill 0 p+            BSU.unsafePackCStringFinalizer p size (Foreign.free p)+        )+  where+    fill i p+        | i >= size = pure ()+        | otherwise = getByte i >>= Foreign.pokeByteOff p i >> fill (i + 1) p+{-# INLINE fillByteStringByWord8 #-}  readSingleByte :: Int64 -> BS.ByteString -> IO Word8 readSingleByte pos buffer = return $ BS.index buffer (fromIntegral pos)

@@ -1,3 +1,6 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE OverloadedStrings #-}  module DataFrame.IO.Parquet.Dictionary where@@ -5,12 +8,15 @@ import Control.Monad import Data.Bits import qualified Data.ByteString as BS-import Data.Char+import Data.IORef import Data.Int import Data.Maybe import qualified Data.Text as T import Data.Text.Encoding import Data.Time+import qualified Data.Vector as V+import qualified Data.Vector.Mutable as VM+import qualified Data.Vector.Unboxed as VU import DataFrame.IO.Parquet.Binary import DataFrame.IO.Parquet.Encoding import DataFrame.IO.Parquet.Levels@@ -20,23 +26,23 @@ import GHC.Float  dictCardinality :: DictVals -> Int-dictCardinality (DBool ds) = length ds-dictCardinality (DInt32 ds) = length ds-dictCardinality (DInt64 ds) = length ds-dictCardinality (DInt96 ds) = length ds-dictCardinality (DFloat ds) = length ds-dictCardinality (DDouble ds) = length ds-dictCardinality (DText ds) = length ds+dictCardinality (DBool ds) = V.length ds+dictCardinality (DInt32 ds) = V.length ds+dictCardinality (DInt64 ds) = V.length ds+dictCardinality (DInt96 ds) = V.length ds+dictCardinality (DFloat ds) = V.length ds+dictCardinality (DDouble ds) = V.length ds+dictCardinality (DText ds) = V.length ds  readDictVals :: ParquetType -> BS.ByteString -> Maybe Int32 -> DictVals-readDictVals PBOOLEAN bs (Just count) = DBool (take (fromIntegral count) $ readPageBool bs)-readDictVals PINT32 bs _ = DInt32 (readPageInt32 bs)-readDictVals PINT64 bs _ = DInt64 (readPageInt64 bs)-readDictVals PINT96 bs _ = DInt96 (readPageInt96Times bs)-readDictVals PFLOAT bs _ = DFloat (readPageFloat bs)-readDictVals PDOUBLE bs _ = DDouble (readPageWord64 bs)-readDictVals PBYTE_ARRAY bs _ = DText (readPageBytes bs)-readDictVals PFIXED_LEN_BYTE_ARRAY bs (Just len) = DText (readPageFixedBytes bs (fromIntegral len))+readDictVals PBOOLEAN bs (Just count) = DBool (V.fromList (take (fromIntegral count) $ readPageBool bs))+readDictVals PINT32 bs _ = DInt32 (V.fromList (readPageInt32 bs))+readDictVals PINT64 bs _ = DInt64 (V.fromList (readPageInt64 bs))+readDictVals PINT96 bs _ = DInt96 (V.fromList (readPageInt96Times bs))+readDictVals PFLOAT bs _ = DFloat (V.fromList (readPageFloat bs))+readDictVals PDOUBLE bs _ = DDouble (V.fromList (readPageWord64 bs))+readDictVals PBYTE_ARRAY bs _ = DText (V.fromList (readPageBytes bs))+readDictVals PFIXED_LEN_BYTE_ARRAY bs (Just len) = DText (V.fromList (readPageFixedBytes bs (fromIntegral len))) readDictVals t _ _ = error $ "Unsupported dictionary type: " ++ show t  readPageInt32 :: BS.ByteString -> [Int32]@@ -57,7 +63,7 @@     | otherwise =         let lenBytes = fromIntegral (littleEndianInt32 $ BS.take 4 xs)             totalBytesRead = lenBytes + 4-         in T.pack (map (chr . fromIntegral) $ take lenBytes (BS.unpack (BS.drop 4 xs)))+         in decodeUtf8Lenient (BS.take lenBytes (BS.drop 4 xs))                 : readPageBytes (BS.drop totalBytesRead xs)  readPageBool :: BS.ByteString -> [Bool]@@ -97,44 +103,154 @@ readPageFixedBytes xs len     | BS.null xs = []     | otherwise =-        decodeUtf8 (BS.take len xs) : readPageFixedBytes (BS.drop len xs) len+        decodeUtf8Lenient (BS.take len xs) : readPageFixedBytes (BS.drop len xs) len +{- | Dispatch to the right multi-level list stitching function.+For maxRep=1 uses stitchList; for 2/3 uses stitchList2/3 with computed thresholds.+Threshold formula: defT_r = maxDef - 2*(maxRep - r).+-}+stitchForRepBool :: Int -> Int -> [Int] -> [Int] -> [Bool] -> DI.Column+stitchForRepBool maxRep maxDef rep def vals = case maxRep of+    2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)+    3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)+    _ -> DI.fromList (stitchList maxDef rep def vals)++stitchForRepInt32 :: Int -> Int -> [Int] -> [Int] -> [Int32] -> DI.Column+stitchForRepInt32 maxRep maxDef rep def vals = case maxRep of+    2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)+    3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)+    _ -> DI.fromList (stitchList maxDef rep def vals)++stitchForRepInt64 :: Int -> Int -> [Int] -> [Int] -> [Int64] -> DI.Column+stitchForRepInt64 maxRep maxDef rep def vals = case maxRep of+    2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)+    3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)+    _ -> DI.fromList (stitchList maxDef rep def vals)++stitchForRepUTCTime :: Int -> Int -> [Int] -> [Int] -> [UTCTime] -> DI.Column+stitchForRepUTCTime maxRep maxDef rep def vals = case maxRep of+    2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)+    3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)+    _ -> DI.fromList (stitchList maxDef rep def vals)++stitchForRepFloat :: Int -> Int -> [Int] -> [Int] -> [Float] -> DI.Column+stitchForRepFloat maxRep maxDef rep def vals = case maxRep of+    2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)+    3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)+    _ -> DI.fromList (stitchList maxDef rep def vals)++stitchForRepDouble :: Int -> Int -> [Int] -> [Int] -> [Double] -> DI.Column+stitchForRepDouble maxRep maxDef rep def vals = case maxRep of+    2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)+    3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)+    _ -> DI.fromList (stitchList maxDef rep def vals)++stitchForRepText :: Int -> Int -> [Int] -> [Int] -> [T.Text] -> DI.Column+stitchForRepText maxRep maxDef rep def vals = case maxRep of+    2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)+    3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)+    _ -> DI.fromList (stitchList maxDef rep def vals)++{- | Build a Column from a dictionary + index vector + def levels in a single+mutable-vector pass, avoiding the intermediate [a] and [Maybe a] lists.+For maxRep > 0 (list columns) the caller must use the rep-stitching path instead.+-}+applyDictToColumn ::+    (DI.Columnable a, DI.Columnable (Maybe a)) =>+    V.Vector a ->+    VU.Vector Int ->+    Int -> -- maxDef+    [Int] -> -- defLvls+    IO DI.Column+applyDictToColumn dict idxs maxDef defLvls+    | maxDef == 0 = do+        -- All rows are required; no nullability to check.+        let n = VU.length idxs+        pure $ DI.fromVector (V.generate n (\i -> dict V.! (idxs VU.! i)))+    | otherwise = do+        let n = length defLvls+        mv <- VM.new n+        hasNullRef <- newIORef False+        let go _ _ [] = pure ()+            go !i !j (d : ds)+                | d == maxDef = do+                    VM.write mv i (Just (dict V.! (idxs VU.! j)))+                    go (i + 1) (j + 1) ds+                | otherwise = do+                    writeIORef hasNullRef True+                    VM.write mv i Nothing+                    go (i + 1) j ds+        go 0 0 defLvls+        vec <- V.freeze mv+        hasNull <- readIORef hasNullRef+        pure $+            if hasNull+                then DI.fromVector vec -- VB.Vector (Maybe a) → OptionalColumn+                else DI.fromVector (V.map fromJust vec) -- VB.Vector a → BoxedColumn/UnboxedColumn+ decodeDictV1 ::-    Maybe DictVals -> Int -> [Int] -> Int -> BS.ByteString -> IO DI.Column-decodeDictV1 dictValsM maxDef defLvls nPresent bytes =+    Maybe DictVals ->+    Int ->+    Int ->+    [Int] ->+    [Int] ->+    Int ->+    BS.ByteString ->+    IO DI.Column+decodeDictV1 dictValsM maxDef maxRep repLvls defLvls nPresent bytes =     case dictValsM of         Nothing -> error "Dictionary-encoded page but dictionary is missing"         Just dictVals ->             let (idxs, _rest) = decodeDictIndicesV1 nPresent (dictCardinality dictVals) bytes              in do-                    when (length idxs /= nPresent) $+                    when (VU.length idxs /= nPresent) $                         error $                             "dict index count mismatch: got "-                                ++ show (length idxs)+                                ++ show (VU.length idxs)                                 ++ ", expected "                                 ++ show nPresent-                    case dictVals of-                        DBool ds -> do-                            let values = [ds !! i | i <- idxs]-                            pure (toMaybeBool maxDef defLvls values)-                        DInt32 ds -> do-                            let values = [ds !! i | i <- idxs]-                            pure (toMaybeInt32 maxDef defLvls values)-                        DInt64 ds -> do-                            let values = [ds !! i | i <- idxs]-                            pure (toMaybeInt64 maxDef defLvls values)-                        DInt96 ds -> do-                            let values = [ds !! i | i <- idxs]-                            pure (toMaybeUTCTime maxDef defLvls values)-                        DFloat ds -> do-                            let values = [ds !! i | i <- idxs]-                            pure (toMaybeFloat maxDef defLvls values)-                        DDouble ds -> do-                            let values = [ds !! i | i <- idxs]-                            pure (toMaybeDouble maxDef defLvls values)-                        DText ds -> do-                            let values = [ds !! i | i <- idxs]-                            pure (toMaybeText maxDef defLvls values)+                    if maxRep > 0+                        then do+                            case dictVals of+                                DBool ds ->+                                    pure $+                                        stitchForRepBool maxRep maxDef repLvls defLvls (map (ds V.!) (VU.toList idxs))+                                DInt32 ds ->+                                    pure $+                                        stitchForRepInt32 maxRep maxDef repLvls defLvls (map (ds V.!) (VU.toList idxs))+                                DInt64 ds ->+                                    pure $+                                        stitchForRepInt64 maxRep maxDef repLvls defLvls (map (ds V.!) (VU.toList idxs))+                                DInt96 ds ->+                                    pure $+                                        stitchForRepUTCTime+                                            maxRep+                                            maxDef+                                            repLvls+                                            defLvls+                                            (map (ds V.!) (VU.toList idxs))+                                DFloat ds ->+                                    pure $+                                        stitchForRepFloat maxRep maxDef repLvls defLvls (map (ds V.!) (VU.toList idxs))+                                DDouble ds ->+                                    pure $+                                        stitchForRepDouble maxRep maxDef repLvls defLvls (map (ds V.!) (VU.toList idxs))+                                DText ds ->+                                    pure $+                                        stitchForRepText maxRep maxDef repLvls defLvls (map (ds V.!) (VU.toList idxs))+                        else case dictVals of+                            -- Fast path: unboxable types, no nulls — one allocation via VU.map+                            DInt32 ds | maxDef == 0 -> pure $ DI.fromUnboxedVector (VU.map (ds V.!) idxs)+                            DInt64 ds | maxDef == 0 -> pure $ DI.fromUnboxedVector (VU.map (ds V.!) idxs)+                            DFloat ds | maxDef == 0 -> pure $ DI.fromUnboxedVector (VU.map (ds V.!) idxs)+                            DDouble ds | maxDef == 0 -> pure $ DI.fromUnboxedVector (VU.map (ds V.!) idxs)+                            DBool ds -> applyDictToColumn ds idxs maxDef defLvls+                            DInt32 ds -> applyDictToColumn ds idxs maxDef defLvls+                            DInt64 ds -> applyDictToColumn ds idxs maxDef defLvls+                            DInt96 ds -> applyDictToColumn ds idxs maxDef defLvls+                            DFloat ds -> applyDictToColumn ds idxs maxDef defLvls+                            DDouble ds -> applyDictToColumn ds idxs maxDef defLvls+                            DText ds -> applyDictToColumn ds idxs maxDef defLvls  toMaybeInt32 :: Int -> [Int] -> [Int32] -> DI.Column toMaybeInt32 maxDef def xs =

@@ -1,8 +1,12 @@+{-# LANGUAGE BangPatterns #-}+ module DataFrame.IO.Parquet.Encoding where  import Data.Bits import qualified Data.ByteString as BS+import qualified Data.ByteString.Unsafe as BSU import Data.List (foldl')+import qualified Data.Vector.Unboxed as VU import Data.Word import DataFrame.IO.Parquet.Binary @@ -22,32 +26,26 @@     | count <= 0 = ([], bs)     | BS.null bs = ([], bs)     | otherwise =-        let totalBits = bw * count-            totalBytes = (totalBits + 7) `div` 8+        let totalBytes = (bw * count + 7) `div` 8             chunk = BS.take totalBytes bs             rest = BS.drop totalBytes bs-            bits =-                BS.concatMap-                    (\b -> BS.map (\i -> (b `shiftR` fromIntegral i) .&. 1) (BS.pack [0 .. 7]))-                    chunk-            toN :: BS.ByteString -> Word32-            toN =-                fst-                    . BS.foldl'-                        (\(a, i) b -> (a .|. (fromIntegral b `shiftL` i), i + 1))-                        (0 :: Word32, 0 :: Int)--            extractValues :: Int -> BS.ByteString -> [Word32]-            extractValues n bitsLeft-                | BS.null bitsLeft = []-                | n <= 0 = []-                | BS.length bitsLeft < bw = []-                | otherwise =-                    let (this, bitsLeft') = BS.splitAt bw bitsLeft-                     in toN this : extractValues (n - 1) bitsLeft'+         in (extractBits bw count chunk, rest) -            vals = extractValues count bits-         in (vals, rest)+-- | LSB-first bit accumulator: reads each byte once with no intermediate ByteString allocation.+extractBits :: Int -> Int -> BS.ByteString -> [Word32]+extractBits bw count bs = go 0 (0 :: Word64) 0 count+  where+    !mask = if bw == 32 then maxBound else (1 `shiftL` bw) - 1 :: Word64+    !len = BS.length bs+    go !byteIdx !acc !accBits !remaining+        | remaining <= 0 = []+        | accBits >= bw =+            fromIntegral (acc .&. mask)+                : go byteIdx (acc `shiftR` bw) (accBits - bw) (remaining - 1)+        | byteIdx >= len = []+        | otherwise =+            let b = fromIntegral (BSU.unsafeIndex bs byteIdx) :: Word64+             in go (byteIdx + 1) (acc .|. (b `shiftL` accBits)) (accBits + 8) remaining  decodeRLEBitPackedHybrid ::     Int -> Int -> BS.ByteString -> ([Word32], BS.ByteString)@@ -81,11 +79,12 @@                             takeN = min n runLen                          in go (n - takeN) afterV (replicate takeN val ++ acc) -decodeDictIndicesV1 :: Int -> Int -> BS.ByteString -> ([Int], BS.ByteString)+decodeDictIndicesV1 ::+    Int -> Int -> BS.ByteString -> (VU.Vector Int, BS.ByteString) decodeDictIndicesV1 need dictCard bs =     case BS.uncons bs of         Nothing -> error "empty dictionary index stream"         Just (w0, rest0) ->             let bw = fromIntegral w0 :: Int                 (u32s, rest1) = decodeRLEBitPackedHybrid bw need rest0-             in (map fromIntegral u32s, rest1)+             in (VU.fromList (map fromIntegral u32s), rest1)

@@ -4,7 +4,6 @@ import Data.Int import Data.List import qualified Data.Text as T-import Data.Word  import DataFrame.IO.Parquet.Binary import DataFrame.IO.Parquet.Encoding@@ -17,7 +16,7 @@     let bwDef = bitWidthForMaxLevel maxDef         bwRep = bitWidthForMaxLevel maxRep -        (repLvlsU32, afterRep) =+        (repLvls, afterRep) =             if bwRep == 0                 then (replicate n 0, bs)                 else@@ -25,9 +24,9 @@                         repData = BS.take (fromIntegral repLength) (BS.drop 4 bs)                         afterRepData = BS.drop (4 + fromIntegral repLength) bs                         (repVals, _) = decodeRLEBitPackedHybrid bwRep n repData-                     in (repVals, afterRepData)+                     in (map fromIntegral repVals, afterRepData) -        (defLvlsU32, afterDef) =+        (defLvls, afterDef) =             if bwDef == 0                 then (replicate n 0, afterRep)                 else@@ -35,8 +34,8 @@                         defData = BS.take (fromIntegral defLength) (BS.drop 4 afterRep)                         afterDefData = BS.drop (4 + fromIntegral defLength) afterRep                         (defVals, _) = decodeRLEBitPackedHybrid bwDef n defData-                     in (defVals, afterDefData)-     in (map fromIntegral defLvlsU32, map fromIntegral repLvlsU32, afterDef)+                     in (map fromIntegral defVals, afterDefData)+     in (defLvls, repLvls, afterDef)  readLevelsV2 ::     Int ->@@ -51,15 +50,15 @@         (defBytes, afterDefBytes) = BS.splitAt (fromIntegral defLen) afterRepBytes         bwDef = bitWidthForMaxLevel maxDef         bwRep = bitWidthForMaxLevel maxRep-        (repLvlsU32, _) =+        (repLvlsRaw, _) =             if bwRep == 0                 then (replicate n 0, repBytes)                 else decodeRLEBitPackedHybrid bwRep n repBytes-        (defLvlsU32, _) =+        (defLvlsRaw, _) =             if bwDef == 0                 then (replicate n 0, defBytes)                 else decodeRLEBitPackedHybrid bwDef n defBytes-     in (map fromIntegral defLvlsU32, map fromIntegral repLvlsU32, afterDefBytes)+     in (map fromIntegral defLvlsRaw, map fromIntegral repLvlsRaw, afterDefBytes)  stitchNullable :: Int -> [Int] -> [a] -> [Maybe a] stitchNullable maxDef = go@@ -102,6 +101,76 @@ parseAll [] = [] parseAll xs = let (n, xs') = parseOne xs in n : parseAll xs' +-- | Tag leaf values as Just/Nothing according to maxDef.+pairWithVals :: Int -> [(Int, Int)] -> [a] -> [(Int, Int, Maybe a)]+pairWithVals _ [] _ = []+pairWithVals maxDef ((r, d) : rds) vs+    | d == maxDef = case vs of+        (v : vs') -> (r, d, Just v) : pairWithVals maxDef rds vs'+        [] -> error "pairWithVals: value stream exhausted"+    | otherwise = (r, d, Nothing) : pairWithVals maxDef rds vs++-- | Split triplets into groups; a new group begins whenever rep <= bound.+splitAtRepBound :: Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]+splitAtRepBound _ [] = []+splitAtRepBound bound (t : ts) =+    let (rest, remaining) = span (\(r, _, _) -> r > bound) ts+     in (t : rest) : splitAtRepBound bound remaining++{- | Reconstruct a list column from Dremel encoding levels.+rep=0 starts a new top-level row; def=0 means the entire list slot is null.+Returns one Maybe [Maybe a] per row.+-}+stitchList :: Int -> [Int] -> [Int] -> [a] -> [Maybe [Maybe a]]+stitchList maxDef repLvls defLvls vals =+    let triplets = pairWithVals maxDef (zip repLvls defLvls) vals+        rows = splitAtRepBound 0 triplets+     in map toRow rows+  where+    toRow [] = Nothing+    toRow ((_, d, _) : _) | d == 0 = Nothing+    toRow grp = Just [v | (_, _, v) <- grp]++{- | Reconstruct a 2-level nested list (maxRep=2) from Dremel triplets.+defT1: def threshold at which the depth-1 element is present (not null).+maxDef: def threshold at which the leaf is present.+-}+stitchList2 :: Int -> Int -> [Int] -> [Int] -> [a] -> [Maybe [Maybe [Maybe a]]]+stitchList2 defT1 maxDef repLvls defLvls vals =+    let triplets = pairWithVals maxDef (zip repLvls defLvls) vals+     in map toRow (splitAtRepBound 0 triplets)+  where+    toRow [] = Nothing+    toRow ((_, d, _) : _) | d == 0 = Nothing+    toRow row = Just (map toOuter (splitAtRepBound 1 row))+    toOuter [] = Nothing+    toOuter ((_, d, _) : _) | d < defT1 = Nothing+    toOuter outer = Just (map toLeaf (splitAtRepBound 2 outer))+    toLeaf [] = Nothing+    toLeaf ((_, _, v) : _) = v++{- | Reconstruct a 3-level nested list (maxRep=3) from Dremel triplets.+defT1, defT2: def thresholds at which depth-1 and depth-2 elements are present.+maxDef: def threshold at which the leaf is present.+-}+stitchList3 ::+    Int -> Int -> Int -> [Int] -> [Int] -> [a] -> [Maybe [Maybe [Maybe [Maybe a]]]]+stitchList3 defT1 defT2 maxDef repLvls defLvls vals =+    let triplets = pairWithVals maxDef (zip repLvls defLvls) vals+     in map toRow (splitAtRepBound 0 triplets)+  where+    toRow [] = Nothing+    toRow ((_, d, _) : _) | d == 0 = Nothing+    toRow row = Just (map toOuter (splitAtRepBound 1 row))+    toOuter [] = Nothing+    toOuter ((_, d, _) : _) | d < defT1 = Nothing+    toOuter outer = Just (map toMiddle (splitAtRepBound 2 outer))+    toMiddle [] = Nothing+    toMiddle ((_, d, _) : _) | d < defT2 = Nothing+    toMiddle middle = Just (map toLeaf (splitAtRepBound 3 middle))+    toLeaf [] = Nothing+    toLeaf ((_, _, v) : _) = v+ levelsForPath :: [SchemaElement] -> [String] -> (Int, Int) levelsForPath schemaTail = go 0 0 (parseAll schemaTail)   where@@ -110,6 +179,6 @@         case find (\n -> sName n == p) nodes of             Nothing -> (defC, repC)             Just n ->-                let defC' = defC + (if sRep n == OPTIONAL then 1 else 0)+                let defC' = defC + (if sRep n == OPTIONAL || sRep n == REPEATED then 1 else 0)                     repC' = repC + (if sRep n == REPEATED then 1 else 0)                  in go defC' repC' (sChildren n) ps

@@ -10,6 +10,7 @@ import qualified Data.ByteString.Lazy as LB import Data.Int import Data.Maybe (fromMaybe)+import qualified Data.Vector.Unboxed as VU import DataFrame.IO.Parquet.Binary import DataFrame.IO.Parquet.Thrift import DataFrame.IO.Parquet.Types@@ -46,9 +47,9 @@                         drainZstd result BS.empty acc                     drainZstd (Zstd.Produce chunk next) _ acc = do                         result <- next-                        drainZstd result BS.empty (acc <> [chunk])+                        drainZstd result BS.empty (chunk : acc)                     drainZstd (Zstd.Done final) _ acc =-                        pure $ BS.concat (acc <> [final])+                        pure $ BS.concat (reverse (final : acc))                     drainZstd (Zstd.Error msg msg2) _ _ =                         error ("ZSTD error: " ++ msg ++ " " ++ msg2)                 SNAPPY -> case Snappy.decompress compressed of@@ -295,6 +296,28 @@                 case maybePage of                     Nothing -> return (reverse acc)                     Just page -> go remaining (page : acc)++-- | Read n Int32 values directly into an unboxed vector (no intermediate list).+readNInt32Vec :: Int -> BS.ByteString -> VU.Vector Int32+readNInt32Vec n bs = VU.generate n (\i -> littleEndianInt32 (BS.drop (4 * i) bs))++-- | Read n Int64 values directly into an unboxed vector.+readNInt64Vec :: Int -> BS.ByteString -> VU.Vector Int64+readNInt64Vec n bs = VU.generate n (\i -> fromIntegral (littleEndianWord64 (BS.drop (8 * i) bs)))++-- | Read n Float values directly into an unboxed vector.+readNFloatVec :: Int -> BS.ByteString -> VU.Vector Float+readNFloatVec n bs =+    VU.generate+        n+        (\i -> castWord32ToFloat (littleEndianWord32 (BS.drop (4 * i) bs)))++-- | Read n Double values directly into an unboxed vector.+readNDoubleVec :: Int -> BS.ByteString -> VU.Vector Double+readNDoubleVec n bs =+    VU.generate+        n+        (\i -> castWord64ToDouble (littleEndianWord64 (BS.drop (8 * i) bs)))  readNInt32 :: Int -> BS.ByteString -> ([Int32], BS.ByteString) readNInt32 0 bs = ([], bs)

@@ -4,6 +4,7 @@ import Data.Int import qualified Data.Text as T import Data.Time+import qualified Data.Vector as V  data ParquetType     = PBOOLEAN@@ -178,13 +179,13 @@     deriving (Show, Eq)  data DictVals-    = DBool [Bool]-    | DInt32 [Int32]-    | DInt64 [Int64]-    | DInt96 [UTCTime]-    | DFloat [Float]-    | DDouble [Double]-    | DText [T.Text]+    = DBool (V.Vector Bool)+    | DInt32 (V.Vector Int32)+    | DInt64 (V.Vector Int64)+    | DInt96 (V.Vector UTCTime)+    | DFloat (V.Vector Float)+    | DDouble (V.Vector Double)+    | DText (V.Vector T.Text)     deriving (Show, Eq)  data Page = Page

@@ -7,6 +7,7 @@     readCsvUnstable,     fastReadTsvUnstable,     readTsvUnstable,+    getDelimiterIndices, ) where  import qualified Data.Vector as Vector@@ -195,30 +196,46 @@ getDelimiterIndices separator originalLen csvFile =     VS.unsafeWith csvFile $ \buffer -> do         let paddedLen = VS.length csvFile-        -- then number of delimiters cannot exceed the size-        -- of the input array (which would be a series of-        -- empty fields)-        indices <- mallocArray paddedLen+        -- GC-managed pinned memory: freed automatically, no leak in streaming use.+        resultMV <- VSM.unsafeNew paddedLen         num_fields <--            get_delimiter_indices-                (castPtr buffer)-                (fromIntegral paddedLen)-                (fromIntegral separator)-                (castPtr indices)+            VSM.unsafeWith resultMV $ \indicesPtr ->+                get_delimiter_indices+                    (castPtr buffer)+                    (fromIntegral paddedLen)+                    (fromIntegral separator)+                    (castPtr indicesPtr)         if num_fields == -1-            then getDelimiterIndices_ separator originalLen csvFile indices+            then do+                -- Haskell state-machine fallback, writing directly into resultMV.+                let trans = stateTransitionTable separator+                    processChar (!state, !idx) i byte =+                        case state of+                            UnEscaped ->+                                if byte == lf || byte == separator+                                    then do+                                        VSM.unsafeWrite resultMV idx (fromIntegral i)+                                        return (toEnum (trans ! (fromEnum state, byte)), idx + 1)+                                    else return (toEnum (trans ! (fromEnum state, byte)), idx)+                            Escaped ->+                                return (toEnum (trans ! (fromEnum state, byte)), idx)+                (_, finalIdx) <- VS.ifoldM' processChar (UnEscaped, 0 :: Int) csvFile+                finalLen <-+                    if originalLen > 0 && csvFile VS.! (originalLen - 1) /= lf+                        then do+                            VSM.unsafeWrite resultMV finalIdx (fromIntegral originalLen)+                            return (finalIdx + 1)+                        else return finalIdx+                VS.unsafeFreeze (VSM.slice 0 finalLen resultMV)             else do-                indices' <- newForeignPtr_ indices-                let resultVector = VSM.unsafeFromForeignPtr0 indices' paddedLen-                -- Handle the case where the file doesn't end with a newline-                -- We need to add a final delimiter for the last field-                finalResultLen <-+                let n = fromIntegral num_fields+                finalLen <-                     if originalLen > 0 && csvFile VS.! (originalLen - 1) /= lf                         then do-                            VSM.write resultVector (fromIntegral num_fields) (fromIntegral originalLen)-                            return (fromIntegral num_fields + 1)-                        else return (fromIntegral num_fields)-                VS.unsafeFreeze $ VSM.slice 0 finalResultLen resultVector+                            VSM.write resultMV n (fromIntegral originalLen)+                            return (n + 1)+                        else return n+                VS.unsafeFreeze (VSM.slice 0 finalLen resultMV)  -- We have a Native version in case the C version -- cannot be used. For example if neither ARM_NEON

@@ -27,6 +27,7 @@ import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM +import Control.DeepSeq (NFData (..), rnf) import Control.Exception (throw) import Control.Monad.ST (runST) import Data.Kind (Type)@@ -36,6 +37,7 @@ import DataFrame.Errors import DataFrame.Internal.Parsing import DataFrame.Internal.Types+import System.IO.Unsafe (unsafePerformIO) import Type.Reflection  {- | Our representation of a column is a GADT that can store data based on the underlying data.@@ -51,6 +53,7 @@ data MutableColumn where     MBoxedColumn :: (Columnable a) => VBM.IOVector a -> MutableColumn     MUnboxedColumn :: (Columnable a, VU.Unbox a) => VUM.IOVector a -> MutableColumn+    MOptionalColumn :: (Columnable a) => VBM.IOVector (Maybe a) -> MutableColumn  {- | A TypedColumn is a wrapper around our type-erased column. It is used to type check expressions on columns.@@ -126,6 +129,11 @@     show :: (Show a) => TypedColumn a -> String     show (TColumn col) = show col +instance NFData Column where+    rnf (BoxedColumn (v :: VB.Vector a)) = rnf v+    rnf (UnboxedColumn v) = v `seq` ()+    rnf (OptionalColumn (v :: VB.Vector (Maybe a))) = rnf v+ instance Show Column where     show :: Column -> String     show (BoxedColumn column) = show column@@ -282,9 +290,6 @@             STrue -> UnboxedColumn (VU.map f col)             SFalse -> fromVector @c (VB.generate (VU.length col) (f . VU.unsafeIndex col))         Nothing -> throwTypeMismatch @a @b-{-# SPECIALIZE mapColumn ::-    (Double -> Double) -> Column -> Either DataFrameException Column-    #-} {-# INLINEABLE mapColumn #-}  -- | Applies a function that returns an unboxed result to an unboxed vector, storing the result in a column.@@ -353,9 +358,17 @@  -- | O(n) Selects the elements at a given set of indices. Does not change the order. atIndicesStable :: VU.Vector Int -> Column -> Column-atIndicesStable indexes (BoxedColumn column) = BoxedColumn $ VG.unsafeBackpermute column (VG.convert indexes)+atIndicesStable indexes (BoxedColumn column) =+    BoxedColumn $+        VB.generate+            (VU.length indexes)+            (\i -> column `VB.unsafeIndex` (indexes `VU.unsafeIndex` i)) atIndicesStable indexes (UnboxedColumn column) = UnboxedColumn $ VU.unsafeBackpermute column indexes-atIndicesStable indexes (OptionalColumn column) = OptionalColumn $ VG.unsafeBackpermute column (VG.convert indexes)+atIndicesStable indexes (OptionalColumn column) =+    OptionalColumn $+        VB.generate+            (VU.length indexes)+            (\i -> column `VB.unsafeIndex` (indexes `VU.unsafeIndex` i)) {-# INLINE atIndicesStable #-}  {- | Like 'atIndicesStable' but treats negative indices as null,@@ -533,6 +546,50 @@                     }                 ) +foldlColumnWith ::+    forall a b.+    (Columnable a) =>+    (b -> a -> b) -> b -> Column -> Either DataFrameException b+foldlColumnWith f acc (BoxedColumn (column :: VB.Vector d)) =+    case testEquality (typeRep @a) (typeRep @d) of+        Just Refl -> pure $ VG.foldl' f acc column+        Nothing ->+            Left $+                TypeMismatchException+                    ( MkTypeErrorContext+                        { userType = Right (typeRep @a)+                        , expectedType = Right (typeRep @d)+                        , callingFunctionName = Just "foldlColumnWith"+                        , errorColumnName = Nothing+                        }+                    )+foldlColumnWith f acc (OptionalColumn (column :: VB.Vector d)) =+    case testEquality (typeRep @a) (typeRep @d) of+        Just Refl -> pure $ VG.foldl' f acc column+        Nothing ->+            Left $+                TypeMismatchException+                    ( MkTypeErrorContext+                        { userType = Right (typeRep @a)+                        , expectedType = Right (typeRep @d)+                        , callingFunctionName = Just "foldlColumnWith"+                        , errorColumnName = Nothing+                        }+                    )+foldlColumnWith f acc (UnboxedColumn (column :: VU.Vector d)) =+    case testEquality (typeRep @a) (typeRep @d) of+        Just Refl -> pure $ VG.foldl' f acc column+        Nothing ->+            Left $+                TypeMismatchException+                    ( MkTypeErrorContext+                        { userType = Right (typeRep @a)+                        , expectedType = Right (typeRep @d)+                        , callingFunctionName = Just "foldlColumnWith"+                        , errorColumnName = Nothing+                        }+                    )+ foldl1Column ::     forall a.     (Columnable a) =>@@ -574,6 +631,292 @@                     }                 ) +{- | O(n) Fold a column over groups without materialising a sorted copy.+Instead of backpermuting the column (expensive: O(n) allocation + random reads),+this iterates @valueIndices@ sequentially and accesses the column at each index.+Avoids one 10M-element allocation per column per aggregation expression.+-}+foldDirectGroups ::+    forall b acc.+    (Columnable b) =>+    (acc -> b -> acc) ->+    acc ->+    Column ->+    VU.Vector Int -> -- valueIndices (sorted row order)+    VU.Vector Int -> -- offsets (group boundaries)+    Either DataFrameException (VB.Vector acc)+foldDirectGroups f seed col valueIndices offsets+    | VU.length offsets <= 1 = Right VB.empty+    | otherwise =+        let !nGroups = VU.length offsets - 1+         in case col of+                UnboxedColumn (vec :: VU.Vector d) ->+                    case testEquality (typeRep @b) (typeRep @d) of+                        Just Refl ->+                            Right $+                                VB.generate nGroups foldGroup+                          where+                            foldGroup k =+                                let !s = VU.unsafeIndex offsets k+                                    !e = VU.unsafeIndex offsets (k + 1)+                                 in go s e seed+                            go !i !e !acc+                                | i >= e = acc+                                | otherwise =+                                    go (i + 1) e $!+                                        f acc (VU.unsafeIndex vec (VU.unsafeIndex valueIndices i))+                        Nothing ->+                            Left $+                                TypeMismatchException+                                    MkTypeErrorContext+                                        { userType = Right (typeRep @b)+                                        , expectedType = Right (typeRep @d)+                                        , callingFunctionName = Just "foldDirectGroups"+                                        , errorColumnName = Nothing+                                        }+                BoxedColumn (vec :: VB.Vector d) ->+                    case testEquality (typeRep @b) (typeRep @d) of+                        Just Refl ->+                            Right $+                                VB.generate nGroups foldGroup+                          where+                            foldGroup k =+                                let !s = VU.unsafeIndex offsets k+                                    !e = VU.unsafeIndex offsets (k + 1)+                                 in go s e seed+                            go !i !e !acc+                                | i >= e = acc+                                | otherwise =+                                    go (i + 1) e $!+                                        f acc (VB.unsafeIndex vec (VU.unsafeIndex valueIndices i))+                        Nothing ->+                            Left $+                                TypeMismatchException+                                    MkTypeErrorContext+                                        { userType = Right (typeRep @b)+                                        , expectedType = Right (typeRep @d)+                                        , callingFunctionName = Just "foldDirectGroups"+                                        , errorColumnName = Nothing+                                        }+                OptionalColumn (vec :: VB.Vector (Maybe d)) ->+                    case testEquality (typeRep @b) (typeRep @(Maybe d)) of+                        Just Refl ->+                            Right $+                                VB.generate nGroups foldGroup+                          where+                            foldGroup k =+                                let !s = VU.unsafeIndex offsets k+                                    !e = VU.unsafeIndex offsets (k + 1)+                                 in go s e seed+                            go !i !e !acc+                                | i >= e = acc+                                | otherwise =+                                    go (i + 1) e $!+                                        f acc (VB.unsafeIndex vec (VU.unsafeIndex valueIndices i))+                        Nothing ->+                            Left $+                                TypeMismatchException+                                    MkTypeErrorContext+                                        { userType = Right (typeRep @b)+                                        , expectedType = Right (typeRep @(Maybe d))+                                        , callingFunctionName = Just "foldDirectGroups"+                                        , errorColumnName = Nothing+                                        }+{-# INLINEABLE foldDirectGroups #-}++{- | O(n) Seedless fold over groups using the first element of each group as seed.+Like 'foldDirectGroups' but for the case where no initial accumulator is available.+-}+foldl1DirectGroups ::+    forall a.+    (Columnable a) =>+    (a -> a -> a) ->+    Column ->+    VU.Vector Int ->+    VU.Vector Int ->+    Either DataFrameException (VB.Vector a)+foldl1DirectGroups f col valueIndices offsets+    | VU.length offsets <= 1 = Right VB.empty+    | otherwise =+        let !nGroups = VU.length offsets - 1+         in case col of+                UnboxedColumn (vec :: VU.Vector d) ->+                    case testEquality (typeRep @a) (typeRep @d) of+                        Just Refl ->+                            Right $+                                VB.generate nGroups foldGroup+                          where+                            foldGroup k =+                                let !s = VU.unsafeIndex offsets k+                                    !e = VU.unsafeIndex offsets (k + 1)+                                    !seed = VU.unsafeIndex vec (VU.unsafeIndex valueIndices s)+                                 in go (s + 1) e seed+                            go !i !e !acc+                                | i >= e = acc+                                | otherwise =+                                    go (i + 1) e $!+                                        f acc (VU.unsafeIndex vec (VU.unsafeIndex valueIndices i))+                        Nothing ->+                            Left $+                                TypeMismatchException+                                    MkTypeErrorContext+                                        { userType = Right (typeRep @a)+                                        , expectedType = Right (typeRep @d)+                                        , callingFunctionName = Just "foldl1DirectGroups"+                                        , errorColumnName = Nothing+                                        }+                BoxedColumn (vec :: VB.Vector d) ->+                    case testEquality (typeRep @a) (typeRep @d) of+                        Just Refl ->+                            Right $+                                VB.generate nGroups foldGroup+                          where+                            foldGroup k =+                                let !s = VU.unsafeIndex offsets k+                                    !e = VU.unsafeIndex offsets (k + 1)+                                    !seed = VB.unsafeIndex vec (VU.unsafeIndex valueIndices s)+                                 in go (s + 1) e seed+                            go !i !e !acc+                                | i >= e = acc+                                | otherwise =+                                    go (i + 1) e $!+                                        f acc (VB.unsafeIndex vec (VU.unsafeIndex valueIndices i))+                        Nothing ->+                            Left $+                                TypeMismatchException+                                    MkTypeErrorContext+                                        { userType = Right (typeRep @a)+                                        , expectedType = Right (typeRep @d)+                                        , callingFunctionName = Just "foldl1DirectGroups"+                                        , errorColumnName = Nothing+                                        }+                OptionalColumn (vec :: VB.Vector (Maybe d)) ->+                    case testEquality (typeRep @a) (typeRep @(Maybe d)) of+                        Just Refl ->+                            Right $+                                VB.generate nGroups foldGroup+                          where+                            foldGroup k =+                                let !s = VU.unsafeIndex offsets k+                                    !e = VU.unsafeIndex offsets (k + 1)+                                    !seed = VB.unsafeIndex vec (VU.unsafeIndex valueIndices s)+                                 in go (s + 1) e seed+                            go !i !e !acc+                                | i >= e = acc+                                | otherwise =+                                    go (i + 1) e $!+                                        f acc (VB.unsafeIndex vec (VU.unsafeIndex valueIndices i))+                        Nothing ->+                            Left $+                                TypeMismatchException+                                    MkTypeErrorContext+                                        { userType = Right (typeRep @a)+                                        , expectedType = Right (typeRep @(Maybe d))+                                        , callingFunctionName = Just "foldl1DirectGroups"+                                        , errorColumnName = Nothing+                                        }+{-# INLINEABLE foldl1DirectGroups #-}++{- | O(n) fold over groups by scanning the column LINEARLY.+rowToGroup[i] = group index for row i.+Avoids random column reads; random writes go to the accumulator array which is+small (nGroups entries) and typically cache-resident.+When @acc@ is unboxable, uses an unboxed mutable vector for the accumulator+array, eliminating pointer indirection on every read/write.+-}+foldLinearGroups ::+    forall b acc.+    (Columnable b, Columnable acc) =>+    (acc -> b -> acc) ->+    acc ->+    Column ->+    VU.Vector Int -> -- rowToGroup (length n)+    Int -> -- nGroups+    Either DataFrameException Column+foldLinearGroups f seed col rowToGroup nGroups+    | nGroups == 0 = Right (fromVector @acc VB.empty)+    | otherwise = case col of+        UnboxedColumn (vec :: VU.Vector d) ->+            case testEquality (typeRep @b) (typeRep @d) of+                Just Refl ->+                    Right $+                        unsafePerformIO $+                            runWith+                                ( \readAt writeAt ->+                                    VU.iforM_ vec $ \row x -> do+                                        let !k = VU.unsafeIndex rowToGroup row+                                        cur <- readAt k+                                        writeAt k $! f cur x+                                )+                Nothing ->+                    Left $+                        TypeMismatchException+                            MkTypeErrorContext+                                { userType = Right (typeRep @b)+                                , expectedType = Right (typeRep @d)+                                , callingFunctionName = Just "foldLinearGroups"+                                , errorColumnName = Nothing+                                }+        BoxedColumn (vec :: VB.Vector d) ->+            case testEquality (typeRep @b) (typeRep @d) of+                Just Refl ->+                    Right $+                        unsafePerformIO $+                            runWith+                                ( \readAt writeAt ->+                                    VB.iforM_ vec $ \row x -> do+                                        let !k = VU.unsafeIndex rowToGroup row+                                        cur <- readAt k+                                        writeAt k $! f cur x+                                )+                Nothing ->+                    Left $+                        TypeMismatchException+                            MkTypeErrorContext+                                { userType = Right (typeRep @b)+                                , expectedType = Right (typeRep @d)+                                , callingFunctionName = Just "foldLinearGroups"+                                , errorColumnName = Nothing+                                }+        OptionalColumn (vec :: VB.Vector (Maybe d)) ->+            case testEquality (typeRep @b) (typeRep @(Maybe d)) of+                Just Refl ->+                    Right $+                        unsafePerformIO $+                            runWith+                                ( \readAt writeAt ->+                                    VB.iforM_ vec $ \row x -> do+                                        let !k = VU.unsafeIndex rowToGroup row+                                        cur <- readAt k+                                        writeAt k $! f cur x+                                )+                Nothing ->+                    Left $+                        TypeMismatchException+                            MkTypeErrorContext+                                { userType = Right (typeRep @b)+                                , expectedType = Right (typeRep @(Maybe d))+                                , callingFunctionName = Just "foldLinearGroups"+                                , errorColumnName = Nothing+                                }+  where+    -- \| Allocate accumulators, run the traversal, return a frozen Column.+    -- When @acc@ is unboxable, uses an unboxed mutable vector (no pointer+    -- indirection per read/write) and returns UnboxedColumn directly —+    -- avoiding a round-trip through VB.Vector.+    runWith :: ((Int -> IO acc) -> (Int -> acc -> IO ()) -> IO ()) -> IO Column+    runWith body = case sUnbox @acc of+        STrue -> do+            accs <- VUM.replicate nGroups seed+            body (VUM.unsafeRead accs) (VUM.unsafeWrite accs)+            UnboxedColumn <$> VU.unsafeFreeze accs+        SFalse -> do+            accs <- VBM.replicate nGroups seed+            body (VBM.unsafeRead accs) (VBM.unsafeWrite accs)+            fromVector @acc <$> VB.unsafeFreeze accs+    {-# INLINE runWith #-}+{-# INLINEABLE foldLinearGroups #-}+ headColumn :: forall a. (Columnable a) => Column -> Either DataFrameException a headColumn (BoxedColumn (col :: VB.Vector b)) = case testEquality (typeRep @a) (typeRep @b) of     Just Refl ->@@ -736,6 +1079,15 @@                     else VBM.unsafeWrite col i value >> return (Right True)                 )             Nothing -> return (Left value)+writeColumn i value (MOptionalColumn (col :: VBM.IOVector (Maybe a))) =+    let+     in case testEquality (typeRep @a) (typeRep @T.Text) of+            Just Refl ->+                ( if isNullish value+                    then VBM.unsafeWrite col i Nothing >> return (Left $! value)+                    else VBM.unsafeWrite col i (Just value) >> return (Right True)+                )+            Nothing -> return (Left value) writeColumn i value (MUnboxedColumn (col :: VUM.IOVector a)) =     case testEquality (typeRep @a) (typeRep @Int) of         Just Refl -> case readInt value of@@ -749,6 +1101,7 @@ {-# INLINE writeColumn #-}  freezeColumn' :: [(Int, T.Text)] -> MutableColumn -> IO Column+freezeColumn' nulls (MOptionalColumn col) = OptionalColumn <$> VB.unsafeFreeze col freezeColumn' nulls (MBoxedColumn col)     | null nulls = BoxedColumn <$> VB.unsafeFreeze col     | all (isNullish . snd) nulls =@@ -874,6 +1227,34 @@ E.g. combining Column containing [1,2] with Column containing ["a","b"] will result in a Column containing [Left 1, Left 2, Right "a", Right "b"]. -}++{- | O(n) Concatenate a list of same-type columns in a single allocation.+All columns must have the same constructor and element type (as they will+within a single Parquet column). Calls 'error' on mismatch.+-}+concatManyColumns :: [Column] -> Column+concatManyColumns [] = fromList ([] :: [Maybe Int])+concatManyColumns [c] = c+concatManyColumns (c0 : cs) = case c0 of+    OptionalColumn v0 ->+        let getVec (OptionalColumn v) = case testEquality (typeOf v0) (typeOf v) of+                Just Refl -> v+                Nothing -> error "concatManyColumns: OptionalColumn type mismatch"+            getVec _ = error "concatManyColumns: column constructor mismatch"+         in OptionalColumn (VB.concat (v0 : map getVec cs))+    BoxedColumn v0 ->+        let getVec (BoxedColumn v) = case testEquality (typeOf v0) (typeOf v) of+                Just Refl -> v+                Nothing -> error "concatManyColumns: BoxedColumn type mismatch"+            getVec _ = error "concatManyColumns: column constructor mismatch"+         in BoxedColumn (VB.concat (v0 : map getVec cs))+    UnboxedColumn v0 ->+        let getVec (UnboxedColumn v) = case testEquality (typeOf v0) (typeOf v) of+                Just Refl -> v+                Nothing -> error "concatManyColumns: UnboxedColumn type mismatch"+            getVec _ = error "concatManyColumns: column constructor mismatch"+         in UnboxedColumn (VU.concat (v0 : map getVec cs))+ concatColumnsEither :: Column -> Column -> Column concatColumnsEither (OptionalColumn left) (OptionalColumn right) = case testEquality (typeOf left) (typeOf right) of     Nothing ->@@ -912,6 +1293,38 @@         Just Refl -> OptionalColumn $ fmap Just (VG.convert left) <> right         Nothing ->             OptionalColumn $ fmap (Just . Left) (VG.convert left) <> fmap (fmap Right) right++-- | Allocate a mutable column of size @n@ matching the constructor/type of the given column.+newMutableColumn :: Int -> Column -> IO MutableColumn+newMutableColumn n (OptionalColumn (_ :: VB.Vector (Maybe a))) =+    MOptionalColumn <$> (VBM.new n :: IO (VBM.IOVector (Maybe a)))+newMutableColumn n (BoxedColumn (_ :: VB.Vector a)) =+    MBoxedColumn <$> (VBM.new n :: IO (VBM.IOVector a))+newMutableColumn n (UnboxedColumn (_ :: VU.Vector a)) =+    MUnboxedColumn <$> (VUM.new n :: IO (VUM.IOVector a))++-- | Copy a column chunk into a mutable column starting at offset @off@.+copyIntoMutableColumn :: MutableColumn -> Int -> Column -> IO ()+copyIntoMutableColumn (MOptionalColumn (mv :: VBM.IOVector (Maybe b))) off (OptionalColumn (v :: VB.Vector (Maybe a))) =+    case testEquality (typeRep @a) (typeRep @b) of+        Just Refl -> VG.imapM_ (\i x -> VBM.unsafeWrite mv (off + i) x) v+        Nothing -> error "copyIntoMutableColumn: Optional type mismatch"+copyIntoMutableColumn (MBoxedColumn (mv :: VBM.IOVector b)) off (BoxedColumn (v :: VB.Vector a)) =+    case testEquality (typeRep @a) (typeRep @b) of+        Just Refl -> VG.imapM_ (\i x -> VBM.unsafeWrite mv (off + i) x) v+        Nothing -> error "copyIntoMutableColumn: Boxed type mismatch"+copyIntoMutableColumn (MUnboxedColumn (mv :: VUM.IOVector b)) off (UnboxedColumn (v :: VU.Vector a)) =+    case testEquality (typeRep @a) (typeRep @b) of+        Just Refl -> VG.imapM_ (\i x -> VUM.unsafeWrite mv (off + i) x) v+        Nothing -> error "copyIntoMutableColumn: Unboxed type mismatch"+copyIntoMutableColumn _ _ _ =+    error "copyIntoMutableColumn: constructor mismatch"++-- | Freeze a mutable column into an immutable column.+freezeMutableColumn :: MutableColumn -> IO Column+freezeMutableColumn (MOptionalColumn mv) = OptionalColumn <$> VB.unsafeFreeze mv+freezeMutableColumn (MBoxedColumn mv) = BoxedColumn <$> VB.unsafeFreeze mv+freezeMutableColumn (MUnboxedColumn mv) = UnboxedColumn <$> VU.unsafeFreeze mv  {- | O(n) Converts a column to a list. Throws an exception if the wrong type is specified. 

@@ -13,6 +13,7 @@ import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU +import Control.DeepSeq (NFData (..), rnf) import Control.Exception (throw) import Data.Function (on) import Data.List (sortBy, transpose, (\\))@@ -36,6 +37,10 @@     , derivingExpressions :: M.Map T.Text UExpr     } +instance NFData DataFrame where+    rnf (DataFrame cols idx dims _exprs) =+        rnf cols `seq` rnf idx `seq` rnf dims+ {- | A record that contains information about how and what rows are grouped in the dataframe. This can only be used with `aggregate`.@@ -45,17 +50,21 @@     , groupedColumns :: [T.Text]     , valueIndices :: VU.Vector Int     , offsets :: VU.Vector Int+    , rowToGroup :: VU.Vector Int+    {- ^ rowToGroup[i] = group index for row i.  Length n (one per row).+    Built once in 'groupBy'; reused by every aggregation.+    -}     }  instance Show GroupedDataFrame where-    show (Grouped df cols indices os) =+    show (Grouped df cols _indices _os _rtg) =         printf             "{ keyColumns: %s groupedColumns: %s }"             (show cols)             (show (M.keys (columnIndices df) \\ cols))  instance Eq GroupedDataFrame where-    (==) (Grouped df cols indices os) (Grouped df' cols' indices' os') = (df == df') && (cols == cols')+    (==) (Grouped df cols _indices _os _rtg) (Grouped df' cols' _indices' _os' _rtg') = (df == df') && (cols == cols')  instance Eq DataFrame where     (==) :: DataFrame -> DataFrame -> Bool

@@ -13,6 +13,7 @@  module DataFrame.Internal.Expression where +import Control.DeepSeq (NFData (..)) import Data.String import qualified Data.Text as T import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))@@ -34,10 +35,24 @@     , binaryPrecedence :: Int     } +data MeanAcc = MeanAcc {-# UNPACK #-} !Double {-# UNPACK #-} !Int+    deriving (Show, Eq, Ord, Read)++instance NFData MeanAcc where+    rnf (MeanAcc _ _) = ()+ data AggStrategy a b where     CollectAgg ::         (VG.Vector v b, Typeable v) => T.Text -> (v b -> a) -> AggStrategy a b     FoldAgg :: T.Text -> Maybe a -> (a -> b -> a) -> AggStrategy a b+    MergeAgg ::+        (Columnable acc) =>+        T.Text ->+        acc ->+        (acc -> b -> acc) ->+        (acc -> acc -> acc) ->+        (acc -> a) ->+        AggStrategy a b  data Expr a where     Col :: (Columnable a) => T.Text -> Expr a@@ -132,7 +147,7 @@                 { binaryFn = (/)                 , binaryName = "divide"                 , binarySymbol = Just "/"-                , binaryCommutative = True+                , binaryCommutative = False                 , binaryPrecedence = 7                 }             )@@ -219,6 +234,7 @@     show (Binary op a b) = "(" ++ T.unpack (binaryName op) ++ " " ++ show a ++ " " ++ show b ++ ")"     show (Agg (CollectAgg op _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"     show (Agg (FoldAgg op _ _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"+    show (Agg (MergeAgg op _ _ _ _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"  normalize :: (Eq a, Ord a, Show a, Typeable a) => Expr a -> Expr a normalize expr = case expr of@@ -251,6 +267,7 @@     exprKey (Binary op e1 e2) = "4:" ++ T.unpack (binaryName op) ++ exprKey e1 ++ exprKey e2     exprKey (Agg (CollectAgg name _) e) = "5:" ++ T.unpack name ++ exprKey e     exprKey (Agg (FoldAgg name _ _) e) = "5:" ++ T.unpack name ++ exprKey e+    exprKey (Agg (MergeAgg name _ _ _ _) e) = "5:" ++ T.unpack name ++ exprKey e  instance (Eq a, Columnable a) => Eq (Expr a) where     (==) l r = eqNormalized (normalize l) (normalize r)@@ -270,6 +287,8 @@             n1 == n2 && e1 `exprEq` e2         eqNormalized (Agg (FoldAgg n1 _ _) e1) (Agg (FoldAgg n2 _ _) e2) =             n1 == n2 && e1 `exprEq` e2+        eqNormalized (Agg (MergeAgg n1 _ _ _ _) e1) (Agg (MergeAgg n2 _ _ _ _) e2) =+            n1 == n2 && e1 `exprEq` e2         eqNormalized _ _ = False  instance (Ord a, Columnable a) => Ord (Expr a) where@@ -284,6 +303,7 @@             compare (binaryName op1) (binaryName op2) <> exprComp a1 a2 <> exprComp b1 b2         (Agg (CollectAgg n1 _) e1', Agg (CollectAgg n2 _) e2') -> compare n1 n2 <> exprComp e1' e2'         (Agg (FoldAgg n1 _ _) e1', Agg (FoldAgg n2 _ _) e2') -> compare n1 n2 <> exprComp e1' e2'+        (Agg (MergeAgg n1 _ _ _ _) e1', Agg (MergeAgg n2 _ _ _ _) e2') -> compare n1 n2 <> exprComp e1' e2'         -- Different constructors - compare by priority         (Col _, _) -> LT         (_, Col _) -> GT@@ -372,3 +392,4 @@              in if prec > p then "(" ++ inner ++ ")" else inner         Agg (CollectAgg op _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"         Agg (FoldAgg op _ _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"+        Agg (MergeAgg op _ _ _ _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"

@@ -6,841 +6,481 @@ {-# LANGUAGE GADTs #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE PatternSynonyms #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE UndecidableInstances #-}--module DataFrame.Internal.Interpreter where--import Control.Monad.ST (runST)-import Data.Bifunctor-import qualified Data.Map as M-import qualified Data.Text as T-import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))-import qualified Data.Vector as V-import qualified Data.Vector.Mutable as VM-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM-import DataFrame.Errors-import DataFrame.Internal.Column-import DataFrame.Internal.DataFrame-import DataFrame.Internal.Expression-import DataFrame.Internal.Types-import Type.Reflection (TypeRep, Typeable, typeOf, typeRep, pattern App)--interpret ::-    forall a.-    (Columnable a) =>-    DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)-interpret df (Lit value) = case sUnbox @a of-    -- Specialize the creation of unboxed columns to avoid an extra allocation.-    STrue -> pure $ TColumn $ fromUnboxedVector $ VU.replicate (numRows df) value-    SFalse -> pure $ TColumn $ fromVector $ V.replicate (numRows df) value-interpret df (Col name) = maybe columnNotFound (pure . TColumn) (getColumn name df)-  where-    columnNotFound = Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df)--- Unary operations.-interpret df expr@(Unary (op :: UnaryOp b c) value) = first (handleInterpretException (show expr)) $ do-    (TColumn value') <- interpret df value-    fmap TColumn (mapColumn (unaryFn op) value')--- Variations of binary operations.-interpret df expr@(Binary (op :: BinaryOp c b a) left right) = first (handleInterpretException (show expr)) $ case (left, right) of-    (Lit left, Lit right) -> interpret df (Lit (binaryFn op left right))-    (Lit left, right) -> do-        -- If we have a literal then we don't have to materialise-        -- the column.-        (TColumn value') <- interpret df right-        fmap TColumn (mapColumn (binaryFn op left) value')-    (left, Lit right) -> do-        -- Same as the above except the right side is the-        -- literl.-        (TColumn value') <- interpret df left-        fmap TColumn (mapColumn (flip (binaryFn op) right) value')-    (_, _) -> do-        -- In the general case we interpret and zip.-        (TColumn left') <- interpret df left-        (TColumn right') <- interpret df right-        fmap TColumn (zipWithColumns (binaryFn op) left' right')--- Conditionals-interpret df expr@(If cond l r) = first (handleInterpretException (show expr)) $ do-    (TColumn conditions) <- interpret df cond-    (TColumn left) <- interpret df l-    (TColumn right) <- interpret df r-    let branch (c :: Bool) (l' :: a, r' :: a) = if c then l' else r'-    fmap TColumn (zipWithColumns branch conditions (zipColumns left right))-interpret df expression@(Agg (CollectAgg op (f :: v b -> c)) expr) = do-    (TColumn column) <- interpret df expr-    -- Helper for errors. Should probably find a way of throwing this-    -- without leaking the fact that we use `Vector` to users.-    let aggTypeError expected =-            TypeMismatchException-                ( MkTypeErrorContext-                    { userType = Right (typeRep @(v b))-                    , expectedType = Left expected :: Either String (TypeRep ())-                    , callingFunctionName = Just "interpret"-                    , errorColumnName = Nothing-                    }-                )-    let processColumn ::-            (Columnable d) => d -> Either DataFrameException (TypedColumn a)-        processColumn col = case testEquality (typeRep @(v b)) (typeOf col) of-            Just Refl -> interpret @c df (Lit (f col))-            Nothing -> Left $ aggTypeError (show (typeOf col))-    case column of-        (BoxedColumn col) -> processColumn col-        (OptionalColumn col) -> processColumn col-        (UnboxedColumn col) -> processColumn col-interpret df expression@(Agg (FoldAgg op (Just v) f) expr) = first (handleInterpretException (show expr)) $ do-    (TColumn column) <- interpret df expr-    value <- foldlColumn f v column-    pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value-interpret df expression@(Agg (FoldAgg op Nothing (f :: a -> b -> a)) expr) = first (handleInterpretException (show expr)) $ do-    (TColumn column) <- interpret df expr-    case testEquality (typeRep @a) (typeRep @b) of-        Just Refl -> do-            value <- foldl1Column f column-            pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value-        Nothing -> error "Type error"--data AggregationResult a-    = UnAggregated Column-    | Aggregated (TypedColumn a)--interpretAggregation ::-    forall a.-    (Columnable a) =>-    GroupedDataFrame -> Expr a -> Either DataFrameException (AggregationResult a)-interpretAggregation gdf (Lit value) =-    Right $-        Aggregated $-            TColumn $-                fromVector $-                    V.replicate (VU.length (offsets gdf) - 1) value-interpretAggregation gdf@(Grouped df names indices os) (Col name) = case getColumn name df of-    Nothing -> Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df)-    Just (BoxedColumn col) -> Right $ UnAggregated $ fromVector $ mkUnaggregatedColumnBoxed col os indices-    Just (OptionalColumn col) -> Right $ UnAggregated $ fromVector $ mkUnaggregatedColumnBoxed col os indices-    Just (UnboxedColumn col) ->-        Right $ UnAggregated $ fromVector $ mkUnaggregatedColumnUnboxed col os indices-interpretAggregation gdf expression@(Unary (op :: UnaryOp b a) expr) =-    case interpretAggregation @b gdf expr of-        Left (TypeMismatchException context) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show expr)-                        }-                    )-        Left e -> Left e-        Right (UnAggregated unaggregated) -> case unaggregated of-            BoxedColumn (col :: V.Vector c) -> case testEquality (typeRep @c) (typeRep @(V.Vector b)) of-                Just Refl -> case sUnbox @a of-                    SFalse -> Right $ UnAggregated $ fromVector $ V.map (V.map (unaryFn op)) col-                    STrue ->-                        Right $-                            UnAggregated $-                                fromVector $-                                    V.map (V.convert @V.Vector @a @VU.Vector . V.map (unaryFn op)) col-                Nothing -> case testEquality (typeRep @c) (typeRep @(VU.Vector b)) of-                    Nothing -> Left $ nestedTypeException @c @a (show expression)-                    Just Refl -> case (sUnbox @b, sUnbox @a) of-                        (SFalse, _) -> Left $ InternalException "Boxed type inside an unboxed column"-                        (STrue, STrue) -> Right $ UnAggregated $ fromVector $ V.map (VU.map (unaryFn op)) col-                        (STrue, _) ->-                            Right $ UnAggregated $ fromVector $ V.map (V.map (unaryFn op) . VU.convert) col-            _ -> Left $ InternalException "Aggregated into a non-boxed column"-        Right (Aggregated (TColumn aggregated)) -> case mapColumn (unaryFn op) aggregated of-            Left e -> Left e-            Right col -> Right $ Aggregated $ TColumn col-interpretAggregation gdf expression@(Binary (op :: t c b a) left (Lit (right :: b))) =-    interpretAggregation-        gdf-        (Unary (MkUnaryOp (flip (binaryFn op) right) "udf" Nothing) left)-interpretAggregation gdf expression@(Binary (op :: BinaryOp c b a) (Lit (left :: c)) right) =-    interpretAggregation-        gdf-        (Unary (MkUnaryOp (binaryFn op left) "udf" Nothing) right)-interpretAggregation gdf expression@(Binary (op :: BinaryOp c b a) left right) =-    case (interpretAggregation gdf left, interpretAggregation gdf right) of-        (Right (Aggregated (TColumn left')), Right (Aggregated (TColumn right'))) -> case zipWithColumns (binaryFn op) left' right' of-            Left e -> Left e-            Right col -> Right $ Aggregated $ TColumn col-        (Right (UnAggregated left'), Right (UnAggregated right')) -> case (left', right') of-            (BoxedColumn (l :: V.Vector m), BoxedColumn (r :: V.Vector n)) -> case testEquality (typeRep @m) (typeRep @(VU.Vector c)) of-                Just Refl -> case testEquality (typeRep @n) (typeRep @(VU.Vector b)) of-                    Just Refl -> case (sUnbox @c, sUnbox @b, sUnbox @a) of-                        (STrue, STrue, STrue) ->-                            Right $ UnAggregated $ fromVector $ V.zipWith (VU.zipWith (binaryFn op)) l r-                        (STrue, STrue, SFalse) ->-                            Right $-                                UnAggregated $-                                    fromVector $-                                        V.zipWith (\l' r' -> V.zipWith (binaryFn op) (V.convert l') (V.convert r')) l r-                        (_, _, _) -> Left $ InternalException "Boxed vectors contain unboxed types"-                    Nothing -> case testEquality (typeRep @n) (typeRep @(V.Vector b)) of-                        Just Refl -> case sUnbox @c of-                            STrue ->-                                Right $-                                    UnAggregated $-                                        fromVector $-                                            V.zipWith (V.zipWith (binaryFn op) . V.convert) l r-                            SFalse -> Left $ InternalException "Unboxed vectors contain boxed types"-                        Nothing -> Left $ nestedTypeException @n @b (show right)-                Nothing -> case testEquality (typeRep @m) (typeRep @(V.Vector c)) of-                    Nothing -> Left $ nestedTypeException @m @c (show left)-                    Just Refl -> case testEquality (typeRep @n) (typeRep @(VU.Vector b)) of-                        Just Refl -> case (sUnbox @b, sUnbox @a) of-                            (STrue, STrue) ->-                                Right $-                                    UnAggregated $-                                        fromVector $-                                            V.zipWith-                                                ( \l' r' ->-                                                    V.convert @V.Vector @a @VU.Vector $ V.zipWith (binaryFn op) l' (V.convert r')-                                                )-                                                l-                                                r-                            (STrue, SFalse) ->-                                Right $-                                    UnAggregated $-                                        fromVector $-                                            V.zipWith (\l' r' -> V.zipWith (binaryFn op) l' (V.convert r')) l r-                            (_, _) -> Left $ InternalException "Unboxed vectors contain boxed types"-                        Nothing -> case testEquality (typeRep @n) (typeRep @(V.Vector b)) of-                            Just Refl -> case sUnbox @a of-                                SFalse ->-                                    Right $-                                        UnAggregated $-                                            fromVector $-                                                V.zipWith (V.zipWith (binaryFn op) . V.convert) l r-                                STrue ->-                                    Right $-                                        UnAggregated $-                                            fromVector $-                                                V.zipWith-                                                    (\l' r' -> V.convert @V.Vector @a @VU.Vector $ V.zipWith (binaryFn op) l' r')-                                                    l-                                                    r-                            Nothing -> Left $ nestedTypeException @n @b (show right)-            _ -> Left $ InternalException "Aggregated into a non-boxed column"-        (Right _, Right _) ->-            Left $-                AggregatedAndNonAggregatedException (T.pack $ show left) (T.pack $ show right)-        (Left (TypeMismatchException context), _) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show left)-                        }-                    )-        (Left e, _) -> Left e-        (_, Left (TypeMismatchException context)) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show right)-                        }-                    )-        (_, Left e) -> Left e-interpretAggregation gdf expression@(If cond (Lit l) (Lit r)) =-    case interpretAggregation @Bool gdf cond of-        Right (Aggregated (TColumn conditions)) -> case mapColumn-            (\(c :: Bool) -> if c then l else r)-            conditions of-            Left e -> Left e-            Right v -> Right $ Aggregated (TColumn v)-        Right (UnAggregated conditions) -> case sUnbox @a of-            STrue -> case mapColumn-                (\(c :: VU.Vector Bool) -> VU.map (\c' -> if c' then l else r) c)-                conditions of-                Left (TypeMismatchException context) ->-                    Left $-                        TypeMismatchException-                            ( context-                                { callingFunctionName = Just "interpretAggregation"-                                , errorColumnName = Just (show expression)-                                }-                            )-                Left e -> Left e-                Right v -> Right $ UnAggregated v-            SFalse -> case mapColumn-                (\(c :: VU.Vector Bool) -> V.map (\c' -> if c' then l else r) (VU.convert c))-                conditions of-                Left (TypeMismatchException context) ->-                    Left $-                        TypeMismatchException-                            ( context-                                { callingFunctionName = Just "interpretAggregation"-                                , errorColumnName = Just (show expression)-                                }-                            )-                Left e -> Left e-                Right v -> Right $ UnAggregated v-        Left (TypeMismatchException context) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show cond)-                        }-                    )-        Left e -> Left e-interpretAggregation gdf expression@(If cond (Lit l) r) =-    case ( interpretAggregation @Bool gdf cond-         , interpretAggregation @a gdf r-         ) of-        ( Right (Aggregated (TColumn conditions))-            , Right (Aggregated (TColumn right))-            ) -> case zipWithColumns-                (\(c :: Bool) (r' :: a) -> if c then l else r')-                conditions-                right of-                Left e -> Left e-                Right v -> Right $ Aggregated (TColumn v)-        ( Right (UnAggregated conditions)-            , Right (UnAggregated right@(BoxedColumn (right' :: V.Vector c)))-            ) -> case testEquality (typeRep @(V.Vector a)) (typeRep @c) of-                Just Refl -> case zipWithColumns-                    ( \(c :: VU.Vector Bool) (r' :: V.Vector a) ->-                        V.zipWith-                            (\c' r'' -> if c' then l else r'')-                            (V.convert c)-                            r'-                    )-                    conditions-                    right of-                    Left (TypeMismatchException context) ->-                        Left $-                            TypeMismatchException-                                ( context-                                    { callingFunctionName = Just "interpretAggregation"-                                    , errorColumnName = Just (show expression)-                                    }-                                )-                    Left e -> Left e-                    Right v -> Right $ UnAggregated v-                Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @c) of-                    Nothing -> Left $ nestedTypeException @c @a (show expression)-                    Just Refl -> case sUnbox @a of-                        SFalse -> Left $ InternalException "Boxed type in unboxed column"-                        STrue -> case zipWithColumns-                            ( \(c :: VU.Vector Bool) (r' :: VU.Vector a) ->-                                VU.zipWith-                                    (\c' r'' -> if c' then l else r'')-                                    c-                                    r'-                            )-                            conditions-                            right of-                            Left (TypeMismatchException context) ->-                                Left $-                                    TypeMismatchException-                                        ( context-                                            { callingFunctionName = Just "interpretAggregation"-                                            , errorColumnName = Just (show expression)-                                            }-                                        )-                            Left e -> Left e-                            Right v -> Right $ UnAggregated v-        (Right _, Right _) ->-            Left $-                AggregatedAndNonAggregatedException (T.pack $ show l) (T.pack $ show r)-        (Left (TypeMismatchException context), _) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show cond)-                        }-                    )-        (Left e, _) -> Left e-        (_, Left (TypeMismatchException context)) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show r)-                        }-                    )-        (_, Left e) -> Left e-interpretAggregation gdf expression@(If cond l (Lit r)) =-    case ( interpretAggregation @Bool gdf cond-         , interpretAggregation @a gdf l-         ) of-        ( Right (Aggregated (TColumn conditions))-            , Right (Aggregated (TColumn left))-            ) -> case zipWithColumns-                (\(c :: Bool) (l' :: a) -> if c then l' else r)-                conditions-                left of-                Left e -> Left e-                Right v -> Right $ Aggregated (TColumn v)-        ( Right (UnAggregated conditions)-            , Right (UnAggregated left@(BoxedColumn (left' :: V.Vector c)))-            ) -> case testEquality (typeRep @(V.Vector a)) (typeRep @c) of-                Just Refl -> case zipWithColumns-                    ( \(c :: VU.Vector Bool) (l' :: V.Vector a) ->-                        V.zipWith-                            (\c' l'' -> if c' then l'' else r)-                            (V.convert c)-                            l'-                    )-                    conditions-                    left of-                    Left (TypeMismatchException context) ->-                        Left $-                            TypeMismatchException-                                ( context-                                    { callingFunctionName = Just "interpretAggregation"-                                    , errorColumnName = Just (show expression)-                                    }-                                )-                    Left e -> Left e-                    Right v -> Right $ UnAggregated v-                Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @c) of-                    Nothing -> Left $ nestedTypeException @c @a (show expression)-                    Just Refl -> case sUnbox @a of-                        SFalse -> Left $ InternalException "Boxed type in unboxed column"-                        STrue -> case zipWithColumns-                            ( \(c :: VU.Vector Bool) (l' :: VU.Vector a) ->-                                VU.zipWith-                                    (\c' l'' -> if c' then l'' else r)-                                    c-                                    l'-                            )-                            conditions-                            left of-                            Left (TypeMismatchException context) ->-                                Left $-                                    TypeMismatchException-                                        ( context-                                            { callingFunctionName = Just "interpretAggregation"-                                            , errorColumnName = Just (show expression)-                                            }-                                        )-                            Left e -> Left e-                            Right v -> Right $ UnAggregated v-        (Right _, Right _) ->-            Left $-                AggregatedAndNonAggregatedException (T.pack $ show l) (T.pack $ show r)-        (Left (TypeMismatchException context), _) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show cond)-                        }-                    )-        (Left e, _) -> Left e-        (_, Left (TypeMismatchException context)) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show r)-                        }-                    )-        (_, Left e) -> Left e-interpretAggregation gdf expression@(If cond l r) =-    case ( interpretAggregation @Bool gdf cond-         , interpretAggregation @a gdf l-         , interpretAggregation @a gdf r-         ) of-        ( Right (Aggregated (TColumn conditions))-            , Right (Aggregated (TColumn left))-            , Right (Aggregated (TColumn right))-            ) -> case zipWithColumns-                (\(c :: Bool) (l' :: a, r' :: a) -> if c then l' else r')-                conditions-                (zipColumns left right) of-                Left e -> Left e-                Right v -> Right $ Aggregated (TColumn v)-        ( Right (UnAggregated conditions)-            , Right (UnAggregated left@(BoxedColumn (left' :: V.Vector b)))-            , Right (UnAggregated right@(BoxedColumn (right' :: V.Vector c)))-            ) -> case testEquality (typeRep @b) (typeRep @c) of-                Nothing ->-                    Left $-                        TypeMismatchException-                            ( MkTypeErrorContext-                                { userType = Right (typeRep @b)-                                , expectedType = Right (typeRep @c)-                                , callingFunctionName = Just "interpretAggregation"-                                , errorColumnName = Just (show expression)-                                }-                            )-                Just Refl -> case testEquality (typeRep @(V.Vector a)) (typeRep @b) of-                    Just Refl -> case zipWithColumns-                        ( \(c :: VU.Vector Bool) (l' :: V.Vector a, r' :: V.Vector a) ->-                            V.zipWith-                                (\c' (l'', r'') -> if c' then l'' else r'')-                                (V.convert c)-                                (V.zip l' r')-                        )-                        conditions-                        (zipColumns left right) of-                        Left (TypeMismatchException context) ->-                            Left $-                                TypeMismatchException-                                    ( context-                                        { callingFunctionName = Just "interpretAggregation"-                                        , errorColumnName = Just (show expression)-                                        }-                                    )-                        Left e -> Left e-                        Right v -> Right $ UnAggregated v-                    Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @b) of-                        Nothing -> Left $ nestedTypeException @b @a (show expression)-                        Just Refl -> case sUnbox @a of-                            SFalse -> Left $ InternalException "Boxed type in unboxed column"-                            STrue -> case zipWithColumns-                                ( \(c :: VU.Vector Bool) (l' :: VU.Vector a, r' :: VU.Vector a) ->-                                    VU.zipWith-                                        (\c' (l'', r'') -> if c' then l'' else r'')-                                        c-                                        (VU.zip l' r')-                                )-                                conditions-                                (zipColumns left right) of-                                Left (TypeMismatchException context) ->-                                    Left $-                                        TypeMismatchException-                                            ( context-                                                { callingFunctionName = Just "interpretAggregation"-                                                , errorColumnName = Just (show expression)-                                                }-                                            )-                                Left e -> Left e-                                Right v -> Right $ UnAggregated v-        (Right _, Right _, Right _) ->-            Left $-                AggregatedAndNonAggregatedException (T.pack $ show l) (T.pack $ show r)-        (Left (TypeMismatchException context), _, _) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show cond)-                        }-                    )-        (Left e, _, _) -> Left e-        (_, Left (TypeMismatchException context), _) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show l)-                        }-                    )-        (_, Left e, _) -> Left e-        (_, _, Left (TypeMismatchException context)) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show r)-                        }-                    )-        (_, _, Left e) -> Left e-interpretAggregation gdf@(Grouped df names indices os) expression@(Agg (CollectAgg op (f :: v b -> c)) expr) =-    case interpretAggregation @b gdf expr of-        Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(v b)) (typeRep @d) of-            Nothing -> Left $ nestedTypeException @d @b (show expr)-            Just Refl -> case testEquality (typeRep @v) (typeRep @V.Vector) of-                Nothing -> Right $ Aggregated $ TColumn $ fromVector $ V.map (f . V.convert) col-                Just Refl -> Right $ Aggregated $ TColumn $ fromVector $ V.map f col-        Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column"-        Right (Aggregated (TColumn (BoxedColumn (col :: V.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of-            Just Refl -> case testEquality (typeRep @v) (typeRep @V.Vector) of-                Just Refl -> interpretAggregation @c gdf (Lit (f col))-                Nothing -> interpretAggregation @c gdf (Lit ((f . V.convert) col))-            Nothing ->-                Left $-                    TypeMismatchException-                        ( MkTypeErrorContext-                            { userType = Right (typeRep @b)-                            , expectedType = Right (typeRep @d)-                            , callingFunctionName = Just "interpretAggregation"-                            , errorColumnName = Just (show expr)-                            }-                        )-        Right (Aggregated (TColumn (UnboxedColumn (col :: VU.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of-            Just Refl -> case testEquality (typeRep @v) (typeRep @VU.Vector) of-                Just Refl -> interpretAggregation @c gdf (Lit (f col))-                Nothing -> interpretAggregation @c gdf (Lit ((f . VU.convert) col))-            Nothing ->-                Left $-                    TypeMismatchException-                        ( MkTypeErrorContext-                            { userType = Right (typeRep @b)-                            , expectedType = Right (typeRep @d)-                            , callingFunctionName = Just "interpretAggregation"-                            , errorColumnName = Just (show expr)-                            }-                        )-        Right (Aggregated (TColumn (OptionalColumn (col :: V.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of-            Just Refl -> case testEquality (typeRep @v) (typeRep @V.Vector) of-                Just Refl -> interpretAggregation @c gdf (Lit (f col))-                Nothing -> interpretAggregation @c gdf (Lit ((f . V.convert) col))-            Nothing ->-                Left $-                    TypeMismatchException-                        ( MkTypeErrorContext-                            { userType = Right (typeRep @b)-                            , expectedType = Right (typeRep @d)-                            , callingFunctionName = Just "interpretAggregation"-                            , errorColumnName = Just (show expr)-                            }-                        )-        (Left (TypeMismatchException context)) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show expression)-                        }-                    )-        (Left e) -> Left e-interpretAggregation gdf@(Grouped df names indices os) expression@(Agg (FoldAgg op Nothing (f :: a -> b -> a)) (Col name)) = case testEquality (typeRep @a) (typeRep @b) of-    Nothing -> error "Type mismatch"-    Just Refl -> case getColumn name df of-        Nothing -> Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df)-        Just (BoxedColumn (col :: V.Vector d)) -> case testEquality (typeRep @a) (typeRep @d) of-            Nothing -> error "Type mismatch"-            Just Refl ->-                Right $-                    Aggregated $-                        TColumn $-                            fromVector $-                                mkReducedColumnBoxed col os indices f-        Just (OptionalColumn (col :: V.Vector d)) -> case testEquality (typeRep @a) (typeRep @d) of-            Nothing -> error "Type mismatch"-            Just Refl ->-                Right $-                    Aggregated $-                        TColumn $-                            fromVector $-                                mkReducedColumnBoxed col os indices f-        Just (UnboxedColumn (col :: VU.Vector d)) -> case testEquality (typeRep @a) (typeRep @d) of-            Just Refl ->-                Right $-                    Aggregated $-                        TColumn $-                            fromUnboxedVector $-                                mkReducedColumnUnboxed col os indices f-            Nothing -> error "Type mismatch"-interpretAggregation gdf@(Grouped df names indices os) expression@(Agg (FoldAgg op Nothing (f :: a -> b -> a)) expr) = case testEquality (typeRep @a) (typeRep @b) of-    Nothing -> error "Type mismatch"-    Just Refl -> case interpretAggregation @a gdf expr of-        (Left (TypeMismatchException context)) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show expression)-                        }-                    )-        (Left e) -> Left e-        Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(V.Vector a)) (typeRep @d) of-            Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @d) of-                Nothing -> Left $ nestedTypeException @d @a (show expr)-                Just Refl -> case sUnbox @a of-                    STrue ->-                        Right $-                            Aggregated $-                                TColumn $-                                    fromVector $-                                        V.map (VU.foldl1' f) col-                    SFalse -> Left $ InternalException "Boxed type inside an unboxed column"-            Just Refl ->-                Right $-                    Aggregated $-                        TColumn $-                            fromVector $-                                V.map (V.foldl1' f) col-        Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column"-        Right (Aggregated (TColumn column)) -> case foldl1Column f column of-            Left e -> Left e-            Right value -> interpretAggregation @a gdf (Lit value)-interpretAggregation gdf@(Grouped df names indices os) expression@(Agg (FoldAgg op (Just s) (f :: a -> b -> a)) expr) =-    case interpretAggregation gdf expr of-        (Left (TypeMismatchException context)) ->-            Left $-                TypeMismatchException-                    ( context-                        { callingFunctionName = Just "interpretAggregation"-                        , errorColumnName = Just (show expression)-                        }-                    )-        (Left e) -> Left e-        Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(V.Vector b)) (typeRep @d) of-            Just Refl -> Right $ Aggregated $ TColumn $ fromVector $ V.map (V.foldl' f s) col-            Nothing -> case testEquality (typeRep @(VU.Vector b)) (typeRep @d) of-                Just Refl -> case sUnbox @b of-                    STrue ->-                        Right $ Aggregated $ TColumn $ fromVector $ V.map (VU.foldl' f s) col-                    SFalse -> Left $ InternalException "Boxed type inside an unboxed column"-                Nothing -> Left $ nestedTypeException @d @b (show expr)-        Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column"-        Right (Aggregated (TColumn column)) -> case foldlColumn f s column of-            Left e -> Left e-            Right value -> interpretAggregation @a gdf (Lit value)--handleInterpretException :: String -> DataFrameException -> DataFrameException-handleInterpretException errorLocation (TypeMismatchException context) = mkTypeMismatchException (Just "interpret") (Just errorLocation) context-handleInterpretException _ e = e--numRows :: DataFrame -> Int-numRows df = fst (dataframeDimensions df)--mkUnaggregatedColumnBoxed ::-    forall a.-    (Columnable a) =>-    V.Vector a -> VU.Vector Int -> VU.Vector Int -> V.Vector (V.Vector a)-mkUnaggregatedColumnBoxed col os indices =-    let-        sorted = V.unsafeBackpermute col (V.convert indices)-        n i = os `VU.unsafeIndex` (i + 1) - (os `VU.unsafeIndex` i)-        start i = os `VU.unsafeIndex` i-     in-        V.generate-            (VU.length os - 1)-            ( \i ->-                V.unsafeSlice (start i) (n i) sorted-            )--mkUnaggregatedColumnUnboxed ::-    forall a.-    (Columnable a, VU.Unbox a) =>-    VU.Vector a -> VU.Vector Int -> VU.Vector Int -> V.Vector (VU.Vector a)-mkUnaggregatedColumnUnboxed col os indices =-    let-        sorted = VU.unsafeBackpermute col indices-        n i = os `VU.unsafeIndex` (i + 1) - (os `VU.unsafeIndex` i)-        start i = os `VU.unsafeIndex` i-     in-        V.generate-            (VU.length os - 1)-            ( \i ->-                VU.unsafeSlice (start i) (n i) sorted-            )--mkAggregatedColumnUnboxed ::-    forall a b.-    (Columnable a, VU.Unbox a, Columnable b, VU.Unbox b) =>-    VU.Vector a ->-    VU.Vector Int ->-    VU.Vector Int ->-    (VU.Vector a -> b) ->-    VU.Vector b-mkAggregatedColumnUnboxed col os indices f =-    let-        sorted = VU.unsafeBackpermute col indices-        n i = os `VU.unsafeIndex` (i + 1) - (os `VU.unsafeIndex` i)-        start i = os `VU.unsafeIndex` i-     in-        VU.generate-            (VU.length os - 1)-            ( \i ->-                f (VU.unsafeSlice (start i) (n i) sorted)-            )--mkReducedColumnUnboxed ::-    forall a.-    (VU.Unbox a) =>-    VU.Vector a ->-    VU.Vector Int ->-    VU.Vector Int ->-    (a -> a -> a) ->-    VU.Vector a-mkReducedColumnUnboxed col os indices f = runST $ do-    let len = VU.length os - 1-    mvec <- VUM.unsafeNew len--    let loopOut i-            | i == len = return ()-            | otherwise = do-                let !start = os `VU.unsafeIndex` i-                let !end = os `VU.unsafeIndex` (i + 1)-                let !initVal = col `VU.unsafeIndex` (indices `VU.unsafeIndex` start)--                let loopIn !acc !idx-                        | idx == end = acc-                        | otherwise =-                            let val = col `VU.unsafeIndex` (indices `VU.unsafeIndex` idx)-                             in loopIn (f acc val) (idx + 1)-                let !finalVal = loopIn initVal (start + 1)-                VUM.unsafeWrite mvec i finalVal-                loopOut (i + 1)--    loopOut 0-    VU.unsafeFreeze mvec-{-# INLINE mkReducedColumnUnboxed #-}--mkReducedColumnBoxed ::-    V.Vector a ->-    VU.Vector Int ->-    VU.Vector Int ->-    (a -> a -> a) ->-    V.Vector a-mkReducedColumnBoxed col os indices f = runST $ do-    let len = VU.length os - 1-    mvec <- VM.unsafeNew len--    let loopOut i-            | i == len = return ()-            | otherwise = do-                let start = os `VU.unsafeIndex` i-                let end = os `VU.unsafeIndex` (i + 1)-                let initVal = col `V.unsafeIndex` (indices `VU.unsafeIndex` start)--                let loopIn !acc idx-                        | idx == end = acc-                        | otherwise =-                            let val = col `V.unsafeIndex` (indices `VU.unsafeIndex` idx)-                             in loopIn (f acc val) (idx + 1)-                let !finalVal = loopIn initVal (start + 1)-                VM.unsafeWrite mvec i finalVal-                loopOut (i + 1)--    loopOut 0-    V.unsafeFreeze mvec-{-# INLINE mkReducedColumnBoxed #-}--nestedTypeException ::-    forall a b. (Typeable a, Typeable b) => String -> DataFrameException-nestedTypeException expression = case typeRep @a of-    App t1 t2 ->-        TypeMismatchException-            ( MkTypeErrorContext-                { userType = Left (show (typeRep @b)) :: Either String (TypeRep ())-                , expectedType = Left (show (typeRep @a)) :: Either String (TypeRep ())-                , callingFunctionName = Just "interpretAggregation"-                , errorColumnName = Just expression-                }-            )-    t ->-        TypeMismatchException-            ( MkTypeErrorContext-                { userType = Right (typeRep @(VU.Vector b))-                , expectedType = Right (typeRep @b)-                , callingFunctionName = Just "interpretAggregation"-                , errorColumnName = Just expression-                }-            )--mkTypeMismatchException ::-    (Typeable a, Typeable b) =>-    Maybe String -> Maybe String -> TypeErrorContext a b -> DataFrameException-mkTypeMismatchException callPoint errorLocation context =-    TypeMismatchException-        ( context-            { callingFunctionName = callPoint-            , errorColumnName = errorLocation-            }-        )+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE UndecidableInstances #-}++module DataFrame.Internal.Interpreter (+    -- * New core API+    Value (..),+    Ctx (..),+    eval,+    materialize,++    -- * Backward-compatible API+    interpret,+    interpretAggregation,+    AggregationResult (..),+) where++import Data.Bifunctor (first)+import qualified Data.Map as M+import qualified Data.Text as T+import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))+import qualified Data.Vector as V+import qualified Data.Vector.Generic as VG+import qualified Data.Vector.Unboxed as VU+import DataFrame.Errors+import DataFrame.Internal.Column+import DataFrame.Internal.DataFrame+import DataFrame.Internal.Expression+import DataFrame.Internal.Types+import Type.Reflection (Typeable, typeRep)++-------------------------------------------------------------------------------+-- Value: the unified result type+-------------------------------------------------------------------------------++{- | The result of interpreting an expression.  Keeps literals as scalars+until the point where a concrete column is needed, avoiding premature+broadcast allocations.+-}+data Value a where+    -- | A single value, not yet broadcast to any length.+    Scalar :: (Columnable a) => a -> Value a+    {- | A flat column (one element per row in the flat case, or one+    element per group after aggregation).+    -}+    Flat :: (Columnable a) => Column -> Value a+    {- | A grouped column: one 'Column' slice per group.  Only produced+    when interpreting inside a 'GroupCtx'.+    -}+    Group :: (Columnable a) => V.Vector Column -> Value a++-- | The interpretation context.+data Ctx+    = FlatCtx DataFrame+    | GroupCtx GroupedDataFrame++-------------------------------------------------------------------------------+-- Materialisation+-------------------------------------------------------------------------------++{- | Force a 'Value' into a flat 'Column' of the given length.  Scalars+are broadcast; flat columns are returned as-is.+-}+materialize :: forall a. (Columnable a) => Int -> Value a -> Column+materialize n (Scalar v) = broadcastScalar @a n v+materialize _ (Flat c) = c+materialize _ (Group _) =+    error "materialize: cannot flatten a grouped value to a single column"++{- | Replicate a scalar to a column of length @n@, choosing the most+efficient representation.+-}+broadcastScalar :: forall a. (Columnable a) => Int -> a -> Column+broadcastScalar n v = case sUnbox @a of+    STrue -> fromUnboxedVector (VU.replicate n v)+    SFalse -> fromVector (V.replicate n v)++-------------------------------------------------------------------------------+-- Lifting: the core combinators+-------------------------------------------------------------------------------++-- | Apply a pure function to a 'Value'.+liftValue ::+    (Columnable b, Columnable a) =>+    (b -> a) -> Value b -> Either DataFrameException (Value a)+liftValue f (Scalar v) = Right (Scalar (f v))+liftValue f (Flat col) = Flat <$> mapColumn f col+liftValue f (Group gs) = Group <$> V.mapM (mapColumn f) gs++{- | Apply a binary function to two 'Value's.  When one side is a+'Scalar' the operation degenerates to a 'liftValue' — this is how the+old @Binary op (Lit l) right@ special cases are recovered without+explicit pattern matches in the evaluator.+-}+liftValue2 ::+    (Columnable c, Columnable b, Columnable a) =>+    (c -> b -> a) ->+    Value c ->+    Value b ->+    Either DataFrameException (Value a)+liftValue2 f (Scalar l) (Scalar r) = Right (Scalar (f l r))+liftValue2 f (Scalar l) v = liftValue (f l) v+liftValue2 f v (Scalar r) = liftValue (`f` r) v+liftValue2 f (Flat l) (Flat r) = Flat <$> zipWithColumns f l r+liftValue2 f (Group ls) (Group rs)+    | V.length ls == V.length rs =+        Group <$> V.zipWithM (zipWithColumns f) ls rs+-- Shape mismatches: aggregated vs. non-aggregated.+liftValue2 _ (Flat _) (Group _) =+    Left $ AggregatedAndNonAggregatedException "aggregated" "non-aggregated"+liftValue2 _ (Group _) (Flat _) =+    Left $ AggregatedAndNonAggregatedException "non-aggregated" "aggregated"+liftValue2 _ (Group _) (Group _) =+    Left $ InternalException "Group count mismatch in binary operation"++-- | Branch on a boolean 'Value', selecting from two same-typed 'Value's.+branchValue ::+    forall a.+    (Columnable a) =>+    Value Bool ->+    Value a ->+    Value a ->+    Either DataFrameException (Value a)+branchValue (Scalar True) l _ = Right l+branchValue (Scalar False) _ r = Right r+branchValue cond (Scalar l) (Scalar r) =+    liftValue (\c -> if c then l else r) cond+branchValue cond (Scalar l) r =+    liftValue2 (\c rv -> if c then l else rv) cond r+branchValue cond l (Scalar r) =+    liftValue2 (\c lv -> if c then lv else r) cond l+branchValue (Flat cc) (Flat lc) (Flat rc) =+    Flat <$> branchColumn @a cc lc rc+branchValue (Group cgs) (Group lgs) (Group rgs)+    | V.length cgs == V.length lgs+        && V.length lgs == V.length rgs =+        Group+            <$> V.generateM+                (V.length cgs)+                ( \i ->+                    branchColumn @a (cgs V.! i) (lgs V.! i) (rgs V.! i)+                )+branchValue _ _ _ =+    Left $+        AggregatedAndNonAggregatedException+            "if-then-else branches"+            "mismatched shapes"++{- | Low-level column branch: given a boolean column and two same-typed+columns, produce the element-wise selection.+-}+branchColumn ::+    forall a.+    (Columnable a) =>+    Column ->+    Column ->+    Column ->+    Either DataFrameException Column+branchColumn cc lc rc = do+    cs <- toVector @Bool @V.Vector cc+    ls <- toVector @a @V.Vector lc+    rs <- toVector @a @V.Vector rc+    pure $+        fromVector @a $+            V.zipWith3 (\c l r -> if c then l else r) cs ls rs++-------------------------------------------------------------------------------+-- Error enrichment+-------------------------------------------------------------------------------++{- | Wrap an interpretation step so that any 'TypeMismatchException' gets+annotated with the expression that was being evaluated.+-}+addContext ::+    (Show a) => Expr a -> Either DataFrameException b -> Either DataFrameException b+addContext expr = first (enrichError (show expr))++enrichError :: String -> DataFrameException -> DataFrameException+enrichError loc (TypeMismatchException ctx) =+    TypeMismatchException+        ctx+            { callingFunctionName =+                callingFunctionName ctx <|+> Just "eval"+            , errorColumnName =+                errorColumnName ctx <|+> Just loc+            }+  where+    -- Prefer the existing value; fall back to the new one.+    Nothing <|+> b = b+    a <|+> _ = a+enrichError _ e = e++-------------------------------------------------------------------------------+-- Group slicing+-------------------------------------------------------------------------------++{- | Given a flat column and grouping metadata, produce one 'Column' per+group.  Each result column is an O(1) slice into a sorted copy of the+input — the sort happens once, not per-group.+-}+sliceGroups :: Column -> VU.Vector Int -> VU.Vector Int -> V.Vector Column+sliceGroups col os indices = case col of+    BoxedColumn vec ->+        let !sorted = V.unsafeBackpermute vec (V.convert indices)+         in V.generate nGroups $ \i ->+                BoxedColumn (V.unsafeSlice (start i) (len i) sorted)+    UnboxedColumn vec ->+        let !sorted = VU.unsafeBackpermute vec indices+         in V.generate nGroups $ \i ->+                UnboxedColumn (VU.unsafeSlice (start i) (len i) sorted)+    OptionalColumn vec ->+        let !sorted = V.unsafeBackpermute vec (V.convert indices)+         in V.generate nGroups $ \i ->+                OptionalColumn (V.unsafeSlice (start i) (len i) sorted)+  where+    !nGroups = VU.length os - 1+    start i = os `VU.unsafeIndex` i+    len i = os `VU.unsafeIndex` (i + 1) - start i+{-# INLINE sliceGroups #-}++numGroups :: GroupedDataFrame -> Int+numGroups gdf = VU.length (offsets gdf) - 1++-------------------------------------------------------------------------------+-- eval: the unified interpreter+-------------------------------------------------------------------------------++{- | Evaluate an expression in a given context, producing a 'Value'.+This single function replaces both the old @interpret@ (flat) and+@interpretAggregation@ (grouped) code paths.+-}+eval ::+    forall a.+    (Columnable a) =>+    Ctx -> Expr a -> Either DataFrameException (Value a)+-- Leaves -----------------------------------------------------------------++eval _ (Lit v) = Right (Scalar v)+eval (FlatCtx df) (Col name) =+    case getColumn name df of+        Nothing ->+            Left $+                ColumnNotFoundException name "" (M.keys $ columnIndices df)+        Just c -> Right (Flat c)+eval (GroupCtx gdf) (Col name) =+    case getColumn name (fullDataframe gdf) of+        Nothing ->+            Left $+                ColumnNotFoundException+                    name+                    ""+                    (M.keys $ columnIndices $ fullDataframe gdf)+        Just c ->+            Right+                ( Group+                    (sliceGroups c (offsets gdf) (valueIndices gdf))+                )+-- Unary ------------------------------------------------------------------++eval ctx expr@(Unary (op :: UnaryOp b a) inner) = addContext expr $ do+    v <- eval @b ctx inner+    liftValue (unaryFn op) v++-- Binary -----------------------------------------------------------------++eval ctx expr@(Binary (op :: BinaryOp c b a) left right) =+    addContext expr $ do+        l <- eval @c ctx left+        r <- eval @b ctx right+        liftValue2 (binaryFn op) l r++-- If ---------------------------------------------------------------------++eval ctx expr@(If cond l r) = addContext expr $ do+    c <- eval @Bool ctx cond+    lv <- eval @a ctx l+    rv <- eval @a ctx r+    branchValue c lv rv++-- Fast path: FoldAgg (seeded) on a bare Col in GroupCtx.+-- Avoids the O(n) backpermute in sliceGroups by folding directly over+-- permuted indices.  Only matches when inner is exactly (Col name).++eval (GroupCtx gdf) expr@(Agg (FoldAgg _ (Just seed) (f :: a -> b -> a)) (Col name :: Expr b)) =+    addContext expr $+        case getColumn name (fullDataframe gdf) of+            Nothing ->+                Left $+                    ColumnNotFoundException+                        name+                        ""+                        (M.keys $ columnIndices $ fullDataframe gdf)+            Just col ->+                Flat <$> foldLinearGroups @b @a f seed col (rowToGroup gdf) (numGroups gdf)+-- Fast path: FoldAgg (seedless) on a bare Col in GroupCtx.++eval (GroupCtx gdf) expr@(Agg (FoldAgg _ Nothing (f :: a -> b -> a)) (Col name :: Expr b)) =+    addContext expr $+        case testEquality (typeRep @a) (typeRep @b) of+            Nothing ->+                Left $+                    InternalException+                        "Type mismatch in seedless fold: \+                        \accumulator and element types must match"+            Just Refl ->+                case getColumn name (fullDataframe gdf) of+                    Nothing ->+                        Left $+                            ColumnNotFoundException+                                name+                                ""+                                (M.keys $ columnIndices $ fullDataframe gdf)+                    Just col ->+                        Flat . fromVector+                            <$> foldl1DirectGroups @b f col (valueIndices gdf) (offsets gdf)+-- Fast path: MergeAgg on a bare Col in GroupCtx.++eval+    (GroupCtx gdf)+    expr@( Agg+                (MergeAgg _ seed (step :: acc -> b -> acc) _ (finalize :: acc -> a))+                (Col name :: Expr b)+            ) =+        addContext expr $+            case getColumn name (fullDataframe gdf) of+                Nothing ->+                    Left $+                        ColumnNotFoundException+                            name+                            ""+                            (M.keys $ columnIndices $ fullDataframe gdf)+                Just col ->+                    Flat+                        <$> ( foldLinearGroups @b step seed col (rowToGroup gdf) (numGroups gdf)+                                >>= mapColumn finalize+                            )+-- Aggregation: CollectAgg ------------------------------------------------++eval ctx expr@(Agg (CollectAgg _ (f :: v b -> a)) inner) =+    addContext expr $ do+        v <- eval @b ctx inner+        case v of+            Scalar _ ->+                Left $+                    InternalException+                        "Cannot apply a collection aggregation to a scalar"+            Flat col ->+                Scalar <$> applyCollect @v @b @a f col+            Group gs ->+                Flat . fromVector+                    <$> V.mapM (applyCollect @v @b @a f) gs++-- Aggregation: FoldAgg with seed -----------------------------------------++eval ctx expr@(Agg (FoldAgg _ (Just seed) (f :: a -> b -> a)) inner) =+    addContext expr $ do+        v <- eval @b ctx inner+        case v of+            Scalar _ ->+                Left $+                    InternalException+                        "Cannot apply a fold aggregation to a scalar"+            Flat col ->+                Scalar <$> foldlColumn @b @a f seed col+            Group gs ->+                Flat . fromVector+                    <$> V.mapM (foldlColumn @b @a f seed) gs++-- Aggregation: MergeAgg --------------------------------------------------++eval+    ctx+    expr@( Agg+                (MergeAgg _ seed (step :: acc -> b -> acc) _ (finalize :: acc -> a))+                (inner :: Expr b)+            ) =+        addContext expr $ do+            v <- eval @b ctx inner+            case v of+                Scalar _ ->+                    Left $+                        InternalException+                            "Cannot apply a merge aggregation to a scalar"+                Flat col ->+                    Scalar . finalize <$> foldlColumnWith @b step seed col+                Group gs ->+                    Flat . fromVector+                        <$> V.mapM (fmap finalize . foldlColumnWith @b step seed) gs++-- Aggregation: FoldAgg without seed (fold1) ------------------------------++eval ctx expr@(Agg (FoldAgg _ Nothing (f :: a -> b -> a)) inner) =+    addContext expr $+        case testEquality (typeRep @a) (typeRep @b) of+            Nothing ->+                Left $+                    InternalException+                        "Type mismatch in seedless fold: \+                        \accumulator and element types must match"+            Just Refl -> do+                v <- eval @b ctx inner+                case v of+                    Scalar _ ->+                        Left $+                            InternalException+                                "Cannot apply a fold aggregation to a scalar"+                    Flat col ->+                        Scalar <$> foldl1Column @a f col+                    Group gs ->+                        Flat . fromVector+                            <$> V.mapM (foldl1Column @a f) gs++-------------------------------------------------------------------------------+-- Aggregation helpers+-------------------------------------------------------------------------------++{- | Apply a 'CollectAgg' function to a single column, extracting the+appropriate vector type and applying the aggregation function.+-}+applyCollect ::+    forall v b a.+    (VG.Vector v b, Typeable v, Columnable b, Columnable a) =>+    (v b -> a) -> Column -> Either DataFrameException a+applyCollect f col = f <$> toVector @b @v col++-------------------------------------------------------------------------------+-- Backward-compatible wrappers+-------------------------------------------------------------------------------++{- | Result of interpreting an expression in a grouped context.+Retained for backward compatibility with 'aggregate' and friends.+-}+data AggregationResult a+    = UnAggregated Column+    | Aggregated (TypedColumn a)++{- | Interpret an expression against a flat 'DataFrame', producing a+typed column.  This is the original top-level entry point; internally+it calls 'eval' and materialises the result.++NOTE: unlike the old implementation, 'Lit' values are no longer+eagerly broadcast.  The broadcast happens here, at the boundary,+via 'materialize'.+-}+interpret ::+    forall a.+    (Columnable a) =>+    DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)+interpret df expr = do+    v <- eval (FlatCtx df) expr+    pure $ TColumn $ materialize @a (fst (dataframeDimensions df)) v++{- | Interpret an expression against a 'GroupedDataFrame',+distinguishing aggregated results from bare column references.+Internally calls 'eval'.+-}+interpretAggregation ::+    forall a.+    (Columnable a) =>+    GroupedDataFrame ->+    Expr a ->+    Either DataFrameException (AggregationResult a)+interpretAggregation gdf expr = do+    v <- eval (GroupCtx gdf) expr+    case v of+        Scalar a ->+            Right $+                Aggregated $+                    TColumn $+                        broadcastScalar @a (numGroups gdf) a+        Flat col ->+            Right $ Aggregated $ TColumn col+        Group _ ->+            -- The Column payload is intentionally unused — the only+            -- call-site ('aggregate') immediately throws+            -- 'UnaggregatedException' on this constructor.+            Right $ UnAggregated $ BoxedColumn @T.Text V.empty

@@ -1,16 +1,25 @@+{-# LANGUAGE LambdaCase #-} {-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}  module DataFrame.Internal.Parsing where  import qualified Data.ByteString.Char8 as C import qualified Data.Set as S import qualified Data.Text as T+import qualified Data.Text.IO as TIO +import Control.Applicative (many, (<|>))+import Data.Attoparsec.Text hiding (decimal, double, signed) import Data.ByteString.Lex.Fractional+import Data.Foldable (fold) import Data.Maybe (fromMaybe)-import Data.Text.Read+import Data.Text.Read (decimal, double, signed)+import Data.Time (Day, defaultTimeLocale, parseTimeM) import GHC.Stack (HasCallStack)+import System.IO (Handle, IOMode (..), hIsEOF, hTell, withFile) import Text.Read (readMaybe)+import Prelude hiding (takeWhile)  isNullish :: T.Text -> Bool isNullish =@@ -19,6 +28,13 @@             ["Nothing", "NULL", "", " ", "nan", "null", "N/A", "NaN", "NAN", "NA"]     ) +isNullishBS :: C.ByteString -> Bool+isNullishBS =+    ( `S.member`+        S.fromList+            ["Nothing", "NULL", "", " ", "nan", "null", "N/A", "NaN", "NAN", "NA"]+    )+ isTrueish :: T.Text -> Bool isTrueish t = t `elem` ["True", "true", "TRUE"] @@ -27,7 +43,7 @@  readValue :: (HasCallStack, Read a) => T.Text -> a readValue s = case readMaybe (T.unpack s) of-    Nothing -> error $ "Could not read value: " ++ T.unpack s+    Nothing -> error ("Could not read value: " <> T.unpack s)     Just value -> value  readBool :: (HasCallStack) => T.Text -> Maybe Bool@@ -36,6 +52,15 @@     | isFalseish s = Just False     | otherwise = Nothing +readByteStringBool :: C.ByteString -> Maybe Bool+readByteStringBool s+    | s `elem` ["True", "true", "TRUE"] = Just True+    | s `elem` ["False", "false", "FALSE"] = Just False+    | otherwise = Nothing++readByteStringDate :: String -> C.ByteString -> Maybe Day+readByteStringDate fmt = parseTimeM True defaultTimeLocale fmt . C.unpack+ readInteger :: (HasCallStack) => T.Text -> Maybe Integer readInteger s = case signed decimal (T.strip s) of     Left _ -> Nothing@@ -102,3 +127,106 @@  readWithDefault :: (HasCallStack, Read a) => a -> T.Text -> a readWithDefault v s = fromMaybe v (readMaybe (T.unpack s))++-- ---------------------------------------------------------------------------+-- Attoparsec CSV parser combinators (shared between Lazy.IO.CSV and others)+-- ---------------------------------------------------------------------------++parseSep :: Char -> T.Text -> [T.Text]+parseSep c s = either error id (parseOnly (record c) s)+{-# INLINE parseSep #-}++record :: Char -> Parser [T.Text]+record c =+    field c `sepBy1` char c+        <?> "record"+{-# INLINE record #-}++parseRow :: Char -> Parser [T.Text]+parseRow c = (record c <* lineEnd) <?> "record-new-line"++field :: Char -> Parser T.Text+field c =+    quotedField <|> unquotedField c+        <?> "field"+{-# INLINE field #-}++unquotedTerminators :: Char -> S.Set Char+unquotedTerminators sep = S.fromList [sep, '\n', '\r', '"']++unquotedField :: Char -> Parser T.Text+unquotedField sep =+    takeWhile (not . (`S.member` terminators)) <?> "unquoted field"+  where+    terminators = unquotedTerminators sep+{-# INLINE unquotedField #-}++quotedField :: Parser T.Text+quotedField = char '"' *> contents <* char '"' <?> "quoted field"+  where+    contents = fold <$> many (unquote <|> unescape)+      where+        unquote = takeWhile1 (notInClass "\"\\")+        unescape =+            char '\\' *> do+                T.singleton <$> do+                    char '\\' <|> char '"'+{-# INLINE quotedField #-}++lineEnd :: Parser ()+lineEnd =+    (endOfLine <|> endOfInput)+        <?> "end of line"+{-# INLINE lineEnd #-}++-- | First pass to count rows for exact allocation.+countRows :: Char -> FilePath -> IO Int+countRows c path = withFile path ReadMode $! go 0 ""+  where+    go n input h = do+        isEOF <- hIsEOF h+        if isEOF && input == mempty+            then pure n+            else+                parseWith (TIO.hGetChunk h) (parseRow c) input >>= \case+                    Fail unconsumed ctx er -> do+                        erpos <- hTell h+                        fail $+                            "Failed to parse CSV file around "+                                <> show erpos+                                <> " byte; due: "+                                <> show er+                                <> "; context: "+                                <> show ctx+                                <> " "+                                <> show unconsumed+                    Partial _ -> fail $ "Partial handler is called; n = " <> show n+                    Done (unconsumed :: T.Text) _ ->+                        go (n + 1) unconsumed h+{-# INLINE countRows #-}++-- | Infer the Haskell type name from a text sample.+inferValueType :: T.Text -> T.Text+inferValueType s = case readInt s of+    Just _ -> "Int"+    Nothing -> case readDouble s of+        Just _ -> "Double"+        Nothing -> "Other"+{-# INLINE inferValueType #-}++-- | Read a single CSV row from a handle using the given separator.+readSingleLine :: Char -> T.Text -> Handle -> IO ([T.Text], T.Text)+readSingleLine c unused handle =+    parseWith (TIO.hGetChunk handle) (parseRow c) unused >>= \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 _ -> fail "Partial handler is called"+        Done (unconsumed :: T.Text) (row :: [T.Text]) ->+            return (row, unconsumed)

@@ -17,6 +17,7 @@ import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU +import Control.DeepSeq (NFData (..)) import Control.Exception (throw) import Control.Monad.ST (runST) import Data.Function (on)@@ -56,6 +57,9 @@ instance Show Any where     show :: Any -> String     show (Value a) = T.unpack (showValue a)++instance NFData Any where+    rnf (Value a) = rnf a  showValue :: forall a. (Columnable a) => a -> T.Text showValue v = case testEquality (typeRep @a) (typeRep @T.Text) of

@@ -13,13 +13,14 @@  module DataFrame.Internal.Types where +import Control.DeepSeq (NFData) import Data.Int (Int16, Int32, Int64, Int8) import Data.Kind (Constraint, Type) import Data.Typeable (Typeable) import qualified Data.Vector.Unboxed as VU import Data.Word (Word16, Word32, Word64, Word8) -type Columnable' a = (Typeable a, Show a, Ord a, Eq a, Read a)+type Columnable' a = (Typeable a, Show a, Ord a, Eq a, Read a, NFData a)  {- | A type with column representations used to select the "right" representation when specializing the `toColumn` function.

@@ -0,0 +1,444 @@+{-# 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.Maybe (fromMaybe, isJust)+import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))+import Data.Word (Word16, Word32, Word64, Word8)+import DataFrame.Internal.Column (Column (..))+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 (_ :: 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 (BoxedColumn _) = tagText+columnTypeTag (OptionalColumn (_ :: V.Vector (Maybe a))) =+    case testEquality (typeRep @a) (typeRep @Int) of+        Just Refl -> tagMaybeInt+        Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+            Just Refl -> tagMaybeDouble+            Nothing -> tagMaybeText++buildColumnData :: Int -> Column -> BSB.Builder+buildColumnData _ (UnboxedColumn (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 _ (BoxedColumn (v :: V.Vector a)) =+    case testEquality (typeRep @a) (typeRep @T.Text) of+        Just Refl -> buildTextVector v+        Nothing -> error "spillToDisk: unsupported BoxedColumn element type"+buildColumnData _ (OptionalColumn (v :: V.Vector (Maybe a))) =+    case testEquality (typeRep @a) (typeRep @Int) of+        Just Refl ->+            buildNullBitmap (V.map isJust v)+                <> buildIntVector (VU.convert (V.map (fromMaybe 0) v))+        Nothing ->+            case testEquality (typeRep @a) (typeRep @Double) of+                Just Refl ->+                    buildNullBitmap (V.map isJust v)+                        <> buildDoubleVector (VU.convert (V.map (fromMaybe 0.0) v))+                Nothing ->+                    let showText x = case testEquality (typeRep @a) (typeRep @T.Text) of+                            Just Refl -> x+                            Nothing -> T.pack (show x)+                        texts = V.map (maybe T.empty showText) v+                     in buildNullBitmap (V.map isJust v) <> 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 v)+    | tag == tagDouble = do+        (off', v) <- readDoubleColumn bs off nrows+        return (off', UnboxedColumn v)+    | tag == tagText = do+        (off', v) <- readTextColumn bs off nrows+        return (off', BoxedColumn v)+    | tag == tagMaybeInt = do+        (off1, bitmap) <- readNullBitmap bs off nrows+        (off2, v) <- readIntColumn bs off1 nrows+        let maybes =+                V.fromList+                    (zipWith (\valid x -> if valid then Just x else Nothing) bitmap (VU.toList v)) ::+                    V.Vector (Maybe Int)+        return (off2, OptionalColumn maybes)+    | tag == tagMaybeDouble = do+        (off1, bitmap) <- readNullBitmap bs off nrows+        (off2, v) <- readDoubleColumn bs off1 nrows+        let maybes =+                V.fromList+                    (zipWith (\valid x -> if valid then Just x else Nothing) bitmap (VU.toList v)) ::+                    V.Vector (Maybe Double)+        return (off2, OptionalColumn maybes)+    | tag == tagMaybeText = do+        (off1, bitmap) <- readNullBitmap bs off nrows+        (off2, v) <- readTextColumn bs off1 nrows+        let maybes =+                V.fromList+                    (zipWith (\valid x -> if valid then Just x else Nothing) bitmap (V.toList v)) ::+                    V.Vector (Maybe T.Text)+        return (off2, OptionalColumn maybes)+    | 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 =+        let b0 = fromIntegral (BSU.unsafeIndex bs off) :: Word32+            b1 = fromIntegral (BSU.unsafeIndex bs (off + 1)) :: Word32+            b2 = fromIntegral (BSU.unsafeIndex bs (off + 2)) :: Word32+            b3 = fromIntegral (BSU.unsafeIndex bs (off + 3)) :: Word32+         in Right+                (off + 4, b0 .|. (b1 `shiftL` 8) .|. (b2 `shiftL` 16) .|. (b3 `shiftL` 24))++readWord64LE :: BS.ByteString -> Int -> ParseResult Word64+readWord64LE bs off+    | off + 8 > BS.length bs = Left "unexpected end of input"+    | otherwise =+        let b0 = fromIntegral (BSU.unsafeIndex bs off) :: Word64+            b1 = fromIntegral (BSU.unsafeIndex bs (off + 1)) :: Word64+            b2 = fromIntegral (BSU.unsafeIndex bs (off + 2)) :: Word64+            b3 = fromIntegral (BSU.unsafeIndex bs (off + 3)) :: Word64+            b4 = fromIntegral (BSU.unsafeIndex bs (off + 4)) :: Word64+            b5 = fromIntegral (BSU.unsafeIndex bs (off + 5)) :: Word64+            b6 = fromIntegral (BSU.unsafeIndex bs (off + 6)) :: Word64+            b7 = fromIntegral (BSU.unsafeIndex bs (off + 7)) :: Word64+         in Right+                ( off + 8+                , b0+                    .|. (b1 `shiftL` 8)+                    .|. (b2 `shiftL` 16)+                    .|. (b3 `shiftL` 24)+                    .|. (b4 `shiftL` 32)+                    .|. (b5 `shiftL` 40)+                    .|. (b6 `shiftL` 48)+                    .|. (b7 `shiftL` 56)+                )++-- | 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+                  b0 = fromIntegral (BSU.unsafeIndex bs i) :: Word32+                  b1 = fromIntegral (BSU.unsafeIndex bs (i + 1)) :: Word32+                  b2 = fromIntegral (BSU.unsafeIndex bs (i + 2)) :: Word32+                  b3 = fromIntegral (BSU.unsafeIndex bs (i + 3)) :: Word32+               in b0 .|. (b1 `shiftL` 8) .|. (b2 `shiftL` 16) .|. (b3 `shiftL` 24)+            | 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

@@ -1,3 +1,4 @@+{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE LambdaCase #-}@@ -8,25 +9,23 @@  module DataFrame.Lazy.IO.CSV where -import qualified Data.List as L+import qualified Data.ByteString as BS import qualified Data.Map as M-import qualified Data.Set as S+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 as VU import qualified Data.Vector.Unboxed.Mutable as VUM -import Control.Applicative (many, (<|>)) import Control.Monad (forM_, unless, when, zipWithM_)-import Data.Attoparsec.Text-import Data.Char-import Data.Foldable (fold)-import Data.Function (on)+import Data.Attoparsec.Text (IResult (..), parseWith)+import Data.Char (intToDigit) import Data.IORef-import Data.Maybe+import Data.Maybe (fromMaybe, isJust) import Data.Type.Equality (TestEquality (testEquality))+import Data.Word (Word8) import DataFrame.Internal.Column (     Column (..),     MutableColumn (..),@@ -36,9 +35,10 @@  ) import DataFrame.Internal.DataFrame (DataFrame (..)) import DataFrame.Internal.Parsing+import DataFrame.Internal.Schema (Schema, SchemaType (..), elements) import System.IO import Type.Reflection-import Prelude hiding (concat, takeWhile)+import Prelude hiding (takeWhile)  -- | Record for CSV read options. data ReadOptions = ReadOptions@@ -94,11 +94,11 @@             Nothing -> (0, totalRows)             Just (start, len) -> (start, min len (totalRows - rowsRead opts))     withFile path ReadMode $ \handle -> do-        firstRow <- map T.strip . parseSep c <$> TIO.hGetLine handle+        firstRow <- fmap T.strip . parseSep c <$> TIO.hGetLine handle         let columnNames =                 if hasHeader opts-                    then map (T.filter (/= '\"')) firstRow-                    else map (T.singleton . intToDigit) [0 .. (length firstRow - 1)]+                    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 @@ -110,13 +110,9 @@         let numColumns = length columnNames         let numRows = len         -- Use this row to infer the types of the rest of the column.-        -- TODO: this isn't robust but in so far as this is a guess anyway-        -- it's probably fine. But we should probably sample n rows and pick-        -- the most likely type from the sample.         (dataRow, remainder) <- readSingleLine c (leftOver opts) handle          -- This array will track the indices of all null values for each column.-        -- If any exist then the column will be an optional type.         nullIndices <- VM.unsafeNew numColumns         VM.set nullIndices []         mutableCols <- VM.unsafeNew numColumns@@ -136,7 +132,7 @@                 { columns = cols                 , columnIndices = M.fromList (zip columnNames [0 ..])                 , dataframeDimensions = (maybe 0 columnLength (cols V.!? 0), V.length cols)-                , derivingExpressions = M.empty -- TODO create base expressions.+                , derivingExpressions = M.empty                 }             , (pos, unconsumed, r + 1)             )@@ -161,35 +157,6 @@         VM.unsafeWrite mCol i col {-# INLINE getInitialDataVectors #-} -inferValueType :: T.Text -> T.Text-inferValueType s =-    let-        example = s-     in-        case readInt example of-            Just _ -> "Int"-            Nothing -> case readDouble example of-                Just _ -> "Double"-                Nothing -> "Other"-{-# INLINE inferValueType #-}--readSingleLine :: Char -> T.Text -> Handle -> IO ([T.Text], T.Text)-readSingleLine c unused handle =-    parseWith (TIO.hGetChunk handle) (parseRow c) unused >>= \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 c -> do-            fail "Partial handler is called"-        Done (unconsumed :: T.Text) (row :: [T.Text]) -> do-            return (row, unconsumed)- -- | Reads rows from the handle and stores values in mutable vectors. fillColumns ::     Int ->@@ -216,7 +183,7 @@                             <> show er                             <> "; context: "                             <> show ctx-                Partial c -> do+                Partial _ -> do                     fail "Partial handler is called"                 Done (unconsumed :: T.Text) (row :: [T.Text]) -> do                     writeIORef input unconsumed@@ -254,142 +221,232 @@     freezeColumn' (nulls V.! colIndex) col {-# INLINE freezeColumn #-} -parseSep :: Char -> T.Text -> [T.Text]-parseSep c s = either error id (parseOnly (record c) s)-{-# INLINE parseSep #-}+-- ---------------------------------------------------------------------------+-- Streaming scan API+-- --------------------------------------------------------------------------- -record :: Char -> Parser [T.Text]-record c =-    field c `sepBy1` char c-        <?> "record"-{-# INLINE record #-}+{- | 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. -parseRow :: Char -> Parser [T.Text]-parseRow c = (record c <* lineEnd) <?> "record-new-line"+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) -field :: Char -> Parser T.Text-field c =-    quotedField <|> unquotedField c-        <?> "field"-{-# INLINE field #-}+{- | 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. -unquotedTerminators :: Char -> S.Set Char-unquotedTerminators sep = S.fromList [sep, '\n', '\r', '"']+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+                    ) -unquotedField :: Char -> Parser T.Text-unquotedField sep =-    takeWhile (not . (`S.member` terminators)) <?> "unquoted field"-  where-    terminators = unquotedTerminators sep-{-# INLINE unquotedField #-}+{- | 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 if isNullish val+                    then VM.unsafeWrite col i T.empty >> return (Left val)+                    else VM.unsafeWrite col i val >> return (Right True)+        Nothing -> return (Left (TextEncoding.decodeUtf8Lenient bs))+writeColumnBs i bs (MOptionalColumn (col :: VM.IOVector (Maybe a))) =+    case testEquality (typeRep @a) (typeRep @T.Text) of+        Just Refl ->+            let val = TextEncoding.decodeUtf8Lenient bs+             in if isNullish val+                    then VM.unsafeWrite col i Nothing >> return (Left val)+                    else VM.unsafeWrite col i (Just 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 #-} -quotedField :: Parser T.Text-quotedField = char '"' *> contents <* char '"' <?> "quoted field"+{- | 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-    contents = fold <$> many (unquote <|> unescape)-      where-        unquote = takeWhile1 (notInClass "\"\\")-        unescape =-            char '\\' *> do-                T.singleton <$> do-                    char '\\' <|> char '"'-{-# INLINE quotedField #-}+    -- 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) -lineEnd :: Parser ()-lineEnd =-    (endOfLine <|> endOfInput)-        <?> "end of line"-{-# INLINE lineEnd #-}+    -- 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) --- | First pass to count rows for exact allocation-countRows :: Char -> FilePath -> IO Int-countRows c path = withFile path ReadMode $! go 0 ""-  where-    go n input h = do-        isEOF <- hIsEOF h-        if isEOF && input == mempty-            then pure n-            else-                parseWith (TIO.hGetChunk h) (parseRow c) input >>= \case-                    Fail unconsumed ctx er -> do-                        erpos <- hTell h-                        fail $-                            "Failed to parse CSV file around "-                                <> show erpos-                                <> " byte; due: "-                                <> show er-                                <> "; context: "-                                <> show ctx-                                <> " "-                                <> show unconsumed-                    Partial c -> do-                        fail $ "Partial handler is called; n = " <> show n-                    Done (unconsumed :: T.Text) _ ->-                        go (n + 1) unconsumed h-{-# INLINE countRows #-}+    extract s e =+        let field = BS.take (e - s) (BS.drop s bs)+         in if BS.length field >= 2+                && BS.head field == quoteChar+                && BS.last field == quoteChar+                then BS.init (BS.tail field)+                else field+{-# INLINE getNthFieldBs #-} -writeCsv :: FilePath -> DataFrame -> IO ()-writeCsv = writeSeparated ','+-- | 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)) -writeSeparated ::-    -- | Separator-    Char ->-    -- | Path to write to-    FilePath ->-    DataFrame ->-    IO ()-writeSeparated c filepath df = withFile filepath WriteMode $ \handle -> do-    let (rows, _) = dataframeDimensions df-    let headers = map fst (L.sortBy (compare `on` snd) (M.toList (columnIndices df)))-    TIO.hPutStrLn handle (T.intercalate ", " headers)-    forM_ [0 .. (rows - 1)] $ \i -> do-        let row = getRowAsText df i-        TIO.hPutStrLn handle (T.intercalate ", " row)+-- | 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)+sliceCol n (MOptionalColumn col) = MOptionalColumn (VM.take n col) -getRowAsText :: DataFrame -> Int -> [T.Text]-getRowAsText df i = V.ifoldr go [] (columns df)+{- | 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-    indexMap = M.fromList (map (\(a, b) -> (b, a)) $ M.toList (columnIndices df))-    go k (BoxedColumn (c :: V.Vector a)) acc = case c V.!? i of-        Just e -> textRep : acc-          where-            textRep = case testEquality (typeRep @a) (typeRep @T.Text) of-                Just Refl -> e-                Nothing -> case typeRep @a of-                    App t1 t2 -> case eqTypeRep t1 (typeRep @Maybe) of-                        Just HRefl -> case testEquality t2 (typeRep @T.Text) of-                            Just Refl -> fromMaybe "null" e-                            Nothing -> (fromOptional . (T.pack . show)) e-                              where-                                fromOptional s-                                    | T.isPrefixOf "Just " s = T.drop (T.length "Just ") s-                                    | otherwise = "null"-                        Nothing -> (T.pack . show) e-                    _ -> (T.pack . show) e-        Nothing ->-            error $-                "Column "-                    ++ T.unpack (indexMap M.! k)-                    ++ " has less items than "-                    ++ "the other columns at index "-                    ++ show i-    go k (UnboxedColumn c) acc = case c VU.!? i of-        Just e -> T.pack (show e) : acc-        Nothing ->-            error $-                "Column "-                    ++ T.unpack (indexMap M.! k)-                    ++ " has less items than "-                    ++ "the other columns at index "-                    ++ show i-    go k (OptionalColumn (c :: V.Vector (Maybe a))) acc = case c V.!? i of-        Just e -> case testEquality (typeRep @a) (typeRep @T.Text) of-            Just Refl -> fromMaybe T.empty e : acc-            Nothing -> maybe T.empty (T.pack . show) e : acc-        Nothing ->-            error $-                "Column "-                    ++ T.unpack (indexMap M.! k)-                    ++ " has less items than "-                    ++ "the other columns at index "-                    ++ show i+    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 #-}

@@ -1,107 +1,138 @@ {-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE ExistentialQuantification #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-}-{-# LANGUAGE InstanceSigs #-} {-# LANGUAGE NumericUnderscores #-} {-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}  module DataFrame.Lazy.Internal.DataFrame where -import Control.Monad (foldM)-import qualified Data.List as L import qualified Data.Text as T import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E-import qualified DataFrame.Lazy.IO.CSV as D-import DataFrame.Operations.Merge ()-import qualified DataFrame.Operations.Subset as D-import qualified DataFrame.Operations.Transformations as D--data LazyOperation where-    Derive :: (C.Columnable a) => T.Text -> E.Expr a -> LazyOperation-    Select :: [T.Text] -> LazyOperation-    Filter :: E.Expr Bool -> LazyOperation--instance Show LazyOperation where-    show :: LazyOperation -> String-    show (Derive name expr) = T.unpack name ++ " := " ++ show expr-    show (Select columns) = "select(" ++ show columns ++ ")"-    show (Filter expr) = "filter(" ++ show expr ++ ")"+import DataFrame.Internal.Schema (Schema)+import DataFrame.Lazy.Internal.Executor (+    ExecutorConfig (..),+    defaultExecutorConfig,+    execute,+ )+import DataFrame.Lazy.Internal.LogicalPlan (+    DataSource (..),+    LogicalPlan (..),+    SortOrder (..),+ )+import qualified DataFrame.Lazy.Internal.Optimizer as Opt+import DataFrame.Operations.Join (JoinType) -data InputType = ICSV deriving (Show)+{- | 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-    { inputPath :: FilePath-    , inputType :: InputType-    , operations :: [LazyOperation]+    { plan :: LogicalPlan     , batchSize :: Int     }-    deriving (Show) -eval :: LazyOperation -> D.DataFrame -> D.DataFrame-eval (Derive name expr) = D.derive name expr-eval (Select columns) = D.select columns-eval (Filter expr) = D.filterWhere expr+instance Show LazyDataFrame where+    show ldf =+        "LazyDataFrame { batchSize = "+            <> (show (batchSize ldf) <> (", plan = " <> (show (plan ldf) <> " }"))) -runDataFrame :: forall a. (C.Columnable a) => LazyDataFrame -> IO D.DataFrame-runDataFrame df = do-    let path = inputPath df-    totalRows <- D.countRows ',' path-    let batches = batchRanges totalRows (batchSize df)-    (df', _) <--        foldM-            ( \(accDf, (pos, unused, r)) (start, end) -> do-                mapM_-                    putStr-                    [ "Scanning: "-                    , show start-                    , " to "-                    , show end-                    , " rows out of "-                    , show totalRows-                    , "\n"-                    ]+-- ---------------------------------------------------------------------------+-- Entry point+-- --------------------------------------------------------------------------- -                (sdf, (pos', unconsumed, rowsRead)) <--                    D.readSeparated-                        ','-                        ( D.defaultOptions-                            { D.rowRange = Just (start, batchSize df)-                            , D.totalRows = Just totalRows-                            , D.seekPos = pos-                            , D.rowsRead = r-                            , D.leftOver = unused-                            }-                        )-                        path-                let rdf = L.foldl' (flip eval) sdf (operations df)-                return (accDf <> rdf, (Just pos', unconsumed, rowsRead + r))-            )-            (D.empty, (Nothing, "", 0))-            batches-    return df'+{- | Execute the lazy query: optimise the logical plan, then stream-execute+the resulting physical plan, returning a fully-materialised 'D.DataFrame'.+-}+runDataFrame :: LazyDataFrame -> IO D.DataFrame+runDataFrame ldf = do+    let physPlan = Opt.optimize (batchSize ldf) (plan ldf)+    execute physPlan defaultExecutorConfig{defaultBatchSize = batchSize ldf} -batchRanges :: Int -> Int -> [(Int, Int)]-batchRanges n inc = go n [0, inc .. n]-  where-    go _ [] = []-    go n [x] = [(x, n)]-    go n (f : s : rest) = (f, s) : go n (s : rest)+-- ---------------------------------------------------------------------------+-- Builders that construct the logical plan tree+-- --------------------------------------------------------------------------- -scanCsv :: T.Text -> LazyDataFrame-scanCsv path = LazyDataFrame (T.unpack path) ICSV [] 512_000+-- | 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} -addOperation :: LazyOperation -> LazyDataFrame -> LazyDataFrame-addOperation op df = df{operations = operations df ++ [op]}+-- | Scan a CSV file with the default comma separator.+scanCsv :: Schema -> T.Text -> LazyDataFrame+scanCsv schema path =+    LazyDataFrame+        { plan = Scan (CsvSource (T.unpack path) ',') schema+        , batchSize = 1_000_000+        } +-- | Scan a character-separated file.+scanSeparated :: Char -> Schema -> T.Text -> LazyDataFrame+scanSeparated sep schema path =+    LazyDataFrame+        { plan = Scan (CsvSource (T.unpack path) sep) 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 = addOperation (Derive name expr)+derive name expr ldf =+    ldf{plan = Derive name (E.UExpr expr) (plan ldf)} -select :: (C.Columnable a) => [T.Text] -> LazyDataFrame -> LazyDataFrame-select columns = addOperation (Select columns)+-- | Retain only the listed columns.+select :: [T.Text] -> LazyDataFrame -> LazyDataFrame+select cols ldf = ldf{plan = Project cols (plan ldf)} -filter :: (C.Columnable a) => E.Expr Bool -> LazyDataFrame -> LazyDataFrame-filter cond = addOperation (Filter cond)+-- | 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.+limit :: Int -> LazyDataFrame -> LazyDataFrame+limit n ldf = ldf{plan = Limit n (plan ldf)}

@@ -0,0 +1,543 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE NumericUnderscores #-}+{-# 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 (+    ExecutorConfig (..),+    defaultExecutorConfig,+    execute,+    foldBatches,+) where++import Control.Concurrent (forkIO)+import Control.Concurrent.STM (atomically)+import Control.Concurrent.STM.TBQueue (newTBQueueIO, readTBQueue, writeTBQueue)+import Control.DeepSeq (force)+import Control.Exception (evaluate)+import Control.Monad (filterM, when)+import qualified Data.ByteString as BS+import Data.IORef+import qualified Data.Map as M+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 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 qualified DataFrame.Lazy.IO.CSV as LCSV+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)+import System.FilePath ((</>))+import System.FilePath.Glob (glob)+import System.IO (hClose)+import Type.Reflection (typeRep)++-- ---------------------------------------------------------------------------+-- Configuration+-- ---------------------------------------------------------------------------++data ExecutorConfig = ExecutorConfig+    { memoryBudgetBytes :: !Int+    -- ^ Per-node spill threshold (currently informational; not enforced yet).+    , spillDirectory :: FilePath+    , defaultBatchSize :: !Int+    }++defaultExecutorConfig :: ExecutorConfig+defaultExecutorConfig =+    ExecutorConfig+        { memoryBudgetBytes = 512 * 1_048_576 -- 512 MiB+        , spillDirectory = "/tmp"+        , defaultBatchSize = 1_000_000+        }++-- ---------------------------------------------------------------------------+-- 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 -> ExecutorConfig -> IO D.DataFrame+execute plan cfg = buildStream plan cfg >>= 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 -> ExecutorConfig -> IO b+foldBatches f seed plan cfg = do+    stream <- buildStream plan cfg+    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 -> ExecutorConfig -> IO Stream+-- Scan -----------------------------------------------------------------------+buildStream (PhysicalScan (CsvSource path sep) cfg) _ =+    executeCsvScan path sep cfg+buildStream (PhysicalScan (ParquetSource path) cfg) _ =+    executeParquetScan path cfg+buildStream (PhysicalSpill child path) execCfg = do+    df <- execute child execCfg+    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) execCfg = do+    childStream <- buildStream child execCfg+    return . Stream $+        ( do+            mb <- pullBatch childStream+            return $ fmap (Sub.filterWhere p) mb+        )+-- Project --------------------------------------------------------------------+buildStream (PhysicalProject cols child) execCfg = do+    childStream <- buildStream child execCfg+    return . Stream $+        ( do+            mb <- pullBatch childStream+            return $ fmap (Sub.select cols) mb+        )+-- Derive ---------------------------------------------------------------------+buildStream (PhysicalDerive name uexpr child) execCfg = do+    childStream <- buildStream child execCfg+    return . Stream $+        ( do+            mb <- pullBatch childStream+            return $ fmap (Trans.deriveMany [(name, uexpr)]) mb+        )+-- Limit ----------------------------------------------------------------------+buildStream (PhysicalLimit n child) execCfg = do+    childStream <- buildStream child execCfg+    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) execCfg = do+    df <- execute child execCfg+    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) execCfg = do+    childStream <- buildStream child execCfg+    if all (isStreamableAgg . snd) aggs+        then do+            -- Streaming partial aggregation: O(|groups|) memory+            let (partialAggs, mergeAggs, finalizer) = buildAggPlan aggs+            accRef <- newIORef (Nothing :: Maybe D.DataFrame)+            let loop = do+                    mb <- pullBatch childStream+                    case mb of+                        Nothing -> return ()+                        Just batch -> do+                            -- Force to NF so the batch DataFrame can be GC'd immediately.+                            -- evaluate . force breaks the thunk chain that would otherwise+                            -- keep every batch (~60 MB each) alive until the end = OOM.+                            !partial <-+                                evaluate . force $ Agg.aggregate partialAggs (Agg.groupBy keys batch)+                            mAcc <- readIORef accRef+                            !newAcc <- case mAcc of+                                Nothing -> return partial+                                Just acc ->+                                    evaluate . force $+                                        Agg.aggregate mergeAggs $+                                            Agg.groupBy keys (acc <> partial)+                            writeIORef accRef (Just newAcc)+                            loop+            loop+            mFinal <- fmap (fmap finalizer) (readIORef accRef)+            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 df) execCfg = do+    let bs = defaultBatchSize execCfg+        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) execCfg =+    case jt of+        Join.INNER -> streamingHashJoin assembleInnerBatch+        Join.LEFT -> streamingHashJoin assembleLeftBatch+        _ -> do+            -- Blocking fallback for RIGHT / FULL_OUTER+            leftDf <- execute leftPlan execCfg+            rightDf <- execute rightPlan execCfg+            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 execCfg+        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 execCfg+        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) execCfg = do+    leftDf <- execute leftPlan execCfg+    rightDf <- execute rightPlan execCfg+    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).+-}+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 Core.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 = 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++-- ---------------------------------------------------------------------------+-- CSV scan implementation+-- ---------------------------------------------------------------------------++{- | CSV scan with pipeline parallelism: a dedicated reader thread fills a+bounded queue while the caller's thread applies pushdown predicates and+delivers batches to the rest of the pipeline.  The queue depth of 8 keeps+at most eight raw batches in flight, bounding memory while hiding I/O latency.+-}+executeCsvScan :: FilePath -> Char -> ScanConfig -> IO Stream+executeCsvScan path sep cfg = do+    (handle, colSpec) <- LCSV.openCsvStream sep (scanSchema cfg) path+    -- Queue carries raw batches; Nothing is the end-of-stream sentinel.+    -- Depth 2: each batch holds ~60 MB (1M Text + Double columns); 8 would be ~480 MB.+    queue <- newTBQueueIO 2+    _ <- forkIO $ do+        let loop lo = do+                result <- LCSV.readBatch sep colSpec (scanBatchSize cfg) lo handle+                case result of+                    Nothing ->+                        hClose handle >> atomically (writeTBQueue queue Nothing)+                    Just (df, lo') ->+                        atomically (writeTBQueue queue (Just df)) >> loop lo'+        loop BS.empty+    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')+        )++-- ---------------------------------------------------------------------------+-- 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)

@@ -0,0 +1,44 @@+{-# LANGUAGE GADTs #-}++module DataFrame.Lazy.Internal.LogicalPlan 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.Operations.Join (JoinType)++-- | Data source for a scan node.+data DataSource+    = -- | path, separator+      CsvSource FilePath Char+    | ParquetSource FilePath+    deriving (Show)++-- | 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)

@@ -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) schema)) =+    PhysicalScan+        (CsvSource path sep)+        (ScanConfig batchSz sep schema (Just p))+toPhysical batchSz (Scan (CsvSource path sep) schema) =+    PhysicalScan+        (CsvSource path sep)+        (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 _ (SourceDF df) = PhysicalSourceDF df

@@ -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.+      PhysicalSourceDF D.DataFrame+    deriving (Show)

@@ -59,12 +59,17 @@             names             VU.empty             (VU.fromList [0])+            VU.empty     | otherwise =-        Grouped-            df-            names-            (VU.map fst valueIndices)-            (changingPoints valueIndices)+        let !vis = VU.map fst valueIndices+            !os = changingPoints valueIndices+            !n = fst (dimensions df)+         in Grouped+                df+                names+                vis+                os+                (buildRowToGroup n vis os)   where     indicesToGroup = M.elems $ M.filterWithKey (\k _ -> k `elem` names) (columnIndices df)     doubleToInt :: Double -> Int@@ -159,6 +164,21 @@         VA.sortBy numPasses bucketSize radixFunc mv         VU.unsafeFreeze mv +{- | Build the rowToGroup lookup vector from valueIndices and offsets.+rowToGroup[i] = k means row i belongs to group k.+-}+buildRowToGroup :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int+buildRowToGroup n vis os = runST $ do+    rtg <- VUM.new n+    let nGroups = VU.length os - 1+    forM_ [0 .. nGroups - 1] $ \k ->+        let s = VU.unsafeIndex os k+            e = VU.unsafeIndex os (k + 1)+         in forM_ [s .. e - 1] $ \i ->+                VUM.unsafeWrite rtg (VU.unsafeIndex vis i) k+    VU.unsafeFreeze rtg+{-# NOINLINE buildRowToGroup #-}+ changingPoints :: VU.Vector (Int, Int) -> VU.Vector Int changingPoints vs =     VU.reverse@@ -260,7 +280,7 @@ All ungrouped columns will be dropped. -} aggregate :: [NamedExpr] -> GroupedDataFrame -> DataFrame-aggregate aggs gdf@(Grouped df groupingColumns valueIndices offsets) =+aggregate aggs gdf@(Grouped df groupingColumns valueIndices offsets _rowToGroup) =     let         df' =             selectIndices@@ -289,4 +309,4 @@ distinct :: DataFrame -> DataFrame distinct df = selectIndices (VU.map (indices VU.!) (VU.init os)) df   where-    (Grouped _ _ indices os) = groupBy (columnNames df) df+    (Grouped _ _ indices os _rtg) = groupBy (columnNames df) df

@@ -10,7 +10,7 @@ module DataFrame.Operations.Join where  import Control.Applicative ((<|>))-import Control.Monad (forM_)+import Control.Monad (forM_, when) import Control.Monad.ST (ST, runST) import qualified Data.HashMap.Strict as HM import Data.List (foldl')@@ -36,6 +36,7 @@     | LEFT     | RIGHT     | FULL_OUTER+    deriving (Show)  -- | Join two dataframes using SQL join semantics. join ::@@ -72,12 +73,7 @@ buildCompactIndex :: VU.Vector Int -> CompactIndex buildCompactIndex hashes =     let n = VU.length hashes-        sorted = runST $ do-            mv <- VU.thaw (VU.zip hashes (VU.enumFromN 0 n))-            VA.sortBy (\(h1, _) (h2, _) -> compare h1 h2) mv-            VU.unsafeFreeze mv-        sortedHashes = VU.map fst sorted-        sortedIndices = VU.map snd sorted+        (sortedHashes, sortedIndices) = sortWithIndices hashes         !offsets = buildOffsets sortedHashes n 0 HM.empty      in CompactIndex sortedIndices offsets   where@@ -102,15 +98,19 @@     | otherwise = j {-# INLINE findGroupEnd #-} --- | Sort a hash vector, returning sorted hashes and corresponding original indices.+{- | Sort a hash vector, returning sorted hashes and corresponding original indices.+Sorts an index array using hash values as the comparison key, avoiding the+intermediate pair vector used by the naive zip-then-sort approach.+-} sortWithIndices :: VU.Vector Int -> (VU.Vector Int, VU.Vector Int)-sortWithIndices hashes =+sortWithIndices hashes = runST $ do     let n = VU.length hashes-        sorted = runST $ do-            mv <- VU.thaw (VU.zip hashes (VU.enumFromN 0 n))-            VA.sortBy (\(h1, _) (h2, _) -> compare h1 h2) mv-            VU.unsafeFreeze mv-     in (VU.map fst sorted, VU.map snd sorted)+    mv <- VU.thaw (VU.enumFromN 0 n)+    VA.sortBy+        (\i j -> compare (hashes `VU.unsafeIndex` i) (hashes `VU.unsafeIndex` j))+        mv+    sortedIdxs <- VU.unsafeFreeze mv+    return (VU.unsafeBackpermute hashes sortedIdxs, sortedIdxs)  -- | Write the cross product of two index ranges into mutable vectors. fillCrossProduct ::@@ -170,126 +170,284 @@ @ -} innerJoin :: [T.Text] -> DataFrame -> DataFrame -> DataFrame-innerJoin cs right left =-    let csSet = S.fromList cs-        leftRows = fst (D.dimensions left)-        rightRows = fst (D.dimensions right)+innerJoin cs left right+    | D.null right || D.null left = D.empty+    | otherwise =+        let+            csSet = S.fromList cs+            leftRows = fst (D.dimensions left)+            rightRows = fst (D.dimensions right) -        leftKeyIdxs = keyColIndices csSet left-        rightKeyIdxs = keyColIndices csSet right-        leftHashes = D.computeRowHashes leftKeyIdxs left-        rightHashes = D.computeRowHashes rightKeyIdxs right+            leftKeyIdxs = keyColIndices csSet left+            rightKeyIdxs = keyColIndices csSet right+            leftHashes = D.computeRowHashes leftKeyIdxs left+            rightHashes = D.computeRowHashes rightKeyIdxs right -        buildRows = min leftRows rightRows-        (leftIxs, rightIxs)-            | buildRows > joinStrategyThreshold =-                sortMergeInnerKernel leftHashes rightHashes-            | rightRows <= leftRows =-                -- Build on right (smaller or equal), probe with left-                hashInnerKernel leftHashes rightHashes-            | otherwise =-                -- Build on left (smaller), probe with right, swap result-                let (!rIxs, !lIxs) = hashInnerKernel rightHashes leftHashes-                 in (lIxs, rIxs)-     in assembleInner csSet left right leftIxs rightIxs+            buildRows = min leftRows rightRows+            (leftIxs, rightIxs)+                | buildRows > joinStrategyThreshold =+                    sortMergeInnerKernel leftHashes rightHashes+                | rightRows <= leftRows =+                    -- Build on right (smaller or equal), probe with left+                    hashInnerKernel leftHashes rightHashes+                | otherwise =+                    -- Build on left (smaller), probe with right, swap result+                    let (!rIxs, !lIxs) = hashInnerKernel rightHashes leftHashes+                     in (lIxs, rIxs)+         in+            assembleInner csSet left right leftIxs rightIxs +-- | Compute hashes for the given key column names in a DataFrame.+buildHashColumn :: [T.Text] -> DataFrame -> VU.Vector Int+buildHashColumn keys df =+    let csSet = S.fromList keys+        keyIdxs = keyColIndices csSet df+     in D.computeRowHashes keyIdxs df++{- | Probe one batch of rows against a pre-built 'CompactIndex'.+Returns @(probeExpandedIxs, buildExpandedIxs)@.+Unlike 'hashInnerKernel', does not build the index (it is pre-built once)+and has no cross-product row guard — the caller controls probe batch size.+-}+hashProbeKernel ::+    -- | Built once from the full right\/build side.+    CompactIndex ->+    -- | Probe hashes (one batch).+    VU.Vector Int ->+    (VU.Vector Int, VU.Vector Int)+hashProbeKernel ci probeHashes =+    let ciIxs = ciSortedIndices ci+        ciOff = ciOffsets ci+        (pFrozen, bFrozen) = runST $ do+            let !probeN = VU.length probeHashes+                initCap = max 1 (min probeN 1_000_000)++            initPv <- VUM.unsafeNew initCap+            initBv <- VUM.unsafeNew initCap+            pvRef <- newSTRef initPv+            bvRef <- newSTRef initBv+            capRef <- newSTRef initCap+            posRef <- newSTRef (0 :: Int)++            let ensureCapacity needed = do+                    cap <- readSTRef capRef+                    when (needed > cap) $ do+                        let newCap = max needed (cap * 2)+                            delta = newCap - cap+                        pv <- readSTRef pvRef+                        bv <- readSTRef bvRef+                        newPv <- VUM.unsafeGrow pv delta+                        newBv <- VUM.unsafeGrow bv delta+                        writeSTRef pvRef newPv+                        writeSTRef bvRef newBv+                        writeSTRef capRef newCap++                go !i+                    | i >= probeN = return ()+                    | otherwise = do+                        let !h = probeHashes `VU.unsafeIndex` i+                        case HM.lookup h ciOff of+                            Nothing -> go (i + 1)+                            Just (!start, !len) -> do+                                !p <- readSTRef posRef+                                ensureCapacity (p + len)+                                pv <- readSTRef pvRef+                                bv <- readSTRef bvRef+                                fillBuild i start len p 0 pv bv+                                writeSTRef posRef (p + len)+                                go (i + 1)+                fillBuild !probeIdx !start !len !p !j !pv !bv+                    | j >= len = return ()+                    | otherwise = do+                        VUM.unsafeWrite pv (p + j) probeIdx+                        VUM.unsafeWrite bv (p + j) (ciIxs `VU.unsafeIndex` (start + j))+                        fillBuild probeIdx start len p (j + 1) pv bv+            go 0++            !total <- readSTRef posRef+            pv <- readSTRef pvRef+            bv <- readSTRef bvRef+            (,)+                <$> VU.unsafeFreeze (VUM.slice 0 total pv)+                <*> VU.unsafeFreeze (VUM.slice 0 total bv)+     in (VU.force pFrozen, VU.force bFrozen)+ {- | Hash-based inner join kernel. Builds compact index on @buildHashes@ (second arg), probes with @probeHashes@ (first arg). Returns @(probeExpandedIndices, buildExpandedIndices)@.+Uses a dynamically growing output buffer to avoid pre-allocating the full+cross-product size (which can be astronomically large for low-cardinality keys). -}++{- | Maximum number of output rows allowed from a join kernel.+Exceeding this limit indicates a cross-product explosion (e.g. low-cardinality keys).+-}+maxJoinOutputRows :: Int+maxJoinOutputRows = 500_000_000+ hashInnerKernel ::     VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)-hashInnerKernel probeHashes buildHashes = runST $ do-    let ci = buildCompactIndex buildHashes-        ciIxs = ciSortedIndices ci-        ciOff = ciOffsets ci-        !probeN = VU.length probeHashes+hashInnerKernel probeHashes buildHashes =+    let (pFrozen, bFrozen) = runST $ do+            let ci = buildCompactIndex buildHashes+                ciIxs = ciSortedIndices ci+                ciOff = ciOffsets ci+                !probeN = VU.length probeHashes+                !buildN = VU.length buildHashes+                initCap = max 1 (min (probeN + buildN) 1_000_000) -    -- Pass 1: count total output rows-    let !totalRows =-            VU.foldl'-                ( \acc h ->-                    case HM.lookup h ciOff of-                        Nothing -> acc-                        Just (_, len) -> acc + len-                )-                0-                probeHashes+            initPv <- VUM.unsafeNew initCap+            initBv <- VUM.unsafeNew initCap+            pvRef <- newSTRef initPv+            bvRef <- newSTRef initBv+            capRef <- newSTRef initCap+            posRef <- newSTRef (0 :: Int) -    -- Pass 2: fill output vectors-    pv <- VUM.unsafeNew totalRows-    bv <- VUM.unsafeNew totalRows-    posRef <- newSTRef (0 :: Int)+            let ensureCapacity needed = do+                    cap <- readSTRef capRef+                    when (needed > cap) $ do+                        let newCap = max needed (cap * 2)+                            delta = newCap - cap+                        pv <- readSTRef pvRef+                        bv <- readSTRef bvRef+                        newPv <- VUM.unsafeGrow pv delta+                        newBv <- VUM.unsafeGrow bv delta+                        writeSTRef pvRef newPv+                        writeSTRef bvRef newBv+                        writeSTRef capRef newCap -    let go !i-            | i >= probeN = return ()-            | otherwise = do-                let !h = probeHashes `VU.unsafeIndex` i-                case HM.lookup h ciOff of-                    Nothing -> go (i + 1)-                    Just (!start, !len) -> do-                        !p <- readSTRef posRef-                        fillBuild i start len p 0-                        writeSTRef posRef (p + len)-                        go (i + 1)-        fillBuild !probeIdx !start !len !p !j-            | j >= len = return ()-            | otherwise = do-                VUM.unsafeWrite pv (p + j) probeIdx-                VUM.unsafeWrite bv (p + j) (ciIxs `VU.unsafeIndex` (start + j))-                fillBuild probeIdx start len p (j + 1)-    go 0+                go !i+                    | i >= probeN = return ()+                    | otherwise = do+                        let !h = probeHashes `VU.unsafeIndex` i+                        case HM.lookup h ciOff of+                            Nothing -> go (i + 1)+                            Just (!start, !len) -> do+                                !p <- readSTRef posRef+                                when (p + len > maxJoinOutputRows) $+                                    error $+                                        "Join output would exceed "+                                            ++ show maxJoinOutputRows+                                            ++ " rows (cross-product explosion). "+                                            ++ "Consider filtering or using higher-cardinality join keys or using the lazy API."+                                ensureCapacity (p + len)+                                pv <- readSTRef pvRef+                                bv <- readSTRef bvRef+                                fillBuild i start len p 0 pv bv+                                writeSTRef posRef (p + len)+                                go (i + 1)+                fillBuild !probeIdx !start !len !p !j !pv !bv+                    | j >= len = return ()+                    | otherwise = do+                        VUM.unsafeWrite pv (p + j) probeIdx+                        VUM.unsafeWrite bv (p + j) (ciIxs `VU.unsafeIndex` (start + j))+                        fillBuild probeIdx start len p (j + 1) pv bv+            go 0 -    (,) <$> VU.unsafeFreeze pv <*> VU.unsafeFreeze bv+            !total <- readSTRef posRef+            pv <- readSTRef pvRef+            bv <- readSTRef bvRef+            (,)+                <$> VU.unsafeFreeze (VUM.slice 0 total pv)+                <*> VU.unsafeFreeze (VUM.slice 0 total bv)+     in -- VU.force copies the slice into a compact array, releasing the oversized+        -- backing buffer allocated by the doubling strategy.+        (VU.force pFrozen, VU.force bFrozen)  {- | Sort-merge inner join kernel. Sorts both sides by hash, walks in lockstep. Returns @(leftExpandedIndices, rightExpandedIndices)@.+Uses a dynamically growing output buffer instead of a two-pass count-then-allocate+strategy, which OOMs when low-cardinality keys produce large cross products. -} sortMergeInnerKernel ::     VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)-sortMergeInnerKernel leftHashes rightHashes = runST $ do-    let (leftSH, leftSI) = sortWithIndices leftHashes-        (rightSH, rightSI) = sortWithIndices rightHashes-        !leftN = VU.length leftHashes-        !rightN = VU.length rightHashes+sortMergeInnerKernel leftHashes rightHashes =+    let (lFrozen, rFrozen) = runST $ do+            let (leftSH, leftSI) = sortWithIndices leftHashes+                (rightSH, rightSI) = sortWithIndices rightHashes+                !leftN = VU.length leftHashes+                !rightN = VU.length rightHashes+                initCap = max 1 (min (leftN + rightN) 1_000_000) -    -- Pass 1: count-    let countLoop !li !ri !c-            | li >= leftN || ri >= rightN = c-            | lh < rh = countLoop (li + 1) ri c-            | lh > rh = countLoop li (ri + 1) c-            | otherwise =-                let !lEnd = findGroupEnd leftSH lh (li + 1) leftN-                    !rEnd = findGroupEnd rightSH rh (ri + 1) rightN-                 in countLoop lEnd rEnd (c + (lEnd - li) * (rEnd - ri))-          where-            !lh = leftSH `VU.unsafeIndex` li-            !rh = rightSH `VU.unsafeIndex` ri-        !totalRows = countLoop 0 0 0+            initLv <- VUM.unsafeNew initCap+            initRv <- VUM.unsafeNew initCap+            lvRef <- newSTRef initLv+            rvRef <- newSTRef initRv+            capRef <- newSTRef initCap+            posRef <- newSTRef (0 :: Int) -    -- Pass 2: fill-    lv <- VUM.unsafeNew totalRows-    rv <- VUM.unsafeNew totalRows+            let ensureCapacity needed = do+                    cap <- readSTRef capRef+                    when (needed > cap) $ do+                        let newCap = max needed (cap * 2)+                            delta = newCap - cap+                        lv <- readSTRef lvRef+                        rv <- readSTRef rvRef+                        newLv <- VUM.unsafeGrow lv delta+                        newRv <- VUM.unsafeGrow rv delta+                        writeSTRef lvRef newLv+                        writeSTRef rvRef newRv+                        writeSTRef capRef newCap -    let fill !li !ri !pos-            | li >= leftN || ri >= rightN = return ()-            | lh < rh = fill (li + 1) ri pos-            | lh > rh = fill li (ri + 1) pos-            | otherwise = do-                let !lEnd = findGroupEnd leftSH lh (li + 1) leftN-                    !rEnd = findGroupEnd rightSH rh (ri + 1) rightN-                    !groupSize = (lEnd - li) * (rEnd - ri)-                fillCrossProduct leftSI rightSI li lEnd ri rEnd lv rv pos-                fill lEnd rEnd (pos + groupSize)-          where-            !lh = leftSH `VU.unsafeIndex` li-            !rh = rightSH `VU.unsafeIndex` ri+                fillGroup !li !lEnd !ri !rEnd = do+                    let !lLen = lEnd - li+                        !rLen = rEnd - ri+                        !groupSize = lLen * rLen+                    !p <- readSTRef posRef+                    when (p + groupSize > maxJoinOutputRows) $+                        error $+                            "Join output would exceed "+                                ++ show maxJoinOutputRows+                                ++ " rows (cross-product explosion with group sizes "+                                ++ show lLen+                                ++ " × "+                                ++ show rLen+                                ++ "). Consider filtering or using higher-cardinality join keys."+                    ensureCapacity (p + groupSize)+                    lv <- readSTRef lvRef+                    rv <- readSTRef rvRef+                    let goL !lIdx !pos+                            | lIdx >= lEnd = return ()+                            | otherwise = do+                                let !lOrig = leftSI `VU.unsafeIndex` lIdx+                                goR lOrig ri pos+                                goL (lIdx + 1) (pos + rLen)+                        goR !lOrig !rIdx !pos+                            | rIdx >= rEnd = return ()+                            | otherwise = do+                                VUM.unsafeWrite lv pos lOrig+                                VUM.unsafeWrite rv pos (rightSI `VU.unsafeIndex` rIdx)+                                goR lOrig (rIdx + 1) (pos + 1)+                    goL li p+                    writeSTRef posRef (p + groupSize) -    fill 0 0 0-    (,) <$> VU.unsafeFreeze lv <*> VU.unsafeFreeze rv+                fill !li !ri+                    | li >= leftN || ri >= rightN = return ()+                    | lh < rh = fill (li + 1) ri+                    | lh > rh = fill li (ri + 1)+                    | otherwise = do+                        let !lEnd = findGroupEnd leftSH lh (li + 1) leftN+                            !rEnd = findGroupEnd rightSH rh (ri + 1) rightN+                        fillGroup li lEnd ri rEnd+                        fill lEnd rEnd+                  where+                    !lh = leftSH `VU.unsafeIndex` li+                    !rh = rightSH `VU.unsafeIndex` ri +            fill 0 0++            !total <- readSTRef posRef+            lv <- readSTRef lvRef+            rv <- readSTRef rvRef+            (,)+                <$> VU.unsafeFreeze (VUM.slice 0 total lv)+                <*> VU.unsafeFreeze (VUM.slice 0 total rv)+     in -- VU.force copies the slice into a compact array, releasing the oversized+        -- backing buffer allocated by the doubling strategy.+        (VU.force lFrozen, VU.force rFrozen)+ -- | Assemble the result DataFrame for an inner join from expanded index vectors. assembleInner ::     S.Set T.Text ->@@ -369,27 +527,34 @@ @ -} leftJoin :: [T.Text] -> DataFrame -> DataFrame -> DataFrame-leftJoin cs right left =-    let csSet = S.fromList cs-        rightRows = fst (D.dimensions right)+leftJoin cs left right+    | D.null right || D.nRows right == 0 = left+    | D.null left || D.nRows left == 0 = D.empty+    | otherwise =+        let+            csSet = S.fromList cs+            rightRows = fst (D.dimensions right) -        leftKeyIdxs = keyColIndices csSet left-        rightKeyIdxs = keyColIndices csSet right-        leftHashes = D.computeRowHashes leftKeyIdxs left-        rightHashes = D.computeRowHashes rightKeyIdxs right+            leftKeyIdxs = keyColIndices csSet left+            rightKeyIdxs = keyColIndices csSet right+            leftHashes = D.computeRowHashes leftKeyIdxs left+            rightHashes = D.computeRowHashes rightKeyIdxs right -        -- Right is always the build side for left join-        (leftIxs, rightIxs)-            | rightRows > joinStrategyThreshold =-                sortMergeLeftKernel leftHashes rightHashes-            | otherwise =-                hashLeftKernel leftHashes rightHashes-     in -- rightIxs uses -1 as sentinel for "no match"-        assembleLeft csSet left right leftIxs rightIxs+            -- Right is always the build side for left join+            (leftIxs, rightIxs)+                | rightRows > joinStrategyThreshold =+                    sortMergeLeftKernel leftHashes rightHashes+                | otherwise =+                    hashLeftKernel leftHashes rightHashes+         in+            -- rightIxs uses -1 as sentinel for "no match"+            assembleLeft csSet left right leftIxs rightIxs  {- | Hash-based left join kernel. Returns @(leftExpandedIndices, rightExpandedIndices)@ where right indices use @-1@ as sentinel for unmatched rows.+Uses a dynamically growing output buffer to avoid pre-allocating the full+cross-product size (which can be astronomically large for low-cardinality keys). -} hashLeftKernel ::     VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)@@ -398,49 +563,68 @@         ciIxs = ciSortedIndices ci         ciOff = ciOffsets ci         !leftN = VU.length leftHashes--    -- Pass 1: count-    let !totalRows =-            VU.foldl'-                ( \acc h ->-                    case HM.lookup h ciOff of-                        Nothing -> acc + 1 -- Unmatched left row still produces one output row-                        Just (_, len) -> acc + len-                )-                0-                leftHashes+        !rightN = VU.length rightHashes+        initCap = max 1 (min (leftN + rightN) 1_000_000) -    -- Pass 2: fill-    lv <- VUM.unsafeNew totalRows-    rv <- VUM.unsafeNew totalRows+    initLv <- VUM.unsafeNew initCap+    initRv <- VUM.unsafeNew initCap+    lvRef <- newSTRef initLv+    rvRef <- newSTRef initRv+    capRef <- newSTRef initCap     posRef <- newSTRef (0 :: Int) -    let go !i+    let ensureCapacity needed = do+            cap <- readSTRef capRef+            when (needed > cap) $ do+                let newCap = max needed (cap * 2)+                    delta = newCap - cap+                lv <- readSTRef lvRef+                rv <- readSTRef rvRef+                newLv <- VUM.unsafeGrow lv delta+                newRv <- VUM.unsafeGrow rv delta+                writeSTRef lvRef newLv+                writeSTRef rvRef newRv+                writeSTRef capRef newCap++        go !i             | i >= leftN = return ()             | otherwise = do                 let !h = leftHashes `VU.unsafeIndex` i                 !p <- readSTRef posRef                 case HM.lookup h ciOff of                     Nothing -> do+                        ensureCapacity (p + 1)+                        lv <- readSTRef lvRef+                        rv <- readSTRef rvRef                         VUM.unsafeWrite lv p i                         VUM.unsafeWrite rv p (-1)                         writeSTRef posRef (p + 1)                     Just (!start, !len) -> do-                        fillBuild i start len p 0+                        ensureCapacity (p + len)+                        lv <- readSTRef lvRef+                        rv <- readSTRef rvRef+                        fillBuild i start len p 0 lv rv                         writeSTRef posRef (p + len)                 go (i + 1)-        fillBuild !leftIdx !start !len !p !j+        fillBuild !leftIdx !start !len !p !j !lv !rv             | j >= len = return ()             | otherwise = do                 VUM.unsafeWrite lv (p + j) leftIdx                 VUM.unsafeWrite rv (p + j) (ciIxs `VU.unsafeIndex` (start + j))-                fillBuild leftIdx start len p (j + 1)+                fillBuild leftIdx start len p (j + 1) lv rv     go 0 -    (,) <$> VU.unsafeFreeze lv <*> VU.unsafeFreeze rv+    !total <- readSTRef posRef+    lv <- readSTRef lvRef+    rv <- readSTRef rvRef+    (,)+        <$> VU.unsafeFreeze (VUM.slice 0 total lv)+        <*> VU.unsafeFreeze (VUM.slice 0 total rv)  {- | Sort-merge left join kernel. Returns @(leftExpandedIndices, rightExpandedIndices)@ with @-1@ sentinel.+Uses a dynamically growing output buffer instead of a two-pass count-then-allocate+strategy, which OOMs when low-cardinality keys produce large cross products. -} sortMergeLeftKernel ::     VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)@@ -449,54 +633,98 @@         (rightSH, rightSI) = sortWithIndices rightHashes         !leftN = VU.length leftHashes         !rightN = VU.length rightHashes+        initCap = max 1 (min (leftN + rightN) 1_000_000) -    -- Pass 1: count-    let countLoop !li !ri !c-            | li >= leftN = c-            | ri >= rightN = c + (leftN - li)-            | lh < rh = countLoop (li + 1) ri (c + 1)-            | lh > rh = countLoop li (ri + 1) c-            | otherwise =-                let !lEnd = findGroupEnd leftSH lh (li + 1) leftN-                    !rEnd = findGroupEnd rightSH rh (ri + 1) rightN-                 in countLoop lEnd rEnd (c + (lEnd - li) * (rEnd - ri))-          where-            !lh = leftSH `VU.unsafeIndex` li-            !rh = rightSH `VU.unsafeIndex` ri-        !totalRows = countLoop 0 0 0+    initLv <- VUM.unsafeNew initCap+    initRv <- VUM.unsafeNew initCap+    lvRef <- newSTRef initLv+    rvRef <- newSTRef initRv+    capRef <- newSTRef initCap+    posRef <- newSTRef (0 :: Int) -    -- Pass 2: fill-    lv <- VUM.unsafeNew totalRows-    rv <- VUM.unsafeNew totalRows+    let ensureCapacity needed = do+            cap <- readSTRef capRef+            when (needed > cap) $ do+                let newCap = max needed (cap * 2)+                    delta = newCap - cap+                lv <- readSTRef lvRef+                rv <- readSTRef rvRef+                newLv <- VUM.unsafeGrow lv delta+                newRv <- VUM.unsafeGrow rv delta+                writeSTRef lvRef newLv+                writeSTRef rvRef newRv+                writeSTRef capRef newCap -    let fill !li !ri !pos+        fillGroup !li !lEnd !ri !rEnd = do+            let !lLen = lEnd - li+                !rLen = rEnd - ri+                !groupSize = lLen * rLen+            !p <- readSTRef posRef+            ensureCapacity (p + groupSize)+            lv <- readSTRef lvRef+            rv <- readSTRef rvRef+            let goL !lIdx !pos+                    | lIdx >= lEnd = return ()+                    | otherwise = do+                        let !lOrig = leftSI `VU.unsafeIndex` lIdx+                        goR lOrig ri pos+                        goL (lIdx + 1) (pos + rLen)+                goR !lOrig !rIdx !pos+                    | rIdx >= rEnd = return ()+                    | otherwise = do+                        VUM.unsafeWrite lv pos lOrig+                        VUM.unsafeWrite rv pos (rightSI `VU.unsafeIndex` rIdx)+                        goR lOrig (rIdx + 1) (pos + 1)+            goL li p+            writeSTRef posRef (p + groupSize)++        fill !li !ri             | li >= leftN = return ()-            | ri >= rightN = fillRemainingLeft li pos+            | ri >= rightN = fillRemainingLeft li             | lh < rh = do-                VUM.unsafeWrite lv pos (leftSI `VU.unsafeIndex` li)-                VUM.unsafeWrite rv pos (-1)-                fill (li + 1) ri (pos + 1)-            | lh > rh = fill li (ri + 1) pos+                !p <- readSTRef posRef+                ensureCapacity (p + 1)+                lv <- readSTRef lvRef+                rv <- readSTRef rvRef+                VUM.unsafeWrite lv p (leftSI `VU.unsafeIndex` li)+                VUM.unsafeWrite rv p (-1)+                writeSTRef posRef (p + 1)+                fill (li + 1) ri+            | lh > rh = fill li (ri + 1)             | otherwise = do                 let !lEnd = findGroupEnd leftSH lh (li + 1) leftN                     !rEnd = findGroupEnd rightSH rh (ri + 1) rightN-                    !groupSize = (lEnd - li) * (rEnd - ri)-                fillCrossProduct leftSI rightSI li lEnd ri rEnd lv rv pos-                fill lEnd rEnd (pos + groupSize)+                fillGroup li lEnd ri rEnd+                fill lEnd rEnd           where             !lh = leftSH `VU.unsafeIndex` li             !rh = rightSH `VU.unsafeIndex` ri -        fillRemainingLeft !i !pos-            | i >= leftN = return ()-            | otherwise = do-                VUM.unsafeWrite lv pos (leftSI `VU.unsafeIndex` i)-                VUM.unsafeWrite rv pos (-1)-                fillRemainingLeft (i + 1) (pos + 1)+        fillRemainingLeft !i = do+            let !remaining = leftN - i+            when (remaining > 0) $ do+                !p <- readSTRef posRef+                ensureCapacity (p + remaining)+                lv <- readSTRef lvRef+                rv <- readSTRef rvRef+                let go !j+                        | j >= remaining = return ()+                        | otherwise = do+                            VUM.unsafeWrite lv (p + j) (leftSI `VU.unsafeIndex` (i + j))+                            VUM.unsafeWrite rv (p + j) (-1)+                            go (j + 1)+                go 0+                writeSTRef posRef (p + remaining) -    fill 0 0 0-    (,) <$> VU.unsafeFreeze lv <*> VU.unsafeFreeze rv+    fill 0 0 +    !total <- readSTRef posRef+    lv <- readSTRef lvRef+    rv <- readSTRef rvRef+    (,)+        <$> VU.unsafeFreeze (VUM.slice 0 total lv)+        <*> VU.unsafeFreeze (VUM.slice 0 total rv)+ {- | Assemble the result DataFrame for a left join. Right index vectors use @-1@ sentinel, gathered via 'gatherWithSentinel'. -}@@ -574,24 +802,29 @@  fullOuterJoin ::     [T.Text] -> DataFrame -> DataFrame -> DataFrame-fullOuterJoin cs right left =-    let csSet = S.fromList cs-        leftRows = fst (D.dimensions left)-        rightRows = fst (D.dimensions right)+fullOuterJoin cs left right+    | D.null right || D.nRows right == 0 = left+    | D.null left || D.nRows left == 0 = right+    | otherwise =+        let+            csSet = S.fromList cs+            leftRows = fst (D.dimensions left)+            rightRows = fst (D.dimensions right) -        leftKeyIdxs = keyColIndices csSet left-        rightKeyIdxs = keyColIndices csSet right-        leftHashes = D.computeRowHashes leftKeyIdxs left-        rightHashes = D.computeRowHashes rightKeyIdxs right+            leftKeyIdxs = keyColIndices csSet left+            rightKeyIdxs = keyColIndices csSet right+            leftHashes = D.computeRowHashes leftKeyIdxs left+            rightHashes = D.computeRowHashes rightKeyIdxs right -        -- Both sides can have nulls in full outer-        (leftIxs, rightIxs)-            | max leftRows rightRows > joinStrategyThreshold =-                sortMergeFullOuterKernel leftHashes rightHashes-            | otherwise =-                hashFullOuterKernel leftHashes rightHashes-     in -- Both index vectors use -1 as sentinel-        assembleFullOuter csSet left right leftIxs rightIxs+            -- Both sides can have nulls in full outer+            (leftIxs, rightIxs)+                | max leftRows rightRows > joinStrategyThreshold =+                    sortMergeFullOuterKernel leftHashes rightHashes+                | otherwise =+                    hashFullOuterKernel leftHashes rightHashes+         in+            -- Both index vectors use -1 as sentinel+            assembleFullOuter csSet left right leftIxs rightIxs  {- | Hash-based full outer join kernel. Builds compact indices on both sides.

@@ -192,9 +192,6 @@     -- * Constraints     KnownSchema (..),     AllKnownSymbol (..),--    -- * Pipe operator-    (|>), ) where  import Prelude hiding (drop, filter, take)

@@ -86,7 +86,6 @@ import Data.Proxy (Proxy (..)) import Data.String (IsString (..)) import qualified Data.Text as T-import qualified Data.Vector.Unboxed as VU import GHC.TypeLits (KnownSymbol, Symbol, symbolVal)  import DataFrame.Internal.Column (Columnable)@@ -94,11 +93,12 @@     AggStrategy (..),     BinaryOp (..),     Expr (..),+    MeanAcc (..),     NamedExpr,     UExpr (..),     UnaryOp (..),  )-import DataFrame.Internal.Statistics+ import DataFrame.Typed.Schema (AssertPresent, Lookup) import DataFrame.Typed.Types (TExpr (..), TSortOrder (..)) import Prelude hiding (maximum, minimum, sum)@@ -233,11 +233,22 @@ sum :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a sum (TExpr e) = TExpr (Agg (FoldAgg "sum" Nothing (+)) e) -mean :: (Columnable a, Real a, VU.Unbox a) => TExpr cols a -> TExpr cols Double-mean (TExpr e) = TExpr (Agg (CollectAgg "mean" mean') e)+mean :: (Columnable a, Real a) => TExpr cols a -> TExpr cols Double+mean (TExpr e) =+    TExpr+        ( Agg+            ( MergeAgg+                "mean"+                (MeanAcc 0.0 0)+                (\(MeanAcc s c) x -> MeanAcc (s + realToFrac x) (c + 1))+                (\(MeanAcc s1 c1) (MeanAcc s2 c2) -> MeanAcc (s1 + s2) (c1 + c2))+                (\(MeanAcc s c) -> if c == 0 then 0 / 0 else s / fromIntegral c)+            )+            e+        )  count :: (Columnable a) => TExpr cols a -> TExpr cols Int-count (TExpr e) = TExpr (Agg (FoldAgg "count" (Just 0) (\acc _ -> acc + 1)) e)+count (TExpr e) = TExpr (Agg (MergeAgg "count" (0 :: Int) (\c _ -> c + 1) (+) id) e)  minimum :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a minimum (TExpr e) = TExpr (Agg (FoldAgg "minimum" Nothing min) e)

@@ -43,7 +43,7 @@     TypedDataFrame right ->     TypedDataFrame (LeftJoinSchema keys left right) leftJoin (TDF l) (TDF r) =-    unsafeFreeze (DJ.leftJoin keyNames r l)+    unsafeFreeze (DJ.leftJoin keyNames l r)   where     keyNames = symbolVals @keys @@ -55,7 +55,7 @@     TypedDataFrame right ->     TypedDataFrame (RightJoinSchema keys left right) rightJoin (TDF l) (TDF r) =-    unsafeFreeze (DJ.rightJoin keyNames r l)+    unsafeFreeze (DJ.rightJoin keyNames l r)   where     keyNames = symbolVals @keys 

@@ -50,12 +50,8 @@      -- * Vertical merge     append,--    -- * Pipe operator-    (|>), ) where -import Data.Function ((&)) import Data.Proxy (Proxy (..)) import qualified Data.Text as T import qualified Data.Vector as V@@ -73,19 +69,11 @@ import qualified DataFrame.Operations.Subset as D import qualified DataFrame.Operations.Transformations as D --- Semigroup instance- import DataFrame.Typed.Freeze (unsafeFreeze) import DataFrame.Typed.Schema import DataFrame.Typed.Types (TExpr (..), TSortOrder (..), TypedDataFrame (..)) import qualified DataFrame.Typed.Types as T --- | Pipe operator, re-exported for convenience.-(|>) :: a -> (a -> b) -> b-(|>) = (&)--infixl 1 |>- ------------------------------------------------------------------------------- -- Schema-preserving operations -------------------------------------------------------------------------------@@ -226,6 +214,7 @@     forall name a cols.     ( KnownSymbol name     , Columnable a+    , Maybe a ~ Lookup name cols     ) =>     a ->     TypedDataFrame cols ->

@@ -0,0 +1,112 @@+{-# LANGUAGE OverloadedStrings #-}++module IO.JSON where++import qualified Data.ByteString.Lazy.Char8 as LBSC+import DataFrame.IO.JSON (readJSONEither)+import qualified DataFrame.Internal.Column as DI+import qualified DataFrame.Internal.DataFrame as DI+import qualified DataFrame.Operations.Core as D+import Test.HUnit (+    Test (TestCase, TestLabel),+    assertBool,+    assertEqual,+    assertFailure,+ )++-- | Happy path: array of objects with string and number columns.+jsonHappyPath :: Test+jsonHappyPath =+    TestCase+        ( case readJSONEither+            (LBSC.pack "[{\"name\":\"Alice\",\"age\":30},{\"name\":\"Bob\",\"age\":25}]") of+            Left err -> assertFailure $ "Unexpected Left: " ++ err+            Right df -> do+                assertEqual "Happy path: 2 rows" 2 (D.nRows df)+                assertEqual "Happy path: 2 columns" 2 (D.nColumns df)+        )++-- | Boolean column is preserved correctly.+jsonBoolColumn :: Test+jsonBoolColumn =+    TestCase+        ( case readJSONEither (LBSC.pack "[{\"flag\":true},{\"flag\":false}]") of+            Left err -> assertFailure $ "Unexpected Left: " ++ err+            Right df -> do+                assertEqual "Bool column: 2 rows" 2 (D.nRows df)+                assertEqual "Bool column: 1 column" 1 (D.nColumns df)+        )++-- | A key absent from some objects produces an Optional column.+jsonMissingKeyBecomesOptional :: Test+jsonMissingKeyBecomesOptional =+    TestCase+        ( case readJSONEither (LBSC.pack "[{\"a\":1,\"b\":2},{\"a\":3}]") of+            Left err -> assertFailure $ "Unexpected Left: " ++ err+            Right df -> do+                assertEqual "Missing key: 2 rows" 2 (D.nRows df)+                assertEqual "Missing key: 2 columns" 2 (D.nColumns df)+                -- 'b' is absent from the second row, so must have a missing value+                assertBool+                    "b column should have missing values"+                    (maybe False DI.hasMissing (DI.getColumn "b" df))+        )++-- | When column values are different types across rows, the column becomes generic.+jsonMixedTypeColumn :: Test+jsonMixedTypeColumn =+    TestCase+        ( case readJSONEither (LBSC.pack "[{\"x\":1},{\"x\":\"hello\"}]") of+            Left err -> assertFailure $ "Unexpected Left: " ++ err+            Right df -> do+                assertEqual "Mixed type: 2 rows" 2 (D.nRows df)+                assertEqual "Mixed type: 1 column" 1 (D.nColumns df)+        )++-- | An empty top-level JSON array is rejected.+jsonEmptyArray :: Test+jsonEmptyArray =+    TestCase+        ( case readJSONEither (LBSC.pack "[]") of+            Left _ -> return ()+            Right _ -> assertFailure "Expected Left for empty array"+        )++-- | A non-array top-level JSON value is rejected.+jsonNonArray :: Test+jsonNonArray =+    TestCase+        ( case readJSONEither (LBSC.pack "{\"name\":\"Alice\"}") of+            Left _ -> return ()+            Right _ -> assertFailure "Expected Left for non-array top-level value"+        )++-- | Array elements that are not objects are rejected.+jsonNonObjectElement :: Test+jsonNonObjectElement =+    TestCase+        ( case readJSONEither (LBSC.pack "[1, 2, 3]") of+            Left _ -> return ()+            Right _ -> assertFailure "Expected Left for non-object array elements"+        )++-- | Completely invalid JSON is rejected.+jsonInvalidJSON :: Test+jsonInvalidJSON =+    TestCase+        ( case readJSONEither (LBSC.pack "not valid json at all") of+            Left _ -> return ()+            Right _ -> assertFailure "Expected Left for invalid JSON"+        )++tests :: [Test]+tests =+    [ TestLabel "jsonHappyPath" jsonHappyPath+    , TestLabel "jsonBoolColumn" jsonBoolColumn+    , TestLabel "jsonMissingKeyBecomesOptional" jsonMissingKeyBecomesOptional+    , TestLabel "jsonMixedTypeColumn" jsonMixedTypeColumn+    , TestLabel "jsonEmptyArray" jsonEmptyArray+    , TestLabel "jsonNonArray" jsonNonArray+    , TestLabel "jsonNonObjectElement" jsonNonObjectElement+    , TestLabel "jsonInvalidJSON" jsonInvalidJSON+    ]

@@ -0,0 +1,237 @@+{-# LANGUAGE OverloadedStrings #-}++module Internal.Parsing where++import DataFrame.Internal.Parsing+import Test.HUnit++-- isNullish: recognized null strings++isNullishEmptyString :: Test+isNullishEmptyString =+    TestCase (assertBool "empty string is nullish" (isNullish ""))++isNullishNA :: Test+isNullishNA = TestCase (assertBool "NA is nullish" (isNullish "NA"))++isNullishNULL :: Test+isNullishNULL = TestCase (assertBool "NULL is nullish" (isNullish "NULL"))++isNullishNull :: Test+isNullishNull = TestCase (assertBool "null is nullish" (isNullish "null"))++isNullishNaN :: Test+isNullishNaN = TestCase (assertBool "nan is nullish" (isNullish "nan"))++isNullishNaNMixed :: Test+isNullishNaNMixed = TestCase (assertBool "NaN is nullish" (isNullish "NaN"))++isNullishNANUpper :: Test+isNullishNANUpper = TestCase (assertBool "NAN is nullish" (isNullish "NAN"))++isNullishNothing :: Test+isNullishNothing =+    TestCase (assertBool "Nothing is nullish" (isNullish "Nothing"))++isNullishSpace :: Test+isNullishSpace =+    TestCase (assertBool "single space is nullish" (isNullish " "))++isNullishNSlashA :: Test+isNullishNSlashA = TestCase (assertBool "N/A is nullish" (isNullish "N/A"))++-- isNullish: values that are NOT null++notNullishNumber :: Test+notNullishNumber =+    TestCase (assertBool "\"42\" is not nullish" (not (isNullish "42")))++notNullishText :: Test+notNullishText =+    TestCase (assertBool "\"hello\" is not nullish" (not (isNullish "hello")))++notNullishTrue :: Test+notNullishTrue =+    TestCase (assertBool "\"True\" is not nullish" (not (isNullish "True")))++notNullishDouble :: Test+notNullishDouble =+    TestCase (assertBool "\"3.14\" is not nullish" (not (isNullish "3.14")))++notNullishZero :: Test+notNullishZero =+    TestCase (assertBool "\"0\" is not nullish" (not (isNullish "0")))++-- readBool: positive cases++readBoolTrue :: Test+readBoolTrue =+    TestCase (assertEqual "readBool \"True\"" (Just True) (readBool "True"))++readBoolTrueLower :: Test+readBoolTrueLower =+    TestCase (assertEqual "readBool \"true\"" (Just True) (readBool "true"))++readBoolTrueUpper :: Test+readBoolTrueUpper =+    TestCase (assertEqual "readBool \"TRUE\"" (Just True) (readBool "TRUE"))++readBoolFalse :: Test+readBoolFalse =+    TestCase (assertEqual "readBool \"False\"" (Just False) (readBool "False"))++readBoolFalseLower :: Test+readBoolFalseLower =+    TestCase (assertEqual "readBool \"false\"" (Just False) (readBool "false"))++readBoolFalseUpper :: Test+readBoolFalseUpper =+    TestCase (assertEqual "readBool \"FALSE\"" (Just False) (readBool "FALSE"))++-- readBool: values that are not booleans++readBoolDigit :: Test+readBoolDigit =+    TestCase (assertEqual "readBool \"1\" is Nothing" Nothing (readBool "1"))++readBoolYes :: Test+readBoolYes =+    TestCase (assertEqual "readBool \"yes\" is Nothing" Nothing (readBool "yes"))++readBoolEmpty :: Test+readBoolEmpty =+    TestCase (assertEqual "readBool \"\" is Nothing" Nothing (readBool ""))++readBoolPartialTrue :: Test+readBoolPartialTrue =+    TestCase (assertEqual "readBool \"Tru\" is Nothing" Nothing (readBool "Tru"))++-- readInt++readIntPositive :: Test+readIntPositive =+    TestCase (assertEqual "readInt \"42\"" (Just 42) (readInt "42"))++readIntNegative :: Test+readIntNegative =+    TestCase (assertEqual "readInt \"-17\"" (Just (-17)) (readInt "-17"))++readIntZero :: Test+readIntZero =+    TestCase (assertEqual "readInt \"0\"" (Just 0) (readInt "0"))++-- readInt strips whitespace before parsing+readIntLeadingSpace :: Test+readIntLeadingSpace =+    TestCase+        ( assertEqual+            "readInt \" 5 \" (strips whitespace)"+            (Just 5)+            (readInt " 5 ")+        )++readIntFloat :: Test+readIntFloat =+    TestCase+        (assertEqual "readInt \"3.14\" is Nothing" Nothing (readInt "3.14"))++readIntText :: Test+readIntText =+    TestCase (assertEqual "readInt \"abc\" is Nothing" Nothing (readInt "abc"))++readIntEmpty :: Test+readIntEmpty =+    TestCase (assertEqual "readInt \"\" is Nothing" Nothing (readInt ""))++-- trailing non-digits must make the parse fail+readIntPartialSuffix :: Test+readIntPartialSuffix =+    TestCase+        ( assertEqual+            "readInt \"42abc\" is Nothing"+            Nothing+            (readInt "42abc")+        )++-- readDouble++readDoublePositive :: Test+readDoublePositive =+    TestCase+        (assertEqual "readDouble \"3.14\"" (Just 3.14) (readDouble "3.14"))++readDoubleNegative :: Test+readDoubleNegative =+    TestCase+        (assertEqual "readDouble \"-1.5\"" (Just (-1.5)) (readDouble "-1.5"))++readDoubleWholeNumber :: Test+readDoubleWholeNumber =+    TestCase+        ( assertEqual+            "readDouble \"42\" parses as 42.0"+            (Just 42.0)+            (readDouble "42")+        )++readDoubleText :: Test+readDoubleText =+    TestCase+        (assertEqual "readDouble \"abc\" is Nothing" Nothing (readDouble "abc"))++readDoubleEmpty :: Test+readDoubleEmpty =+    TestCase+        (assertEqual "readDouble \"\" is Nothing" Nothing (readDouble ""))++readDoublePartialSuffix :: Test+readDoublePartialSuffix =+    TestCase+        ( assertEqual+            "readDouble \"3.14abc\" is Nothing"+            Nothing+            (readDouble "3.14abc")+        )++tests :: [Test]+tests =+    [ TestLabel "isNullishEmptyString" isNullishEmptyString+    , TestLabel "isNullishNA" isNullishNA+    , TestLabel "isNullishNULL" isNullishNULL+    , TestLabel "isNullishNull" isNullishNull+    , TestLabel "isNullishNaN" isNullishNaN+    , TestLabel "isNullishNaNMixed" isNullishNaNMixed+    , TestLabel "isNullishNANUpper" isNullishNANUpper+    , TestLabel "isNullishNothing" isNullishNothing+    , TestLabel "isNullishSpace" isNullishSpace+    , TestLabel "isNullishNSlashA" isNullishNSlashA+    , TestLabel "notNullishNumber" notNullishNumber+    , TestLabel "notNullishText" notNullishText+    , TestLabel "notNullishTrue" notNullishTrue+    , TestLabel "notNullishDouble" notNullishDouble+    , TestLabel "notNullishZero" notNullishZero+    , TestLabel "readBoolTrue" readBoolTrue+    , TestLabel "readBoolTrueLower" readBoolTrueLower+    , TestLabel "readBoolTrueUpper" readBoolTrueUpper+    , TestLabel "readBoolFalse" readBoolFalse+    , TestLabel "readBoolFalseLower" readBoolFalseLower+    , TestLabel "readBoolFalseUpper" readBoolFalseUpper+    , TestLabel "readBoolDigit" readBoolDigit+    , TestLabel "readBoolYes" readBoolYes+    , TestLabel "readBoolEmpty" readBoolEmpty+    , TestLabel "readBoolPartialTrue" readBoolPartialTrue+    , TestLabel "readIntPositive" readIntPositive+    , TestLabel "readIntNegative" readIntNegative+    , TestLabel "readIntZero" readIntZero+    , TestLabel "readIntLeadingSpace" readIntLeadingSpace+    , TestLabel "readIntFloat" readIntFloat+    , TestLabel "readIntText" readIntText+    , TestLabel "readIntEmpty" readIntEmpty+    , TestLabel "readIntPartialSuffix" readIntPartialSuffix+    , TestLabel "readDoublePositive" readDoublePositive+    , TestLabel "readDoubleNegative" readDoubleNegative+    , TestLabel "readDoubleWholeNumber" readDoubleWholeNumber+    , TestLabel "readDoubleText" readDoubleText+    , TestLabel "readDoubleEmpty" readDoubleEmpty+    , TestLabel "readDoublePartialSuffix" readDoublePartialSuffix+    ]