dataframe-core 1.1.0.4 → 2.0.0.0

raw patch · 64 files changed

+9919/−10147 lines, 64 filesdep +dataframe-coredep ~containersdep ~primitivedep ~vectorPVP ok

version bump matches the API change (PVP)

Dependencies added: dataframe-core

Dependency ranges changed: containers, primitive, vector

API changes (from Hackage documentation)

Files

• dataframe-core.cabal +22/−7
• src-internal/DataFrame/Display/Terminal/Colours.hs +7/−0
• src-internal/DataFrame/Display/Terminal/PrettyPrint.hs +114/−0
• src-internal/DataFrame/Errors.hs +187/−0
• src-internal/DataFrame/Internal/AggKernel.hs +260/−0
• src-internal/DataFrame/Internal/AggKernelDirect.hs +338/−0
• src-internal/DataFrame/Internal/AggKernelPar.hs +391/−0
• src-internal/DataFrame/Internal/AggPlan.hs +299/−0
• src-internal/DataFrame/Internal/Column.hs +1744/−0
• src-internal/DataFrame/Internal/ColumnBuilder.hs +299/−0
• src-internal/DataFrame/Internal/ColumnMerge.hs +176/−0
• src-internal/DataFrame/Internal/DataFrame.hs +362/−0
• src-internal/DataFrame/Internal/DictEncode.hs +129/−0
• src-internal/DataFrame/Internal/Expression.hs +511/−0
• src-internal/DataFrame/Internal/Grouping.hs +410/−0
• src-internal/DataFrame/Internal/GroupingDirect.hs +234/−0
• src-internal/DataFrame/Internal/GroupingPar.hs +308/−0
• src-internal/DataFrame/Internal/Hash.hs +113/−0
• src-internal/DataFrame/Internal/HashTable.hs +94/−0
• src-internal/DataFrame/Internal/Interpreter.hs +1051/−0
• src-internal/DataFrame/Internal/Nullable.hs +467/−0
• src-internal/DataFrame/Internal/PackedText.hs +149/−0
• src-internal/DataFrame/Internal/ParRadixSort.hs +272/−0
• src-internal/DataFrame/Internal/Pretty.hs +129/−0
• src-internal/DataFrame/Internal/RadixRank.hs +101/−0
• src-internal/DataFrame/Internal/Row.hs +181/−0
• src-internal/DataFrame/Internal/RowHash.hs +216/−0
• src-internal/DataFrame/Internal/Simplify.hs +417/−0
• src-internal/DataFrame/Internal/Types.hs +161/−0
• src-internal/DataFrame/Internal/Utf8.hs +95/−0
• src-internal/DataFrame/Operators.hs +425/−0
• src/DataFrame/Core.hs +104/−0
• src/DataFrame/Display/Terminal/Colours.hs +0/−7
• src/DataFrame/Display/Terminal/PrettyPrint.hs +0/−118
• src/DataFrame/Errors.hs +0/−187
• src/DataFrame/Internal/AggKernel.hs +0/−308
• src/DataFrame/Internal/AggKernelDirect.hs +0/−373
• src/DataFrame/Internal/AggKernelPar.hs +0/−454
• src/DataFrame/Internal/AggPlan.hs +0/−316
• src/DataFrame/Internal/Column.hs +0/−1852
• src/DataFrame/Internal/ColumnBuilder.hs +0/−299
• src/DataFrame/Internal/ColumnMerge.hs +0/−179
• src/DataFrame/Internal/DataFrame.hs +0/−378
• src/DataFrame/Internal/DictEncode.hs +0/−159
• src/DataFrame/Internal/Expression.hs +0/−472
• src/DataFrame/Internal/Grouping.hs +0/−454
• src/DataFrame/Internal/GroupingDirect.hs +0/−260
• src/DataFrame/Internal/GroupingPar.hs +0/−355
• src/DataFrame/Internal/Hash.hs +0/−128
• src/DataFrame/Internal/HashTable.hs +0/−113
• src/DataFrame/Internal/Interpreter.hs +0/−1103
• src/DataFrame/Internal/Nullable.hs +0/−500
• src/DataFrame/Internal/PackedText.hs +0/−169
• src/DataFrame/Internal/ParRadixSort.hs +0/−294
• src/DataFrame/Internal/RadixRank.hs +0/−114
• src/DataFrame/Internal/Row.hs +0/−209
• src/DataFrame/Internal/RowHash.hs +0/−232
• src/DataFrame/Internal/Simplify.hs +0/−422
• src/DataFrame/Internal/Types.hs +0/−161
• src/DataFrame/Internal/Utf8.hs +0/−98
• src/DataFrame/Operators.hs +0/−425
• src/DataFrame/Typed/Freeze.hs +23/−0
• src/DataFrame/Typed/Schema.hs +105/−0
• src/DataFrame/Typed/Types.hs +25/−1

@@ -1,6 +1,6 @@-cabal-version:      2.4+cabal-version:      3.4 name:               dataframe-core-version:            1.1.0.4+version:            2.0.0.0 synopsis:           Core data structures for the dataframe library. description:     Minimal interchange-format types for the @dataframe@ ecosystem:@@ -28,8 +28,9 @@         -Wunused-local-binds         -Wunused-packages -library+library internal     import:             warnings+    visibility:         public     exposed-modules:                         DataFrame.Errors                         DataFrame.Operators@@ -54,23 +55,37 @@                         DataFrame.Internal.Nullable                         DataFrame.Internal.PackedText                         DataFrame.Internal.ParRadixSort+                        DataFrame.Internal.Pretty                         DataFrame.Internal.RadixRank                         DataFrame.Internal.RowHash                         DataFrame.Internal.Row                         DataFrame.Internal.Simplify                         DataFrame.Internal.Types                         DataFrame.Internal.Utf8+    build-depends:      base >= 4 && < 5,+                        containers >= 0.6.7 && < 0.10,+                        primitive >= 0.7 && < 0.11,+                        random >= 1 && < 2,+                        text >= 2.1 && < 3,+                        vector >= 0.13 && < 0.15+    hs-source-dirs:     src-internal+    default-language:   Haskell2010++library+    import:             warnings+    exposed-modules:+                        DataFrame.Core                         DataFrame.Typed.Freeze                         DataFrame.Typed.Generic                         DataFrame.Typed.Record                         DataFrame.Typed.Schema                         DataFrame.Typed.Types                         DataFrame.Typed.Util+    reexported-modules: DataFrame.Errors,+                        DataFrame.Operators     build-depends:      base >= 4 && < 5,-                        containers >= 0.6.7 && < 0.9,-                        primitive >= 0.7 && < 0.10,-                        random >= 1 && < 2,                         text >= 2.1 && < 3,-                        vector ^>= 0.13+                        vector >= 0.13 && < 0.15,+                        dataframe-core:internal     hs-source-dirs:     src     default-language:   Haskell2010

@@ -0,0 +1,7 @@+module DataFrame.Display.Terminal.Colours where++red, green, brightGreen, brightBlue :: String -> String+red = id+green = id+brightGreen = id+brightBlue = id

@@ -0,0 +1,114 @@+{-# LANGUAGE OverloadedStrings #-}++module DataFrame.Display.Terminal.PrettyPrint where++import qualified Data.Text as T+import qualified Data.Vector as V++{- | Output format for 'showTable'. 'Plain' renders a terminal-style table with+ASCII borders; 'Markdown' renders a GitHub-flavoured pipe table suitable for+notebooks.+-}+data RenderFormat = Plain | Markdown+    deriving (Show, Eq)++-- Utility functions to show a DataFrame as a Markdown-ish table.++-- Fill functions, adapted from:+-- https://stackoverflow.com/questions/5929377/format-list-output-in-haskell+type Filler = Int -> T.Text -> T.Text++-- A description of one table column.+data ColDesc t = ColDesc+    { colTitleFill :: Filler+    , colTitle :: T.Text+    , colValueFill :: Filler+    }++-- Fill text to a given width with a pad character, aligning left, right, or center.+fillLeft :: Char -> Int -> T.Text -> T.Text+fillLeft c n s = s <> T.replicate (n - T.length s) (T.singleton c)++fillRight :: Char -> Int -> T.Text -> T.Text+fillRight c n s = T.replicate (n - T.length s) (T.singleton c) <> s++fillCenter :: Char -> Int -> T.Text -> T.Text+fillCenter c n s =+    T.replicate l (T.singleton c) <> s <> T.replicate r (T.singleton c)+  where+    x = n - T.length s+    l = x `div` 2+    r = x - l++-- Fill with spaces.+left :: Int -> T.Text -> T.Text+left = fillLeft ' '++right :: Int -> T.Text -> T.Text+right = fillRight ' '++center :: Int -> T.Text -> T.Text+center = fillCenter ' '++{- | Render a table from column-major data. @columns@ has one 'V.Vector' per+column; widths are computed in one pass per column (no row-major transpose),+and row lines are built by indexing each column at row @i@.+-}+showTable ::+    RenderFormat ->+    [T.Text] ->+    [T.Text] ->+    [V.Vector T.Text] ->+    T.Text+showTable fmt header types columns =+    let isMarkdown = fmt == Markdown+        esc = if isMarkdown then escapeMarkdownCell else id+        hdr = map esc header+        tys = map esc types+        cols = map (V.map esc) columns+        consolidatedHeader =+            if isMarkdown+                then zipWith (\h t -> h <> "<br>" <> t) hdr tys+                else hdr+        cs = map (\h -> ColDesc center h left) consolidatedHeader+        nRows = case cols of+            (c : _) -> V.length c+            [] -> 0+        columnMaxWidth col+            | V.null col = 0+            | otherwise = V.foldl' (\acc x -> max acc (T.length x)) 0 col+        widths =+            zipWith3+                (\h t col -> T.length h `max` T.length t `max` columnMaxWidth col)+                consolidatedHeader+                tys+                cols+        dashesOf w = T.replicate w "-"+        border = T.intercalate "---" (map dashesOf widths)+        separator = T.intercalate "-|-" (map dashesOf widths)+        fillCells fill cells =+            T.intercalate " | " (zipWith3 fill cs widths cells)+        rowCells i = map (V.! i) cols+        rowLines = [fillCells colValueFill (rowCells i) | i <- [0 .. nRows - 1]]+        wrapMd t = T.concat ["| ", t, " |"]+        outputLines =+            if isMarkdown+                then+                    wrapMd (fillCells colTitleFill consolidatedHeader)+                        : wrapMd separator+                        : map wrapMd rowLines+                else+                    border+                        : fillCells colTitleFill consolidatedHeader+                        : separator+                        : fillCells colTitleFill tys+                        : separator+                        : rowLines+     in T.unlines outputLines++{- | Escape a value for a GitHub-flavoured Markdown table cell: a bare @|@ would+end the cell and a newline would end the row, so both are neutralised (the pipe+backslash-escaped, the newline turned into a @\<br\>@).+-}+escapeMarkdownCell :: T.Text -> T.Text+escapeMarkdownCell = T.replace "\n" "<br>" . T.replace "|" "\\|"

@@ -0,0 +1,187 @@+{-# LANGUAGE DeriveAnyClass #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RankNTypes #-}++module DataFrame.Errors where++import qualified Data.Map.Lazy as ML+import qualified Data.Text as T+import qualified Data.Vector as V+import qualified Data.Vector.Unboxed as VU++import Control.Exception+import qualified Data.List as L+import Data.Typeable (Typeable)+import DataFrame.Display.Terminal.Colours+import Type.Reflection (TypeRep)++data TypeErrorContext a b = MkTypeErrorContext+    { userType :: Either String (TypeRep a)+    , expectedType :: Either String (TypeRep b)+    , errorColumnName :: Maybe String+    , callingFunctionName :: Maybe String+    }++data DataFrameException where+    TypeMismatchException ::+        forall a b.+        (Typeable a, Typeable b) =>+        TypeErrorContext a b ->+        DataFrameException+    AggregatedAndNonAggregatedException :: T.Text -> T.Text -> DataFrameException+    ColumnsNotFoundException :: [T.Text] -> T.Text -> [T.Text] -> DataFrameException+    EmptyDataSetException :: T.Text -> DataFrameException+    InternalException :: T.Text -> DataFrameException+    NonColumnReferenceException :: T.Text -> DataFrameException+    UnaggregatedException :: T.Text -> DataFrameException+    WrongQuantileNumberException :: Int -> DataFrameException+    WrongQuantileIndexException :: VU.Vector Int -> Int -> DataFrameException+    deriving (Exception)++instance Show DataFrameException where+    show :: DataFrameException -> String+    show (TypeMismatchException context) =+        let+            errorString =+                typeMismatchError+                    (either id show (userType context))+                    (either id show (expectedType context))+         in+            addCallPointInfo+                (errorColumnName context)+                (callingFunctionName context)+                errorString+    show (ColumnsNotFoundException columnNames callPoint availableColumns) = columnsNotFound columnNames callPoint availableColumns+    show (EmptyDataSetException callPoint) = emptyDataSetError callPoint+    show (WrongQuantileNumberException q) = wrongQuantileNumberError q+    show (WrongQuantileIndexException qs q) = wrongQuantileIndexError qs q+    show (InternalException msg) = "Internal error: " ++ T.unpack msg+    show (NonColumnReferenceException msg) = "Expression must be a column reference in: " ++ T.unpack msg+    show (UnaggregatedException expr) = "Expression is not fully aggregated: " ++ T.unpack expr+    show (AggregatedAndNonAggregatedException expr1 expr2) =+        "Cannot combine aggregated and non-aggregated expressions: \n"+            ++ T.unpack expr1+            ++ "\n"+            ++ T.unpack expr2++columnNotFound :: T.Text -> T.Text -> [T.Text] -> String+columnNotFound missingColumn = columnsNotFound [missingColumn]++columnsNotFound :: [T.Text] -> T.Text -> [T.Text] -> String+columnsNotFound missingColumns callPoint availableColumns =+    red "\n\n[ERROR] "+        ++ missingColumnsLabel missingColumns+        ++ ": "+        ++ T.unpack (T.intercalate ", " missingColumns)+        ++ " for operation "+        ++ T.unpack callPoint+        ++ formatSuggestions missingColumns availableColumns+        ++ "\n\n"+  where+    missingColumnsLabel [_] = "Column not found"+    missingColumnsLabel _ = "Columns not found"++    formatSuggestions [missingColumn] columns =+        case guessColumnName missingColumn columns of+            "" -> ""+            guessed ->+                "\n\tDid you mean "+                    ++ T.unpack guessed+                    ++ "?"+    formatSuggestions names columns =+        case traverse (`suggestColumnName` columns) names of+            Just guessedColumns+                | not (null guessedColumns) ->+                    "\n\tDid you mean "+                        ++ formatColumnSuggestions guessedColumns+                        ++ "?"+            _ -> ""++    suggestColumnName missingColumn columns = case guessColumnName missingColumn columns of+        "" -> Nothing+        guessed -> Just guessed++    formatColumnSuggestions guessedColumns =+        "["+            ++ L.intercalate ", " (map (show . T.unpack) guessedColumns)+            ++ "]"++typeMismatchError :: String -> String -> String+typeMismatchError givenType expType =+    red $+        red "\n\n[Error]: Type Mismatch"+            ++ "\n\tWhile running your code I tried to "+            ++ "get a column of type: "+            ++ red (show givenType)+            ++ " but the column in the dataframe was actually of type: "+            ++ green (show expType)++emptyDataSetError :: T.Text -> String+emptyDataSetError callPoint =+    red "\n\n[ERROR] "+        ++ T.unpack callPoint+        ++ " cannot be called on empty data sets"++wrongQuantileNumberError :: Int -> String+wrongQuantileNumberError q =+    red "\n\n[ERROR] "+        ++ "Quantile number q should satisfy "+        ++ "q >= 2, but here q is "+        ++ show q++wrongQuantileIndexError :: VU.Vector Int -> Int -> String+wrongQuantileIndexError qs q =+    red "\n\n[ERROR] "+        ++ "For quantile number q, "+        ++ "each quantile index i "+        ++ "should satisfy 0 <= i <= q, "+        ++ "but here q is "+        ++ show q+        ++ " and indexes are "+        ++ show qs++addCallPointInfo :: Maybe String -> Maybe String -> String -> String+addCallPointInfo (Just name) (Just cp) err =+    err+        ++ ( "\n\tThis happened when calling function "+                ++ brightGreen cp+                ++ " on "+                ++ brightGreen name+           )+addCallPointInfo Nothing (Just cp) err =+    err+        ++ ( "\n\tThis happened when calling function "+                ++ brightGreen cp+           )+addCallPointInfo (Just name) Nothing err =+    err+        ++ ( "\n\tOn "+                ++ name+                ++ "\n\n"+           )+addCallPointInfo Nothing Nothing err = err++guessColumnName :: T.Text -> [T.Text] -> T.Text+guessColumnName userInput columns = case map (\k -> (editDistance userInput k, k)) columns of+    [] -> ""+    res -> (snd . minimum) res++editDistance :: T.Text -> T.Text -> Int+editDistance xs ys = table ML.! (m, n)+  where+    (m, n) = (T.length xs, T.length ys)+    xv = V.fromList (T.unpack xs)+    yv = V.fromList (T.unpack ys)+    table :: ML.Map (Int, Int) Int+    table = ML.fromList [((i, j), dist i j) | i <- [0 .. m], j <- [0 .. n]]+    dist 0 j = j+    dist i 0 = i+    dist i j =+        minimum+            [ table ML.! (i - 1, j) + 1+            , table ML.! (i, j - 1) + 1+            , (if xv V.! (i - 1) == yv V.! (j - 1) then 0 else 1)+                + table ML.! (i - 1, j - 1)+            ]

@@ -0,0 +1,260 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | Vectorized scatter-accumulate aggregation kernel.+module DataFrame.Internal.AggKernel (+    Reduction (..),+    scatterReduce,+    scatterColumnToDouble,+) where++import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM++import Control.Monad (when)+import Control.Monad.ST (ST, runST)+import DataFrame.Internal.Column (+    Column (..),+    Columnable,+    fromUnboxedVector,+    materializePacked,+ )+import Type.Reflection (typeRep)++{- | A recognised fast-path reduction over a single value column. The element+type (Int vs Double) is resolved at scatter time; sum/min/max preserve the+column's element type, everything else produces a Double column.+-}+data Reduction+    = RSum+    | RCount+    | RMin+    | RMax+    | RMean+    | RStd+    | RVar+    | RTop2Sum+    deriving (Eq, Show)++{- | Coerce an unboxed Int or Double column to an unboxed Double vector for the+moment/mean/sd/median family. Returns 'Nothing' for boxed, nullable, or other+element types (the caller then falls back to the interpreter).+-}+scatterColumnToDouble :: Column -> Maybe (VU.Vector Double)+scatterColumnToDouble = \case+    UnboxedColumn Nothing (v :: VU.Vector a) ->+        case testEquality (typeRep @a) (typeRep @Double) of+            Just Refl -> Just v+            Nothing -> case testEquality (typeRep @a) (typeRep @Int) of+                Just Refl -> Just (VU.map fromIntegral v)+                Nothing -> Nothing+    p@(PackedText _ _) -> scatterColumnToDouble (materializePacked p)+    _ -> Nothing++scatterReduce ::+    Reduction -> VU.Vector Int -> Int -> Column -> Maybe Column+scatterReduce red g nGroups col = case col of+    UnboxedColumn Nothing (v :: VU.Vector a) ->+        case testEquality (typeRep @a) (typeRep @Int) of+            Just Refl -> Just (reduceTyped red g nGroups v intIdent)+            Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+                Just Refl -> Just (reduceTyped red g nGroups v dblIdent)+                Nothing -> Nothing+    p@(PackedText _ _) -> scatterReduce red g nGroups (materializePacked p)+    _ -> Nothing+{-# INLINEABLE scatterReduce #-}++-- | Per-type seed identities for the order-preserving reductions.+data Idents a = Idents {minSeed :: !a, maxSeed :: !a}++intIdent :: Idents Int+intIdent = Idents maxBound minBound++dblIdent :: Idents Double+dblIdent = Idents (1 / 0) (negate (1 / 0))++reduceTyped ::+    forall a.+    (Columnable a, VU.Unbox a, Num a, Ord a, Real a) =>+    Reduction -> VU.Vector Int -> Int -> VU.Vector a -> Idents a -> Column+reduceTyped red g nGroups v idents = case red of+    RCount -> fromUnboxedVector (countScatter g nGroups)+    RSum -> fromUnboxedVector (sumScatter g nGroups v)+    RMin -> fromUnboxedVector (extremaScatter min (minSeed idents) g nGroups v)+    RMax -> fromUnboxedVector (extremaScatter max (maxSeed idents) g nGroups v)+    RMean -> fromUnboxedVector (meanScatter g nGroups v)+    RVar -> fromUnboxedVector (varScatter False g nGroups v)+    RStd -> fromUnboxedVector (varScatter True g nGroups v)+    RTop2Sum -> fromUnboxedVector (top2Scatter g nGroups v)+{-# INLINE reduceTyped #-}++countScatter :: VU.Vector Int -> Int -> VU.Vector Int+countScatter g nGroups = runST $ do+    cnt <- VUM.replicate nGroups (0 :: Int)+    let n = VU.length g+        go !i+            | i >= n = pure ()+            | otherwise = do+                let !k = VU.unsafeIndex g i+                c <- VUM.unsafeRead cnt k+                VUM.unsafeWrite cnt k (c + 1)+                go (i + 1)+    go 0+    VU.unsafeFreeze cnt++sumScatter ::+    (VU.Unbox a, Num a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector a+sumScatter g nGroups v = runST $ do+    s <- VUM.replicate nGroups 0+    let n = VU.length v+        go !i+            | i >= n = pure ()+            | otherwise = do+                let !k = VU.unsafeIndex g i+                cur <- VUM.unsafeRead s k+                VUM.unsafeWrite s k (cur + VU.unsafeIndex v i)+                go (i + 1)+    go 0+    VU.unsafeFreeze s+{-# INLINE sumScatter #-}++extremaScatter ::+    (VU.Unbox a) =>+    (a -> a -> a) -> a -> VU.Vector Int -> Int -> VU.Vector a -> VU.Vector a+extremaScatter combine seed g nGroups v = runST $ do+    m <- VUM.replicate nGroups seed+    let n = VU.length v+        go !i+            | i >= n = pure ()+            | otherwise = do+                let !k = VU.unsafeIndex g i+                cur <- VUM.unsafeRead m k+                VUM.unsafeWrite m k (combine cur (VU.unsafeIndex v i))+                go (i + 1)+    go 0+    VU.unsafeFreeze m+{-# INLINE extremaScatter #-}++meanScatter ::+    (VU.Unbox a, Real a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double+meanScatter g nGroups v = runST $ do+    s <- VUM.replicate nGroups (0 :: Double)+    cnt <- VUM.replicate nGroups (0 :: Int)+    scatterSumCount g v s cnt+    finalizeMean nGroups s cnt+{-# INLINE meanScatter #-}++scatterSumCount ::+    (VU.Unbox a, Real a) =>+    VU.Vector Int ->+    VU.Vector a ->+    VUM.MVector s Double ->+    VUM.MVector s Int ->+    ST s ()+scatterSumCount g v s cnt = go 0+  where+    n = VU.length v+    go !i+        | i >= n = pure ()+        | otherwise = do+            let !k = VU.unsafeIndex g i+                !x = realToFrac (VU.unsafeIndex v i)+            curS <- VUM.unsafeRead s k+            VUM.unsafeWrite s k (curS + x)+            curC <- VUM.unsafeRead cnt k+            VUM.unsafeWrite cnt k (curC + 1)+            go (i + 1)+{-# INLINE scatterSumCount #-}++finalizeMean ::+    Int -> VUM.MVector s Double -> VUM.MVector s Int -> ST s (VU.Vector Double)+finalizeMean nGroups s cnt = do+    out <- VUM.new nGroups+    let go !k+            | k >= nGroups = pure ()+            | otherwise = do+                sv <- VUM.unsafeRead s k+                c <- VUM.unsafeRead cnt k+                VUM.unsafeWrite out k (if c == 0 then 0 / 0 else sv / fromIntegral c)+                go (k + 1)+    go 0+    VU.unsafeFreeze out++varScatter ::+    (VU.Unbox a, Real a) =>+    Bool -> VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double+varScatter takeSqrt g nGroups v = runST $ do+    cnt <- VUM.replicate nGroups (0 :: Int)+    meanV <- VUM.replicate nGroups (0 :: Double)+    m2 <- VUM.replicate nGroups (0 :: Double)+    let n = VU.length v+        go !i+            | i >= n = pure ()+            | otherwise = do+                let !k = VU.unsafeIndex g i+                    !x = realToFrac (VU.unsafeIndex v i)+                c <- VUM.unsafeRead cnt k+                mu <- VUM.unsafeRead meanV k+                mm <- VUM.unsafeRead m2 k+                let !c' = c + 1+                    !delta = x - mu+                    !mu' = mu + delta / fromIntegral c'+                    !mm' = mm + delta * (x - mu')+                VUM.unsafeWrite cnt k c'+                VUM.unsafeWrite meanV k mu'+                VUM.unsafeWrite m2 k mm'+                go (i + 1)+    go 0+    out <- VUM.new nGroups+    let fin !k+            | k >= nGroups = pure ()+            | otherwise = do+                c <- VUM.unsafeRead cnt k+                mm <- VUM.unsafeRead m2 k+                let var = if c < 2 then 0 else mm / fromIntegral (c - 1)+                VUM.unsafeWrite out k (if takeSqrt then sqrt var else var)+                fin (k + 1)+    fin 0+    VU.unsafeFreeze out+{-# INLINE varScatter #-}++top2Scatter ::+    (VU.Unbox a, Real a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double+top2Scatter g nGroups v = runST $ do+    let ninf = negate (1 / 0) :: Double+    m1 <- VUM.replicate nGroups ninf+    m2 <- VUM.replicate nGroups ninf+    let n = VU.length v+        go !i+            | i >= n = pure ()+            | otherwise = do+                let !k = VU.unsafeIndex g i+                    !x = realToFrac (VU.unsafeIndex v i)+                a1 <- VUM.unsafeRead m1 k+                if x > a1+                    then do+                        VUM.unsafeWrite m1 k x+                        VUM.unsafeWrite m2 k a1+                    else do+                        a2 <- VUM.unsafeRead m2 k+                        when (x > a2) (VUM.unsafeWrite m2 k x)+                go (i + 1)+    go 0+    out <- VUM.new nGroups+    let fin !k+            | k >= nGroups = pure ()+            | otherwise = do+                a1 <- VUM.unsafeRead m1 k+                a2 <- VUM.unsafeRead m2 k+                let s = (if isInfinite a1 then 0 else a1) + (if isInfinite a2 then 0 else a2)+                VUM.unsafeWrite out k s+                fin (k + 1)+    fin 0+    VU.unsafeFreeze out+{-# INLINE top2Scatter #-}

@@ -0,0 +1,338 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module DataFrame.Internal.AggKernelDirect (+    directThreshold,+    directReduce,+) where++import Control.Concurrent (forkIO, getNumCapabilities)+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)+import Control.Exception (SomeException, throwIO, try)+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import System.IO.Unsafe (unsafePerformIO)+import Type.Reflection (typeRep)++import DataFrame.Internal.AggKernel (Reduction (..))+import DataFrame.Internal.Column (+    Column (..),+    fromUnboxedVector,+    materializePacked,+ )++{- | Group-domain size at or below which the direct-indexed accumulator path is+taken; wider domains keep the group-range kernel. The admitted reductions are+order-independent, so the per-worker accumulator merge is exact.+-}+directThreshold :: Int+directThreshold = 262144++capabilities :: Int+capabilities = unsafePerformIO getNumCapabilities+{-# NOINLINE capabilities #-}++{- | Below this many rows the parallel fan-out is not worth it; a single+sequential direct pass runs instead (tiny accumulator, one tight loop). Matches+the grouping/scatter parallel threshold.+-}+parThreshold :: Int+parThreshold = 200000++{- | Run a recognised reduction through the direct-indexed path. 'Nothing' (so+the caller falls back to the order-preserving kernel) unless the reduction is+order-independent at this element type AND the column is a clean unboxed Int/Double.+-}+directReduce :: Reduction -> VU.Vector Int -> Int -> Column -> Maybe Column+directReduce red g nGroups col = case col of+    UnboxedColumn Nothing (v :: VU.Vector a) ->+        case testEquality (typeRep @a) (typeRep @Int) of+            Just Refl -> directInt red g nGroups v+            Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+                Just Refl -> directDouble red g nGroups v+                Nothing -> Nothing+    p@(PackedText _ _) -> directReduce red g nGroups (materializePacked p)+    _ -> Nothing+{-# INLINEABLE directReduce #-}++-- | The order-independent reductions over an Int column.+directInt :: Reduction -> VU.Vector Int -> Int -> VU.Vector Int -> Maybe Column+directInt red g nGroups v = case red of+    RCount -> Just (fromUnboxedVector (countDirect g nGroups (VU.length v)))+    RSum -> Just (fromUnboxedVector (sumIntDirect g nGroups v))+    RMin -> Just (fromUnboxedVector (extremaIntDirect True g nGroups v))+    RMax -> Just (fromUnboxedVector (extremaIntDirect False g nGroups v))+    RMean -> Just (fromUnboxedVector (meanIntDirect g nGroups v))+    _ -> Nothing++{- | Over a Double column only @count@ is order-independent; the float+sum/mean/variance reductions must keep the order-preserving kernel.+-}+directDouble ::+    Reduction -> VU.Vector Int -> Int -> VU.Vector Double -> Maybe Column+directDouble red g nGroups v = case red of+    RCount -> Just (fromUnboxedVector (countDirect g nGroups (VU.length v)))+    _ -> Nothing++-- | Whether to fan out at this row count.+shouldPar :: Int -> Bool+shouldPar n = n >= parThreshold && capabilities > 1++{- | Fork @caps@ workers over disjoint contiguous row ranges of @[0, n)@, each+producing its own private accumulator (no shared array, no sync). Returns the+partials in worker order for the caller's merge; rethrows the first failure.+-}+runPartialsOver ::+    Int -> Int -> (Int -> Int -> IO (VUM.IOVector Int)) -> IO [VUM.IOVector Int]+runPartialsOver n caps fill = do+    let !per = (n + caps - 1) `div` caps+        spawn w = do+            var <- newEmptyMVar+            let !lo = min n (w * per)+                !hi = min n (lo + per)+            _ <- forkIO (try (fill lo hi) >>= putMVar var)+            pure var+    vars <- mapM spawn [0 .. caps - 1]+    results <- mapM takeMVar vars+    mapM (either (throwIO @SomeException) pure) results++{- | As 'runPartialsOver' but each worker produces a PAIR of accumulators (e.g.+sum and count for the fused integer mean).+-}+runPartialsPairOver ::+    Int ->+    Int ->+    (Int -> Int -> IO (VUM.IOVector Int, VUM.IOVector Int)) ->+    IO [(VUM.IOVector Int, VUM.IOVector Int)]+runPartialsPairOver n caps fill = do+    let !per = (n + caps - 1) `div` caps+        spawn w = do+            var <- newEmptyMVar+            let !lo = min n (w * per)+                !hi = min n (lo + per)+            _ <- forkIO (try (fill lo hi) >>= putMVar var)+            pure var+    vars <- mapM spawn [0 .. caps - 1]+    results <- mapM takeMVar vars+    mapM (either (throwIO @SomeException) pure) results++-------------------------------------------------------------------------------+-- Count (order-independent: per-group row count)+-------------------------------------------------------------------------------++countDirect :: VU.Vector Int -> Int -> Int -> VU.Vector Int+countDirect g nGroups n+    | not (shouldPar n) =+        unsafePerformIO (countChunk g nGroups 0 n >>= VU.unsafeFreeze)+    | otherwise = unsafePerformIO $ do+        parts <- runPartialsOver n capabilities (countChunk g nGroups)+        mergeIntSum nGroups parts+{-# NOINLINE countDirect #-}++countChunk :: VU.Vector Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)+countChunk g nGroups lo hi = do+    acc <- VUM.replicate nGroups (0 :: Int)+    let go !i+            | i >= hi = pure ()+            | otherwise = do+                let !k = VU.unsafeIndex g i+                c <- VUM.unsafeRead acc k+                VUM.unsafeWrite acc k (c + 1)+                go (i + 1)+    go lo+    pure acc++-------------------------------------------------------------------------------+-- Integer sum (exact: merge order irrelevant)+-------------------------------------------------------------------------------++sumIntDirect :: VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int+sumIntDirect g nGroups v+    | not (shouldPar n) =+        unsafePerformIO (sumIntChunk g v nGroups 0 n >>= VU.unsafeFreeze)+    | otherwise = unsafePerformIO $ do+        parts <- runPartialsOver n capabilities (sumIntChunk g v nGroups)+        mergeIntSum nGroups parts+  where+    !n = VU.length v+{-# NOINLINE sumIntDirect #-}++sumIntChunk ::+    VU.Vector Int -> VU.Vector Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)+sumIntChunk g v nGroups lo hi = do+    acc <- VUM.replicate nGroups (0 :: Int)+    let go !i+            | i >= hi = pure ()+            | otherwise = do+                let !k = VU.unsafeIndex g i+                c <- VUM.unsafeRead acc k+                VUM.unsafeWrite acc k (c + VU.unsafeIndex v i)+                go (i + 1)+    go lo+    pure acc++-------------------------------------------------------------------------------+-- Integer min / max (order-independent)+-------------------------------------------------------------------------------++extremaIntDirect ::+    Bool -> VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int+extremaIntDirect isMin g nGroups v+    | not (shouldPar n) =+        unsafePerformIO (extremaIntChunk isMin g v nGroups 0 n >>= VU.unsafeFreeze)+    | otherwise = unsafePerformIO $ do+        parts <- runPartialsOver n capabilities (extremaIntChunk isMin g v nGroups)+        mergeExtremaInt isMin nGroups parts+  where+    !n = VU.length v+{-# NOINLINE extremaIntDirect #-}++extremaIntChunk ::+    Bool ->+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    Int ->+    Int ->+    IO (VUM.IOVector Int)+extremaIntChunk isMin g v nGroups lo hi = do+    let !seed = if isMin then maxBound else minBound+        combine a b = if isMin then min a b else max a b+    acc <- VUM.replicate nGroups seed+    let go !i+            | i >= hi = pure ()+            | otherwise = do+                let !k = VU.unsafeIndex g i+                c <- VUM.unsafeRead acc k+                VUM.unsafeWrite acc k (combine c (VU.unsafeIndex v i))+                go (i + 1)+    go lo+    pure acc++-------------------------------------------------------------------------------+-- Integer mean (exact integer sum + count, divided once -> order-independent)+-------------------------------------------------------------------------------++{- | Integer mean in ONE fused pass: a running integer sum and count per group,+divided once at finalize. The integer sum is exact, so the parallel partial+merge is byte-identical to the sequential single pass at any @-N@.+-}+meanIntDirect :: VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Double+meanIntDirect g nGroups v+    | not (shouldPar n) = unsafePerformIO $ do+        (s, c) <- meanIntChunk g v nGroups 0 n+        finalizeMeanInt nGroups s c+    | otherwise = unsafePerformIO $ do+        parts <- runPartialsPairOver n capabilities (meanIntChunk g v nGroups)+        (s, c) <- mergePair nGroups parts+        finalizeMeanInt nGroups s c+  where+    !n = VU.length v+{-# NOINLINE meanIntDirect #-}++meanIntChunk ::+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    Int ->+    Int ->+    IO (VUM.IOVector Int, VUM.IOVector Int)+meanIntChunk g v nGroups lo hi = do+    s <- VUM.replicate nGroups (0 :: Int)+    c <- VUM.replicate nGroups (0 :: Int)+    let go !i+            | i >= hi = pure ()+            | otherwise = do+                let !k = VU.unsafeIndex g i+                sv <- VUM.unsafeRead s k+                VUM.unsafeWrite s k (sv + VU.unsafeIndex v i)+                cv <- VUM.unsafeRead c k+                VUM.unsafeWrite c k (cv + 1)+                go (i + 1)+    go lo+    pure (s, c)++finalizeMeanInt ::+    Int -> VUM.IOVector Int -> VUM.IOVector Int -> IO (VU.Vector Double)+finalizeMeanInt nGroups s c = do+    out <- VUM.new nGroups+    let go !k+            | k >= nGroups = pure ()+            | otherwise = do+                sv <- VUM.unsafeRead s k+                cv <- VUM.unsafeRead c k+                VUM.unsafeWrite+                    out+                    k+                    (if cv == 0 then 0 / 0 else fromIntegral sv / fromIntegral cv)+                go (k + 1)+    go 0+    VU.unsafeFreeze out++-------------------------------------------------------------------------------+-- Partial accumulation + merge+-------------------------------------------------------------------------------++mergeIntSum :: Int -> [VUM.IOVector Int] -> IO (VU.Vector Int)+mergeIntSum nGroups parts = case parts of+    [] -> VU.unsafeFreeze =<< VUM.replicate nGroups 0+    (p0 : rest) -> do+        let add !p = do+                let go !k+                        | k >= nGroups = pure ()+                        | otherwise = do+                            a <- VUM.unsafeRead p0 k+                            b <- VUM.unsafeRead p k+                            VUM.unsafeWrite p0 k (a + b)+                            go (k + 1)+                go 0+        mapM_ add rest+        VU.unsafeFreeze p0++{- | Merge per-worker (sum, count) partials into the first worker's pair by+exact integer addition; returns the accumulated pair for finalize.+-}+mergePair ::+    Int ->+    [(VUM.IOVector Int, VUM.IOVector Int)] ->+    IO (VUM.IOVector Int, VUM.IOVector Int)+mergePair nGroups parts = case parts of+    [] -> (,) <$> VUM.replicate nGroups 0 <*> VUM.replicate nGroups 0+    ((s0, c0) : rest) -> do+        let add (s, c) = do+                let go !k+                        | k >= nGroups = pure ()+                        | otherwise = do+                            sa <- VUM.unsafeRead s0 k+                            sb <- VUM.unsafeRead s k+                            VUM.unsafeWrite s0 k (sa + sb)+                            ca <- VUM.unsafeRead c0 k+                            cb <- VUM.unsafeRead c k+                            VUM.unsafeWrite c0 k (ca + cb)+                            go (k + 1)+                go 0+        mapM_ add rest+        pure (s0, c0)++mergeExtremaInt :: Bool -> Int -> [VUM.IOVector Int] -> IO (VU.Vector Int)+mergeExtremaInt isMin nGroups parts = case parts of+    [] ->+        VU.unsafeFreeze =<< VUM.replicate nGroups (if isMin then maxBound else minBound)+    (p0 : rest) -> do+        let combine a b = if isMin then min a b else max a b+            add !p = do+                let go !k+                        | k >= nGroups = pure ()+                        | otherwise = do+                            a <- VUM.unsafeRead p0 k+                            b <- VUM.unsafeRead p k+                            VUM.unsafeWrite p0 k (combine a b)+                            go (k + 1)+                go 0+        mapM_ add rest+        VU.unsafeFreeze p0

@@ -0,0 +1,391 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | Parallel scatter-accumulate aggregation kernel.+module DataFrame.Internal.AggKernelPar (+    scatterReducePar,+    momentScatterPar,+) where++import Control.Concurrent (forkIO, getNumCapabilities)+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)+import Control.Exception (SomeException, throwIO, try)+import Control.Monad (when)+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import System.IO.Unsafe (unsafePerformIO)+import Type.Reflection (typeRep)++import DataFrame.Internal.AggKernel (+    Reduction (..),+    scatterColumnToDouble,+    scatterReduce,+ )+import DataFrame.Internal.AggPlan (Moments (..), momentScatter)+import DataFrame.Internal.Column (+    Column (..),+    Columnable,+    fromUnboxedVector,+    materializePacked,+ )++parThreshold :: Int+parThreshold = 200000++capabilities :: Int+capabilities = unsafePerformIO getNumCapabilities+{-# NOINLINE capabilities #-}++-- | Whether to take the parallel path at this row count.+shouldPar :: Int -> Bool+shouldPar n = n >= parThreshold && capabilities > 1++groupRangeBounds :: VU.Vector Int -> Int -> Int -> VU.Vector Int+groupRangeBounds offs nGroups caps = VU.create $ do+    b <- VUM.new (caps + 1)+    let !nRows = VU.unsafeIndex offs nGroups+        !per = max 1 ((nRows + caps - 1) `div` caps)+        adv !target !gg+            | gg >= nGroups = nGroups+            | VU.unsafeIndex offs gg >= target = gg+            | otherwise = adv target (gg + 1)+        go !w !prev+            | w >= caps = VUM.unsafeWrite b caps nGroups+            | otherwise = do+                let !target = min nRows (w * per)+                    !g = adv target prev+                VUM.unsafeWrite b w g+                go (w + 1) g+    VUM.unsafeWrite b 0 0+    go 1 0+    pure b++forEachRange :: VU.Vector Int -> Int -> (Int -> Int -> IO ()) -> IO ()+forEachRange bounds caps act+    | caps <= 1 = act (VU.unsafeIndex bounds 0) (VU.unsafeIndex bounds caps)+    | otherwise = do+        vars <- mapM spawn [0 .. caps - 1]+        results <- mapM takeMVar vars+        mapM_ (either (throwIO :: SomeException -> IO ()) pure) results+  where+    spawn w = do+        var <- newEmptyMVar+        let !s = VU.unsafeIndex bounds w+            !e = VU.unsafeIndex bounds (w + 1)+        _ <- forkIO (try (act s e) >>= putMVar var)+        pure var++scatterReducePar ::+    Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> Maybe Column+scatterReducePar red vis offs nGroups col+    | not (shouldPar (VU.length vis)) || nGroups <= 1 =+        scatterReduce red (rtgFromVis vis offs nGroups) nGroups col+    | otherwise = case col of+        UnboxedColumn Nothing (v :: VU.Vector a) ->+            case testEquality (typeRep @a) (typeRep @Int) of+                Just Refl -> Just (reduceParTyped red vis offs nGroups v intIdent)+                Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+                    Just Refl -> Just (reduceParTyped red vis offs nGroups v dblIdent)+                    Nothing -> Nothing+        p@(PackedText _ _) -> scatterReducePar red vis offs nGroups (materializePacked p)+        _ -> Nothing+{-# NOINLINE scatterReducePar #-}++rtgFromVis :: VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Int+rtgFromVis vis offs nGroups = VU.create $ do+    let n = VU.length vis+    rtg <- VUM.new (max 1 n)+    let go !g+            | g >= nGroups = pure ()+            | otherwise = do+                let !e = VU.unsafeIndex offs (g + 1)+                    inner !pos+                        | pos >= e = pure ()+                        | otherwise = do+                            VUM.unsafeWrite rtg (VU.unsafeIndex vis pos) g+                            inner (pos + 1)+                inner (VU.unsafeIndex offs g)+                go (g + 1)+    go 0+    pure rtg++data Idents a = Idents {minSeed :: !a, maxSeed :: !a}++intIdent :: Idents Int+intIdent = Idents maxBound minBound++dblIdent :: Idents Double+dblIdent = Idents (1 / 0) (negate (1 / 0))++reduceParTyped ::+    forall a.+    (Columnable a, VU.Unbox a, Num a, Ord a, Real a) =>+    Reduction ->+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector a ->+    Idents a ->+    Column+reduceParTyped red vis offs nGroups v idents =+    let !caps = capabilities+        !bounds = groupRangeBounds offs nGroups caps+     in case red of+            RCount -> fromUnboxedVector (unsafePerformIO (countPar vis offs nGroups caps bounds))+            RSum -> fromUnboxedVector (unsafePerformIO (sumPar vis offs nGroups v caps bounds))+            RMin ->+                fromUnboxedVector+                    (unsafePerformIO (extremaPar min (minSeed idents) vis offs nGroups v caps bounds))+            RMax ->+                fromUnboxedVector+                    (unsafePerformIO (extremaPar max (maxSeed idents) vis offs nGroups v caps bounds))+            RMean -> fromUnboxedVector (unsafePerformIO (meanPar vis offs nGroups v caps bounds))+            RVar ->+                fromUnboxedVector+                    (unsafePerformIO (varPar False vis offs nGroups v caps bounds))+            RStd ->+                fromUnboxedVector (unsafePerformIO (varPar True vis offs nGroups v caps bounds))+            RTop2Sum -> fromUnboxedVector (unsafePerformIO (top2Par vis offs nGroups v caps bounds))+{-# INLINE reduceParTyped #-}++-- | Iterate the rows of groups @[gs, ge)@ in @valueIndices@/group order.+overGroups ::+    VU.Vector Int -> VU.Vector Int -> Int -> Int -> (Int -> Int -> IO ()) -> IO ()+overGroups vis offs gs ge step = grp gs+  where+    grp !g+        | g >= ge = pure ()+        | otherwise = do+            let !e = VU.unsafeIndex offs (g + 1)+                inner !pos+                    | pos >= e = pure ()+                    | otherwise = step g (VU.unsafeIndex vis pos) >> inner (pos + 1)+            inner (VU.unsafeIndex offs g)+            grp (g + 1)+{-# INLINE overGroups #-}++countPar ::+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    Int ->+    VU.Vector Int ->+    IO (VU.Vector Int)+countPar _vis offs nGroups caps bounds = do+    out <- VUM.replicate nGroups (0 :: Int)+    forEachRange bounds caps $ \gs ge ->+        let grp !g+                | g >= ge = pure ()+                | otherwise = do+                    let !c = VU.unsafeIndex offs (g + 1) - VU.unsafeIndex offs g+                    VUM.unsafeWrite out g c+                    grp (g + 1)+         in grp gs+    VU.unsafeFreeze out++sumPar ::+    (VU.Unbox a, Num a) =>+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector a ->+    Int ->+    VU.Vector Int ->+    IO (VU.Vector a)+sumPar vis offs nGroups v caps bounds = do+    out <- VUM.replicate nGroups 0+    forEachRange bounds caps $ \gs ge ->+        overGroups vis offs gs ge $ \g row -> do+            cur <- VUM.unsafeRead out g+            VUM.unsafeWrite out g (cur + VU.unsafeIndex v row)+    VU.unsafeFreeze out+{-# INLINE sumPar #-}++extremaPar ::+    (VU.Unbox a) =>+    (a -> a -> a) ->+    a ->+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector a ->+    Int ->+    VU.Vector Int ->+    IO (VU.Vector a)+extremaPar combine seed vis offs nGroups v caps bounds = do+    out <- VUM.replicate nGroups seed+    forEachRange bounds caps $ \gs ge ->+        overGroups vis offs gs ge $ \g row -> do+            cur <- VUM.unsafeRead out g+            VUM.unsafeWrite out g (combine cur (VU.unsafeIndex v row))+    VU.unsafeFreeze out+{-# INLINE extremaPar #-}++meanPar ::+    (VU.Unbox a, Real a) =>+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector a ->+    Int ->+    VU.Vector Int ->+    IO (VU.Vector Double)+meanPar vis offs nGroups v caps bounds = do+    s <- VUM.replicate nGroups (0 :: Double)+    cnt <- VUM.replicate nGroups (0 :: Int)+    forEachRange bounds caps $ \gs ge ->+        overGroups vis offs gs ge $ \g row -> do+            let !x = realToFrac (VU.unsafeIndex v row)+            cs <- VUM.unsafeRead s g+            VUM.unsafeWrite s g (cs + x)+            cc <- VUM.unsafeRead cnt g+            VUM.unsafeWrite cnt g (cc + 1)+    out <- VUM.new nGroups+    let fin !k+            | k >= nGroups = pure ()+            | otherwise = do+                sv <- VUM.unsafeRead s k+                c <- VUM.unsafeRead cnt k+                VUM.unsafeWrite out k (if c == 0 then 0 / 0 else sv / fromIntegral c)+                fin (k + 1)+    fin 0+    VU.unsafeFreeze out+{-# INLINE meanPar #-}++varPar ::+    (VU.Unbox a, Real a) =>+    Bool ->+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector a ->+    Int ->+    VU.Vector Int ->+    IO (VU.Vector Double)+varPar takeSqrt vis offs nGroups v caps bounds = do+    cnt <- VUM.replicate nGroups (0 :: Int)+    meanV <- VUM.replicate nGroups (0 :: Double)+    m2 <- VUM.replicate nGroups (0 :: Double)+    forEachRange bounds caps $ \gs ge ->+        overGroups vis offs gs ge $ \g row -> do+            let !x = realToFrac (VU.unsafeIndex v row)+            c <- VUM.unsafeRead cnt g+            mu <- VUM.unsafeRead meanV g+            mm <- VUM.unsafeRead m2 g+            let !c' = c + 1+                !delta = x - mu+                !mu' = mu + delta / fromIntegral c'+                !mm' = mm + delta * (x - mu')+            VUM.unsafeWrite cnt g c'+            VUM.unsafeWrite meanV g mu'+            VUM.unsafeWrite m2 g mm'+    out <- VUM.new nGroups+    let fin !k+            | k >= nGroups = pure ()+            | otherwise = do+                c <- VUM.unsafeRead cnt k+                mm <- VUM.unsafeRead m2 k+                let var = if c < 2 then 0 else mm / fromIntegral (c - 1)+                VUM.unsafeWrite out k (if takeSqrt then sqrt var else var)+                fin (k + 1)+    fin 0+    VU.unsafeFreeze out+{-# INLINE varPar #-}++top2Par ::+    (VU.Unbox a, Real a) =>+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector a ->+    Int ->+    VU.Vector Int ->+    IO (VU.Vector Double)+top2Par vis offs nGroups v caps bounds = do+    let ninf = negate (1 / 0) :: Double+    m1 <- VUM.replicate nGroups ninf+    m2 <- VUM.replicate nGroups ninf+    forEachRange bounds caps $ \gs ge ->+        overGroups vis offs gs ge $ \g row -> do+            let !x = realToFrac (VU.unsafeIndex v row)+            a1 <- VUM.unsafeRead m1 g+            if x > a1+                then do+                    VUM.unsafeWrite m1 g x+                    VUM.unsafeWrite m2 g a1+                else do+                    a2 <- VUM.unsafeRead m2 g+                    when (x > a2) (VUM.unsafeWrite m2 g x)+    out <- VUM.new nGroups+    let fin !k+            | k >= nGroups = pure ()+            | otherwise = do+                a1 <- VUM.unsafeRead m1 k+                a2 <- VUM.unsafeRead m2 k+                let sm = (if isInfinite a1 then 0 else a1) + (if isInfinite a2 then 0 else a2)+                VUM.unsafeWrite out k sm+                fin (k + 1)+    fin 0+    VU.unsafeFreeze out+{-# INLINE top2Par #-}++-------------------------------------------------------------------------------+-- Parallel fused two-column moments (Q9)+-------------------------------------------------------------------------------++{- | Parallel counterpart of 'momentScatter': one fused pass over both columns,+each group's six sums accumulated within one worker's range. Byte-identical to+'momentScatter'. 'Nothing' unless both columns are non-null unboxed Int/Double.+-}+momentScatterPar ::+    VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -> Maybe Moments+momentScatterPar vis offs nGroups colX colY+    | not (shouldPar (VU.length vis)) || nGroups <= 1 =+        momentScatter (rtgFromVis vis offs nGroups) nGroups colX colY+    | otherwise = do+        xs <- scatterColumnToDouble colX+        ys <- scatterColumnToDouble colY+        let !caps = capabilities+            !bounds = groupRangeBounds offs nGroups caps+        pure (unsafePerformIO (momentPar vis offs nGroups xs ys caps bounds))+{-# NOINLINE momentScatterPar #-}++momentPar ::+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector Double ->+    VU.Vector Double ->+    Int ->+    VU.Vector Int ->+    IO Moments+momentPar vis offs nGroups xs ys caps bounds = do+    cnt <- VUM.replicate nGroups (0 :: Int)+    sx <- VUM.replicate nGroups (0 :: Double)+    sy <- VUM.replicate nGroups (0 :: Double)+    sxx <- VUM.replicate nGroups (0 :: Double)+    syy <- VUM.replicate nGroups (0 :: Double)+    sxy <- VUM.replicate nGroups (0 :: Double)+    let bump arr g d = VUM.unsafeRead arr g >>= \c -> VUM.unsafeWrite arr g (c + d)+    forEachRange bounds caps $ \gs ge ->+        overGroups vis offs gs ge $ \g row -> do+            let !x = VU.unsafeIndex xs row+                !y = VU.unsafeIndex ys row+            VUM.unsafeRead cnt g >>= \c -> VUM.unsafeWrite cnt g (c + 1)+            bump sx g x+            bump sy g y+            bump sxx g (x * x)+            bump syy g (y * y)+            bump sxy g (x * y)+    Moments . fromUnboxedVector+        <$> VU.unsafeFreeze cnt+        <*> (fromUnboxedVector <$> VU.unsafeFreeze sx)+        <*> (fromUnboxedVector <$> VU.unsafeFreeze sy)+        <*> (fromUnboxedVector <$> VU.unsafeFreeze sxx)+        <*> (fromUnboxedVector <$> VU.unsafeFreeze syy)+        <*> (fromUnboxedVector <$> VU.unsafeFreeze sxy)

@@ -0,0 +1,299 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | The aggregation fast-path planner and the two-column moment scatter.+'planAgg' recognises a supported aggregate shape over a clean unboxed Int/Double+column and returns an 'AggPlan'; 'momentScatter' fuses the six regression sums.+-}+module DataFrame.Internal.AggPlan (+    AggPlan (..),+    planAgg,+    Moments (..),+    momentScatter,+    MomentPlan (..),+    planMoments,+) where++import qualified Data.Map.Strict as M+import qualified Data.Text as T+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM++import Control.Monad.ST (runST)+import DataFrame.Internal.AggKernel (Reduction (..), scatterColumnToDouble)+import DataFrame.Internal.Column (Column (..), fromUnboxedVector)+import DataFrame.Internal.DataFrame (+    DataFrame (derivingExpressions),+    GroupedDataFrame (..),+    getColumn,+ )+import DataFrame.Internal.Expression (+    AggStrategy (..),+    BinaryOp (binaryCommutative, binaryName),+    Expr (..),+    UExpr (..),+ )+import Type.Reflection (Typeable, typeRep)++{- | The plan 'planAgg' produces for a recognised output expression. The median+plan carries only the column name (the holistic grouped sort lives in the+operations layer, where @vector-algorithms@ is available).+-}+data AggPlan+    = -- | A single scatter reduction over one named column.+      PlanScatter Reduction T.Text+    | -- | @max a - min b@ (Q7): two scatters then a vectorized combine.+      PlanMaxMinusMin T.Text T.Text+    | -- | Holistic median over one named column.+      PlanMedian T.Text++{- | Inspect a named output expression; return @Just plan@ on a recognised shape+over a present clean column, else 'Nothing'. Nullable or non-Int/Double columns+are rejected here so the scatter only sees a clean unboxed vector.+-}+planAgg :: GroupedDataFrame -> UExpr -> Maybe AggPlan+planAgg gdf (UExpr (expr :: Expr a)) = case expr of+    Agg (FoldAgg tag _ _) (Col name) -> foldPlan tag name+    Agg (MergeAgg tag _ _ _ _) (Col name) -> mergePlan tag name+    Agg (CollectAgg tag _) (Col name) -> collectPlan tag name+    Binary+        op+        (Agg (FoldAgg lt Nothing _) (Col a))+        (Agg (FoldAgg rt Nothing _) (Col b)) ->+            if binaryName op == "sub" && lt == "maximum" && rt == "minimum"+                then requireBoth a b (PlanMaxMinusMin a b)+                else Nothing+    _ -> Nothing+  where+    foldPlan tag name = case tag of+        "sum" -> require name (PlanScatter RSum name)+        "minimum" -> require name (PlanScatter RMin name)+        "maximum" -> require name (PlanScatter RMax name)+        _ -> Nothing+    mergePlan tag name = case tag of+        "mean" -> outputType @Double >> require name (PlanScatter RMean name)+        "count" -> outputType @Int >> require name (PlanScatter RCount name)+        _ -> Nothing+    outputType :: forall t. (Typeable t) => Maybe ()+    outputType = case testEquality (typeRep @a) (typeRep @t) of+        Just Refl -> Just ()+        Nothing -> Nothing+    collectPlan tag name = case tag of+        "stddev" -> require name (PlanScatter RStd name)+        "variance" -> require name (PlanScatter RVar name)+        "top2Sum" -> require name (PlanScatter RTop2Sum name)+        "median" -> require name (PlanMedian name)+        _ -> Nothing+    require name plan = colUnboxedNumeric name >> Just plan+    requireBoth a b plan = colUnboxedNumeric a >> colUnboxedNumeric b >> Just plan+    colUnboxedNumeric name = case getColumn name (fullDataframe gdf) of+        Just c | isUnboxedNumeric c -> Just ()+        _ -> Nothing++-- | The matcher only fires on non-null unboxed Int/Double columns.+isUnboxedNumeric :: Column -> Bool+isUnboxedNumeric = \case+    UnboxedColumn Nothing (_ :: VU.Vector a) ->+        case testEquality (typeRep @a) (typeRep @Int) of+            Just Refl -> True+            Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+                Just Refl -> True+                Nothing -> False+    _ -> False++{- | A recognised moment (Q9 regression) aggregate group: six output columns that+form the sufficient statistics of two base columns @x@ and @y@. The caller runs+'momentScatter' once and binds each output name to a field of the result.+-}+data MomentPlan = MomentPlan+    { mpColX :: T.Text+    , mpColY :: T.Text+    , mpNName :: T.Text+    , mpSxName :: T.Text+    , mpSyName :: T.Text+    , mpSxxName :: T.Text+    , mpSyyName :: T.Text+    , mpSxyName :: T.Text+    }++{- | The shape of a sum's argument once unary coercions are peeled and derived+columns are resolved through @derivingExpressions@: either linear in one base+column or the product of two base columns (sorted).+-}+data Term+    = Lin T.Text+    | Prod T.Text T.Text+    deriving (Eq, Ord, Show)++{- | Recognise the moment shape across a whole @aggregate@ list: exactly+@count@, @sum(x)@, @sum(y)@, @sum(x*x)@, @sum(y*y)@, @sum(x*y)@ over two distinct+clean unboxed base columns. 'Nothing' on any other set.+-}+planMoments :: GroupedDataFrame -> [(T.Text, UExpr)] -> Maybe MomentPlan+planMoments gdf aggs+    | length aggs /= 6 = Nothing+    | otherwise = do+        let exprs = derivingExpressions (fullDataframe gdf)+        roles <- traverse (classify exprs) aggs+        let names = M.fromList [(r, nm) | (nm, r) <- roles]+        nName <- M.lookup RoleN names+        (x, y) <- pickBaseColumns roles+        sxName <- M.lookup (RoleLin x) names+        syName <- M.lookup (RoleLin y) names+        sxxName <- M.lookup (RoleProd x x) names+        syyName <- M.lookup (RoleProd y y) names+        sxyName <- M.lookup (RoleProd x y) names+        _ <- if x /= y then Just () else Nothing+        _ <- colUnboxedNumeric x+        _ <- colUnboxedNumeric y+        pure+            MomentPlan+                { mpColX = x+                , mpColY = y+                , mpNName = nName+                , mpSxName = sxName+                , mpSyName = syName+                , mpSxxName = sxxName+                , mpSyyName = syyName+                , mpSxyName = sxyName+                }+  where+    colUnboxedNumeric name = case getColumn name (fullDataframe gdf) of+        Just c | isUnboxedNumeric c -> Just ()+        _ -> Nothing++-- | The output role each named aggregation plays in the moment shape.+data Role+    = RoleN+    | RoleLin T.Text+    | RoleProd T.Text T.Text+    deriving (Eq, Ord, Show)++-- | Tag a single named aggregation with its moment role, or reject the group.+classify :: M.Map T.Text UExpr -> (T.Text, UExpr) -> Maybe (T.Text, Role)+classify exprs (name, UExpr expr) = case expr of+    Agg (MergeAgg "count" _ _ _ _) _ -> Just (name, RoleN)+    Agg (FoldAgg "sum" _ _) arg -> (\t -> (name, termRole t)) <$> resolveTerm exprs (UExpr arg)+    _ -> Nothing++termRole :: Term -> Role+termRole (Lin a) = RoleLin a+termRole (Prod a b) = RoleProd a b++{- | Resolve a (sum-argument) expression to its 'Term'. Peels @toDouble@-style+unary coercions, follows a derived column to its stored expression, and+recognises a commutative product of two linear terms.+-}+resolveTerm :: M.Map T.Text UExpr -> UExpr -> Maybe Term+resolveTerm exprs = go (8 :: Int)+  where+    go 0 _ = Nothing+    go fuel (UExpr e) = case e of+        Col nm -> case M.lookup nm exprs of+            Just ue -> go (fuel - 1) ue+            Nothing -> Just (Lin nm)+        Unary _ inner -> go (fuel - 1) (UExpr inner)+        Binary op l r+            | binaryName op == "mult" && binaryCommutative op -> do+                Lin a <- go (fuel - 1) (UExpr l)+                Lin b <- go (fuel - 1) (UExpr r)+                Just (sortProd a b)+        _ -> Nothing++-- | Products are unordered: store the pair sorted so @x*y@ and @y*x@ unify.+sortProd :: T.Text -> T.Text -> Term+sortProd a b+    | a <= b = Prod a b+    | otherwise = Prod b a++{- | From the classified roles, find the unordered pair of base columns that the+linear sums name. There must be exactly two distinct linear-sum columns.+-}+pickBaseColumns :: [(T.Text, Role)] -> Maybe (T.Text, T.Text)+pickBaseColumns roles =+    case lins of+        [a, b] | a /= b -> Just (a, b)+        _ -> Nothing+  where+    lins = M.keys (M.fromList [(c, ()) | (_, RoleLin c) <- roles])++{- | The additive moment sums of two columns, each an @nGroups@-length column:+@(n, Sx, Sy, Sxx, Syy, Sxy)@.+-}+data Moments = Moments+    { mN :: Column+    , mSx :: Column+    , mSy :: Column+    , mSxx :: Column+    , mSyy :: Column+    , mSxy :: Column+    }++{- | One pass over two Double-coercible columns @x@ and @y@ filling the count and+five sums, collapsing the Q9 regression family's six folds into a single pass.+'Nothing' unless both columns are non-null unboxed Int/Double.+-}+momentScatter :: VU.Vector Int -> Int -> Column -> Column -> Maybe Moments+momentScatter g nGroups colX colY = do+    xs <- scatterColumnToDouble colX+    ys <- scatterColumnToDouble colY+    let (cnt, sx, sy, sxx, syy, sxy) = momentPass g nGroups xs ys+    pure+        Moments+            { mN = fromUnboxedVector cnt+            , mSx = fromUnboxedVector sx+            , mSy = fromUnboxedVector sy+            , mSxx = fromUnboxedVector sxx+            , mSyy = fromUnboxedVector syy+            , mSxy = fromUnboxedVector sxy+            }++momentPass ::+    VU.Vector Int ->+    Int ->+    VU.Vector Double ->+    VU.Vector Double ->+    ( VU.Vector Int+    , VU.Vector Double+    , VU.Vector Double+    , VU.Vector Double+    , VU.Vector Double+    , VU.Vector Double+    )+momentPass g nGroups xs ys = runST $ do+    cnt <- VUM.replicate nGroups (0 :: Int)+    sx <- VUM.replicate nGroups (0 :: Double)+    sy <- VUM.replicate nGroups (0 :: Double)+    sxx <- VUM.replicate nGroups (0 :: Double)+    syy <- VUM.replicate nGroups (0 :: Double)+    sxy <- VUM.replicate nGroups (0 :: Double)+    let n = VU.length xs+        bump arr k d = VUM.unsafeRead arr k >>= \c -> VUM.unsafeWrite arr k (c + d)+        go !i+            | i >= n = pure ()+            | otherwise = do+                let !k = VU.unsafeIndex g i+                    !x = VU.unsafeIndex xs i+                    !y = VU.unsafeIndex ys i+                VUM.unsafeRead cnt k >>= \c -> VUM.unsafeWrite cnt k (c + 1)+                bump sx k x+                bump sy k y+                bump sxx k (x * x)+                bump syy k (y * y)+                bump sxy k (x * y)+                go (i + 1)+    go 0+    (,,,,,)+        <$> VU.unsafeFreeze cnt+        <*> VU.unsafeFreeze sx+        <*> VU.unsafeFreeze sy+        <*> VU.unsafeFreeze sxx+        <*> VU.unsafeFreeze syy+        <*> VU.unsafeFreeze sxy

@@ -0,0 +1,1744 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}++module DataFrame.Internal.Column where++import qualified Data.Text as T+import qualified Data.Vector as VB+import qualified Data.Vector.Generic as VG+import qualified Data.Vector.Mutable as VBM+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM++import Control.Exception (throw)+import Control.Monad (forM_, when)+import Control.Monad.ST (ST, runST)+import Data.Bits (+    complement,+    popCount,+    setBit,+    shiftL,+    shiftR,+    testBit,+    (.&.),+ )+import Data.Kind (Type)+import Data.Maybe+import Data.Type.Equality (TestEquality (..))+import Data.Word (Word8)+import DataFrame.Errors+import DataFrame.Internal.PackedText (+    PackedTextData (..),+    packedGather,+    packedIndexText,+    packedLength,+    packedRowOffsetVec,+    packedSlice,+    packedTake,+    sliceEqBytes,+ )+import DataFrame.Internal.Types+import DataFrame.Internal.Utf8 (sliceTextVector)+import System.IO.Unsafe (unsafePerformIO)+import System.Random+import Type.Reflection++-- | A bit-packed validity bitmap. Bit @i@ = 1 means row @i@ is valid (not null).+type Bitmap = VU.Vector Word8++{- | Type-erased column GADT. Pattern-matching on the constructor recovers the+representation; nullability is an optional bit-packed 'Bitmap' (@Nothing@ = no+nulls, @Just bm@ = bit @i@ set iff row @i@ is valid).+-}+data Column where+    BoxedColumn :: (Columnable a) => Maybe Bitmap -> VB.Vector a -> Column+    UnboxedColumn ::+        (Columnable a, VU.Unbox a) => Maybe Bitmap -> VU.Vector a -> Column+    -- Bit-packed Text: shared UTF-8 byte buffer + row offsets + optional bitmap;+    -- Text is materialized on demand. Only CSV ingest emits this; user-built+    -- Text columns stay 'BoxedColumn'.+    PackedText :: Maybe Bitmap -> {-# UNPACK #-} !PackedTextData -> Column++{- | A mutable companion struct to dataframe columns.++Used mostly as an intermediate structure for I/O.+-}+data MutableColumn where+    MBoxedColumn :: (Columnable a) => VBM.IOVector a -> MutableColumn+    MUnboxedColumn :: (Columnable a, VU.Unbox a) => VUM.IOVector a -> MutableColumn++-- ---------------------------------------------------------------------------+-- Bitmap helpers+-- ---------------------------------------------------------------------------++-- | Test whether row @i@ is valid (not null) in a bitmap.+bitmapTestBit :: Bitmap -> Int -> Bool+bitmapTestBit bm i = testBit (VU.unsafeIndex bm (i `shiftR` 3)) (i .&. 7)+{-# INLINE bitmapTestBit #-}++-- | Build a fully-valid bitmap for @n@ rows (all bits set).+allValidBitmap :: Int -> Bitmap+allValidBitmap n =+    let bytes = (n + 7) `shiftR` 3+        lastBits = n .&. 7+        full = VU.replicate (bytes - 1) 0xFF+        lastByte = if lastBits == 0 then 0xFF else (1 `shiftL` lastBits) - 1+     in if bytes == 0 then VU.empty else VU.snoc full lastByte+{-# INLINE allValidBitmap #-}++{- | Build a bitmap from a @VU.Vector Word8@ validity vector+(1 = valid, 0 = null), as produced by Arrow / Parquet decoders.+-}+buildBitmapFromValid :: VU.Vector Word8 -> Bitmap+buildBitmapFromValid valid =+    let n = VU.length valid+        bytes = (n + 7) `shiftR` 3+     in VU.generate bytes $ \b ->+            let base = b `shiftL` 3+                setBitIf acc bit =+                    let idx = base + bit+                     in if idx < n && VU.unsafeIndex valid idx /= 0+                            then setBit acc bit+                            else acc+             in foldl setBitIf (0 :: Word8) [0 .. 7]++{- | Build a bitmap from a list of null-row indices.+@nullIdxs@ are the positions that are NULL.+-}+buildBitmapFromNulls :: Int -> [Int] -> Bitmap+buildBitmapFromNulls n nullIdxs =+    let base = allValidBitmap n+     in VU.modify+            ( \mv ->+                forM_ nullIdxs $ \i -> do+                    let byteIdx = i `shiftR` 3+                        bitIdx = i .&. 7+                    v <- VUM.unsafeRead mv byteIdx+                    VUM.unsafeWrite mv byteIdx (clearBit8 v bitIdx)+            )+            base+  where+    clearBit8 :: Word8 -> Int -> Word8+    clearBit8 b bit = b .&. complement (1 `shiftL` bit)++-- | Slice a bitmap for rows @[start .. start+len-1]@.+bitmapSlice :: Int -> Int -> Bitmap -> Bitmap+bitmapSlice start len bm+    | start .&. 7 == 0 =+        let startByte = start `shiftR` 3+            bytes = min ((len + 7) `shiftR` 3) (VU.length bm - startByte)+         in VU.slice startByte bytes bm+    | otherwise =+        let n = min len (VU.length bm `shiftL` 3 - start)+         in buildBitmapFromValid $+                VU.generate n $+                    \i -> if bitmapTestBit bm (start + i) then 1 else 0++-- | Concatenate two bitmaps covering @n1@ and @n2@ rows respectively.+bitmapConcat :: Int -> Bitmap -> Int -> Bitmap -> Bitmap+bitmapConcat n1 bm1 n2 bm2 =+    buildBitmapFromValid $+        VU.generate (n1 + n2) $ \i ->+            if i < n1+                then if bitmapTestBit bm1 i then 1 else 0+                else if bitmapTestBit bm2 (i - n1) then 1 else 0++-- | Combine two bitmaps with AND (both must be valid for result to be valid).+mergeBitmaps :: Bitmap -> Bitmap -> Bitmap+mergeBitmaps = VU.zipWith (.&.)++{- | Materialize a nullable column from @VB.Vector (Maybe a)@; picks 'UnboxedColumn'+when @a@ is unboxable, else 'BoxedColumn'. Always attaches a bitmap so the column+reads as nullable even with no 'Nothing' values.+-}+fromMaybeVec :: forall a. (Columnable a) => VB.Vector (Maybe a) -> Column+fromMaybeVec v = case sUnbox @a of+    STrue -> fromMaybeVecUnboxed v+    SFalse ->+        let n = VB.length v+            nullIdxs = [i | i <- [0 .. n - 1], isNothing (VB.unsafeIndex v i)]+            bm = if null nullIdxs then allValidBitmap n else buildBitmapFromNulls n nullIdxs+            dat = VB.map (fromMaybe (errorWithoutStackTrace "fromMaybeVec: Nothing slot")) v+         in BoxedColumn (Just bm) dat++{- | Materialize a nullable 'UnboxedColumn' to @VB.Vector (Maybe a)@ using runST.+Always attaches a bitmap so the column is recognized as nullable even when+no 'Nothing' values are present (preserves the Maybe type marker).+-}+fromMaybeVecUnboxed ::+    forall a. (Columnable a, VU.Unbox a) => VB.Vector (Maybe a) -> Column+fromMaybeVecUnboxed v =+    let n = VB.length v+        nullIdxs = [i | i <- [0 .. n - 1], isNothing (VB.unsafeIndex v i)]+        bm = if null nullIdxs then allValidBitmap n else buildBitmapFromNulls n nullIdxs+        dat = runST $ do+            mv <- VUM.new n+            VG.iforM_ v $ \i mx -> forM_ mx (VUM.unsafeWrite mv i)+            VU.unsafeFreeze mv+     in UnboxedColumn (Just bm) dat++-- | Whether row @i@ is null, respecting the bitmap.+columnElemIsNull :: Column -> Int -> Bool+columnElemIsNull (BoxedColumn (Just bm) _) i = not (bitmapTestBit bm i)+columnElemIsNull (UnboxedColumn (Just bm) _) i = not (bitmapTestBit bm i)+columnElemIsNull (PackedText (Just bm) _) i = not (bitmapTestBit bm i)+columnElemIsNull _ _ = False++-- | Return the 'Maybe Bitmap' from a column.+columnBitmap :: Column -> Maybe Bitmap+columnBitmap (BoxedColumn bm _) = bm+columnBitmap (UnboxedColumn bm _) = bm+columnBitmap (PackedText bm _) = bm++{- | Decode a 'PackedText' into a @BoxedColumn Text@ (bit-identical to+materializing at freeze). Identity on every other column.+-}+materializePacked :: Column -> Column+materializePacked (PackedText bm p) = case packedRowOffsetVec p of+    Just (arr, offs) -> BoxedColumn bm (sliceTextVector arr offs)+    Nothing -> BoxedColumn bm (VB.generate (packedLength p) (packedIndexText p))+materializePacked c = c+{-# INLINE materializePacked #-}++-- | Whether a column is a 'PackedText'.+isPackedText :: Column -> Bool+isPackedText (PackedText _ _) = True+isPackedText _ = False+{-# INLINE isPackedText #-}++-- ---------------------------------------------------------------------------+-- End bitmap helpers+-- ---------------------------------------------------------------------------++{- | A wrapper around the type-erased 'Column' carrying a phantom element type,+used to type-check expressions. The phantom is not guaranteed to match the+underlying vector's type.+-}+data TypedColumn a where+    TColumn :: (Columnable a) => Column -> TypedColumn a++instance (Eq a) => Eq (TypedColumn a) where+    (==) :: (Eq a) => TypedColumn a -> TypedColumn a -> Bool+    (==) (TColumn a) (TColumn b) = a == b++-- | Gets the underlying value from a TypedColumn.+unwrapTypedColumn :: TypedColumn a -> Column+unwrapTypedColumn (TColumn value) = value++-- | Gets the underlying vector from a TypedColumn.+vectorFromTypedColumn :: TypedColumn a -> VB.Vector a+vectorFromTypedColumn (TColumn value) = either throw id (toVector value)++-- | Checks if a column contains missing values (has a bitmap).+hasMissing :: Column -> Bool+hasMissing (BoxedColumn (Just _) _) = True+hasMissing (UnboxedColumn (Just _) _) = True+hasMissing (PackedText (Just _) _) = True+hasMissing _ = False++-- | Checks if a column contains only missing values.+allMissing :: Column -> Bool+allMissing (BoxedColumn (Just bm) col) = VU.all (== 0) bm && not (VB.null col)+allMissing (UnboxedColumn (Just bm) col) = VU.all (== 0) bm && not (VU.null col)+allMissing (PackedText (Just bm) p) = VU.all (== 0) bm && packedLength p > 0+allMissing _ = False++-- | Checks if a column contains numeric values.+isNumeric :: Column -> Bool+isNumeric (UnboxedColumn _ (_vec :: VU.Vector a)) = case sNumeric @a of+    STrue -> True+    _ -> False+isNumeric (BoxedColumn _ (_vec :: VB.Vector a)) = case testEquality (typeRep @a) (typeRep @Integer) of+    Nothing -> False+    Just Refl -> True+isNumeric (PackedText _ _) = False++{- | Whether the column stores element type @a@. For nullable columns, also+'True' when @a = Maybe b@ and the column stores @b@ internally.+-}+hasElemType :: forall a. (Columnable a) => Column -> Bool+hasElemType = \case+    BoxedColumn bm (_column :: VB.Vector b) -> checkBoxed bm (typeRep @b)+    UnboxedColumn bm (_column :: VU.Vector b) -> checkUnboxed bm (typeRep @b)+    PackedText bm _ -> checkBoxed bm (typeRep @T.Text)+  where+    directMatch :: forall (b :: Type). TypeRep b -> Bool+    directMatch = isJust . testEquality (typeRep @a)+    checkMaybe :: forall (b :: Type). TypeRep b -> Bool+    checkMaybe tb = case typeRep @a of+        App tMaybe tInner -> case eqTypeRep tMaybe (typeRep @Maybe) of+            Just HRefl -> isJust (testEquality tInner tb)+            Nothing -> False+        _ -> False+    checkBoxed :: forall (b :: Type). Maybe Bitmap -> TypeRep b -> Bool+    checkBoxed bm tb = directMatch tb || (isJust bm && checkMaybe tb)+    checkUnboxed :: forall (b :: Type). Maybe Bitmap -> TypeRep b -> Bool+    checkUnboxed bm tb = directMatch tb || (isJust bm && checkMaybe tb)++-- | An internal/debugging function to get the column type of a column.+columnVersionString :: Column -> String+columnVersionString column = case column of+    BoxedColumn Nothing _ -> "Boxed"+    BoxedColumn (Just _) _ -> "NullableBoxed"+    UnboxedColumn Nothing _ -> "Unboxed"+    UnboxedColumn (Just _) _ -> "NullableUnboxed"+    PackedText Nothing _ -> "Boxed"+    PackedText (Just _) _ -> "NullableBoxed"++{- | An internal/debugging function to get the type stored in the outermost vector+of a column.+-}+columnTypeString :: Column -> String+columnTypeString column = case column of+    BoxedColumn Nothing (_ :: VB.Vector a) -> show (typeRep @a)+    BoxedColumn (Just _) (_ :: VB.Vector a) -> showMaybeType @a+    UnboxedColumn Nothing (_ :: VU.Vector a) -> show (typeRep @a)+    UnboxedColumn (Just _) (_ :: VU.Vector a) -> showMaybeType @a+    PackedText Nothing _ -> show (typeRep @T.Text)+    PackedText (Just _) _ -> showMaybeType @T.Text+  where+    showMaybeType :: forall a. (Typeable a) => String+    showMaybeType =+        let s = show (typeRep @a)+         in "Maybe " ++ if ' ' `elem` s then "(" ++ s ++ ")" else s++instance (Show a) => Show (TypedColumn a) where+    show :: (Show a) => TypedColumn a -> String+    show (TColumn col) = show col++{- | Force evaluation of all elements in a column. Replacement for the removed+@instance NFData Column@; used by the IO and lazy-executor strict paths.+-}+forceColumn :: Column -> ()+forceColumn (BoxedColumn Nothing (v :: VB.Vector a)) = VB.foldl' (const (`seq` ())) () v+forceColumn (BoxedColumn (Just bm) (v :: VB.Vector a)) =+    let n = VB.length v+        go !i+            | i >= n = ()+            | bitmapTestBit bm i = VB.unsafeIndex v i `seq` go (i + 1)+            | otherwise = go (i + 1)+     in go 0+forceColumn (UnboxedColumn _ v) = v `seq` ()+forceColumn (PackedText _ (PackedTextData arr offs sel)) = arr `seq` offs `seq` sel `seq` ()++instance Show Column where+    show :: Column -> String+    show (BoxedColumn Nothing column) = show column+    show (BoxedColumn (Just bm) column) =+        let n = VB.length column+            elems =+                [ if bitmapTestBit bm i then show (VB.unsafeIndex column i) else "null"+                | i <- [0 .. n - 1]+                ]+         in "[" ++ foldl (\acc e -> if null acc then e else acc ++ "," ++ e) "" elems ++ "]"+    show (UnboxedColumn Nothing column) = show column+    show (UnboxedColumn (Just bm) column) =+        let n = VU.length column+            elems =+                [ if bitmapTestBit bm i then show (VU.unsafeIndex column i) else "null"+                | i <- [0 .. n - 1]+                ]+         in "[" ++ foldl (\acc e -> if null acc then e else acc ++ "," ++ e) "" elems ++ "]"+    show c@(PackedText _ _) = show (materializePacked c)++{- | Compare two nullable boxed columns element by element, skipping null slots.+Uses a manual loop to avoid stream fusion forcing null-slot error thunks.+-}+eqBoxedCols ::+    (Eq a) => Maybe Bitmap -> VB.Vector a -> Maybe Bitmap -> VB.Vector a -> Bool+eqBoxedCols bm1 a bm2 b+    | VB.length a /= VB.length b = False+    | otherwise = go 0+  where+    !n = VB.length a+    go !i+        | i >= n = True+        | nullA || nullB = (nullA == nullB) && go (i + 1)+        | VB.unsafeIndex a i == VB.unsafeIndex b i = go (i + 1)+        | otherwise = False+      where+        nullA = maybe False (\bm -> not (bitmapTestBit bm i)) bm1+        nullB = maybe False (\bm -> not (bitmapTestBit bm i)) bm2+{-# INLINE eqBoxedCols #-}++instance Eq Column where+    (==) :: Column -> Column -> Bool+    (==) (BoxedColumn bm1 (a :: VB.Vector t1)) (BoxedColumn bm2 (b :: VB.Vector t2)) =+        case testEquality (typeRep @t1) (typeRep @t2) of+            Nothing -> False+            Just Refl -> eqBoxedCols bm1 a bm2 b+    (==) (UnboxedColumn bm1 (a :: VU.Vector t1)) (UnboxedColumn bm2 (b :: VU.Vector t2)) =+        case testEquality (typeRep @t1) (typeRep @t2) of+            Nothing -> False+            Just Refl ->+                VU.length a == VU.length b+                    && VU.and+                        ( VU.imap+                            ( \i x ->+                                let nullA = maybe False (\bm -> not (bitmapTestBit bm i)) bm1+                                    nullB = maybe False (\bm -> not (bitmapTestBit bm i)) bm2+                                 in if nullA || nullB then nullA == nullB else x == VU.unsafeIndex b i+                            )+                            a+                        )+    (==) (PackedText bm1 p1) (PackedText bm2 p2) = eqPackedCols bm1 p1 bm2 p2+    (==) lhs@(PackedText _ _) rhs = materializePacked lhs == rhs+    (==) lhs rhs@(PackedText _ _) = lhs == materializePacked rhs+    (==) _ _ = False++{- | Byte-slice equality of two packed-text columns, skipping null slots+(a null compares equal only to a null), mirroring 'eqBoxedCols'.+-}+eqPackedCols ::+    Maybe Bitmap -> PackedTextData -> Maybe Bitmap -> PackedTextData -> Bool+eqPackedCols bm1 p1 bm2 p2+    | packedLength p1 /= packedLength p2 = False+    | otherwise = go 0+  where+    !n = packedLength p1+    go !i+        | i >= n = True+        | nullA || nullB = (nullA == nullB) && go (i + 1)+        | otherwise =+            let (a1, o1, l1) = packedSlice p1 i+                (a2, o2, l2) = packedSlice p2 i+             in sliceEqBytes a1 o1 l1 a2 o2 l2 && go (i + 1)+      where+        nullA = maybe False (\bm -> not (bitmapTestBit bm i)) bm1+        nullB = maybe False (\bm -> not (bitmapTestBit bm i)) bm2+{-# INLINE eqPackedCols #-}++{- | A class for converting a vector to a column of the appropriate type.+Given each Rep we tell the `toColumnRep` function which Column type to pick.+-}+class ColumnifyRep (r :: Rep) a where+    toColumnRep :: VB.Vector a -> Column++-- | Constraint synonym for what we can put into columns.+type Columnable a =+    ( Columnable' a+    , ColumnifyRep (KindOf a) a+    , UnboxIf a+    , IntegralIf a+    , FloatingIf a+    , SBoolI (Unboxable a)+    , SBoolI (Numeric a)+    , SBoolI (IntegralTypes a)+    , SBoolI (FloatingTypes a)+    )++instance+    (Columnable a, VU.Unbox a) =>+    ColumnifyRep 'RUnboxed a+    where+    toColumnRep :: (Columnable a, VUM.Unbox a) => VB.Vector a -> Column+    toColumnRep v = UnboxedColumn Nothing (VU.convert v)++instance+    (Columnable a) =>+    ColumnifyRep 'RBoxed a+    where+    toColumnRep :: (Columnable a) => VB.Vector a -> Column+    toColumnRep = BoxedColumn Nothing++instance+    (Columnable a) =>+    ColumnifyRep 'RNullableBoxed (Maybe a)+    where+    toColumnRep :: (Columnable a) => VB.Vector (Maybe a) -> Column+    toColumnRep = fromMaybeVec++{- | O(n) Convert a vector to a column. Automatically picks the best representation of a vector to store the underlying data in.++__Examples:__++@+> import qualified Data.Vector as V+> fromVector (VB.fromList [(1 :: Int), 2, 3, 4])+[1,2,3,4]+@+-}+fromVector ::+    forall a.+    (Columnable a, ColumnifyRep (KindOf a) a) =>+    VB.Vector a -> Column+fromVector = toColumnRep @(KindOf a)++{- | O(n) Convert an unboxed vector to a column. This avoids the extra conversion if you already have the data in an unboxed vector.++__Examples:__++@+> import qualified Data.Vector.Unboxed as V+> fromUnboxedVector (VB.fromList [(1 :: Int), 2, 3, 4])+[1,2,3,4]+@+-}+fromUnboxedVector ::+    forall a. (Columnable a, VU.Unbox a) => VU.Vector a -> Column+fromUnboxedVector = UnboxedColumn Nothing++{- | O(n) Convert a list to a column. Automatically picks the best representation of a vector to store the underlying data in.++__Examples:__++@+> fromList [(1 :: Int), 2, 3, 4]+[1,2,3,4]+@+-}+fromList ::+    forall a.+    (Columnable a, ColumnifyRep (KindOf a) a) =>+    [a] -> Column+fromList = toColumnRep @(KindOf a) . VB.fromList++{- | O(n) Create a column of random elements within a range.++Takes a random number generator, a length, and a lower and upper bound for the random values.++__Examples:__++@+> import System.Random (mkStdGen)+> mkRandom (mkStdGen 42) 4 0 10+[4,2,6,5]+@+-}+mkRandom ::+    (RandomGen g, Columnable a, ColumnifyRep (KindOf a) a, UniformRange a) =>+    g -> Int -> a -> a -> Column+mkRandom pureGen k lo hi = fromList $ go pureGen k+  where+    go _g 0 = []+    go g n =+        let+            (!v, !g') = uniformR (lo, hi) g+         in+            v : go g' (n - 1)++-- An internal helper for type errors+throwTypeMismatch ::+    forall (a :: Type) (b :: Type).+    (Typeable a, Typeable b) => Either DataFrameException Column+throwTypeMismatch =+    Left $+        TypeMismatchException+            MkTypeErrorContext+                { userType = Right (typeRep @b)+                , expectedType = Right (typeRep @a)+                , callingFunctionName = Nothing+                , errorColumnName = Nothing+                }++-- | An internal function to map a function over the values of a column.+mapColumn ::+    forall b c.+    (Columnable b, Columnable c) =>+    (b -> c) -> Column -> Either DataFrameException Column+mapColumn f = \case+    BoxedColumn bm (col :: VB.Vector a) -> runBoxed bm col+    UnboxedColumn bm (col :: VU.Vector a) -> runUnboxed bm col+    c@(PackedText _ _) -> mapColumn f (materializePacked c)+  where+    runBoxed ::+        forall a.+        (Columnable a) =>+        Maybe Bitmap -> VB.Vector a -> Either DataFrameException Column+    runBoxed bm col = case testEquality (typeRep @b) (typeRep @(Maybe a)) of+        Just Refl ->+            let !n = VB.length col+             in Right $ case sUnbox @c of+                    STrue -> UnboxedColumn Nothing $+                        VU.generate n $ \i ->+                            f+                                ( if maybe True (`bitmapTestBit` i) bm+                                    then Just (VB.unsafeIndex col i)+                                    else Nothing+                                )+                    SFalse -> fromVector @c $+                        VB.generate n $ \i ->+                            f+                                ( if maybe True (`bitmapTestBit` i) bm+                                    then Just (VB.unsafeIndex col i)+                                    else Nothing+                                )+        Nothing -> case testEquality (typeRep @a) (typeRep @b) of+            Just Refl ->+                Right $ case sUnbox @c of+                    STrue -> UnboxedColumn bm (VU.generate (VB.length col) (f . VB.unsafeIndex col))+                    SFalse -> case bm of+                        Nothing -> fromVector @c (VB.map f col)+                        Just _ -> BoxedColumn bm (VB.map f col)+            Nothing -> throwTypeMismatch @a @b++    runUnboxed ::+        forall a.+        (Columnable a, VU.Unbox a) =>+        Maybe Bitmap -> VU.Vector a -> Either DataFrameException Column+    runUnboxed bm col = case testEquality (typeRep @b) (typeRep @(Maybe a)) of+        Just Refl ->+            let !n = VU.length col+             in Right $ case sUnbox @c of+                    STrue -> UnboxedColumn Nothing $+                        VU.generate n $ \i ->+                            f+                                ( if maybe True (`bitmapTestBit` i) bm+                                    then Just (VU.unsafeIndex col i)+                                    else Nothing+                                )+                    SFalse -> fromVector @c $+                        VB.generate n $ \i ->+                            f+                                ( if maybe True (`bitmapTestBit` i) bm+                                    then Just (VU.unsafeIndex col i)+                                    else Nothing+                                )+        Nothing -> case testEquality (typeRep @a) (typeRep @b) of+            Just Refl -> Right $ case sUnbox @c of+                STrue -> UnboxedColumn bm (VU.map f col)+                SFalse -> case bm of+                    Nothing -> fromVector @c (VB.generate (VU.length col) (f . VU.unsafeIndex col))+                    Just _ -> BoxedColumn bm (VB.generate (VU.length col) (f . VU.unsafeIndex col))+            Nothing -> throwTypeMismatch @a @b+{-# INLINEABLE mapColumn #-}++-- | Applies a function that returns an unboxed result to an unboxed vector, storing the result in a column.+imapColumn ::+    forall b c.+    (Columnable b, Columnable c) =>+    (Int -> b -> c) -> Column -> Either DataFrameException Column+imapColumn f = \case+    BoxedColumn bm (col :: VB.Vector a) -> runBoxed bm col+    UnboxedColumn bm (col :: VU.Vector a) -> runUnboxed bm col+    c@(PackedText _ _) -> imapColumn f (materializePacked c)+  where+    runBoxed ::+        forall a.+        (Columnable a) =>+        Maybe Bitmap -> VB.Vector a -> Either DataFrameException Column+    runBoxed bm col = case testEquality (typeRep @a) (typeRep @b) of+        Just Refl -> Right $ case sUnbox @c of+            STrue ->+                UnboxedColumn+                    bm+                    (VU.generate (VB.length col) (\i -> f i (VB.unsafeIndex col i)))+            SFalse -> BoxedColumn bm (VB.imap f col)+        Nothing -> throwTypeMismatch @a @b++    runUnboxed ::+        forall a.+        (Columnable a, VU.Unbox a) =>+        Maybe Bitmap -> VU.Vector a -> Either DataFrameException Column+    runUnboxed bm col = case testEquality (typeRep @a) (typeRep @b) of+        Just Refl -> Right $ case sUnbox @c of+            STrue -> UnboxedColumn bm (VU.imap f col)+            SFalse -> BoxedColumn bm (VB.imap f (VG.convert col))+        Nothing -> throwTypeMismatch @a @b++-- | O(1) Gets the number of elements in the column.+columnLength :: Column -> Int+columnLength (BoxedColumn _ xs) = VB.length xs+columnLength (UnboxedColumn _ xs) = VU.length xs+columnLength (PackedText _ p) = packedLength p+{-# INLINE columnLength #-}++-- | O(n) Gets the number of non-null elements in the column.+numElements :: Column -> Int+numElements (BoxedColumn Nothing xs) = VB.length xs+numElements (BoxedColumn (Just bm) _xs) = VU.foldl' (\acc b -> acc + popCount b) 0 bm+numElements (UnboxedColumn Nothing xs) = VU.length xs+numElements (UnboxedColumn (Just bm) _xs) = VU.foldl' (\acc b -> acc + popCount b) 0 bm+numElements (PackedText Nothing p) = packedLength p+numElements (PackedText (Just bm) _p) = VU.foldl' (\acc b -> acc + popCount b) 0 bm+{-# INLINE numElements #-}++-- | O(n) Takes the first n values of a column.+takeColumn :: Int -> Column -> Column+takeColumn n (BoxedColumn bm xs) =+    BoxedColumn (fmap (bitmapSlice 0 n) bm) (VG.take n xs)+takeColumn n (UnboxedColumn bm xs) =+    UnboxedColumn (fmap (bitmapSlice 0 n) bm) (VG.take n xs)+takeColumn n (PackedText bm p) =+    PackedText (fmap (bitmapSlice 0 n) bm) (packedTake n p)+{-# INLINE takeColumn #-}++-- | O(n) Takes the last n values of a column.+takeLastColumn :: Int -> Column -> Column+takeLastColumn n column = sliceColumn (columnLength column - n) n column+{-# INLINE takeLastColumn #-}++-- | O(n) Takes n values after a given column index.+sliceColumn :: Int -> Int -> Column -> Column+sliceColumn start n (BoxedColumn bm xs) =+    BoxedColumn (fmap (bitmapSlice start n) bm) (VG.slice start n xs)+sliceColumn start n (UnboxedColumn bm xs) =+    UnboxedColumn (fmap (bitmapSlice start n) bm) (VG.slice start n xs)+sliceColumn start n c@(PackedText _ _) = sliceColumn start n (materializePacked c)+{-# INLINE sliceColumn #-}++-- | 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 bm column) =+    BoxedColumn+        ( fmap+            ( \bm0 ->+                buildBitmapFromValid $+                    VU.map (\i -> if bitmapTestBit bm0 i then 1 else 0) indexes+            )+            bm+        )+        ( VB.generate+            (VU.length indexes)+            ((column `VB.unsafeIndex`) . (indexes `VU.unsafeIndex`))+        )+atIndicesStable indexes (UnboxedColumn bm column) =+    UnboxedColumn+        ( fmap+            ( \bm0 ->+                buildBitmapFromValid $+                    VU.map (\i -> if bitmapTestBit bm0 i then 1 else 0) indexes+            )+            bm+        )+        (VU.unsafeBackpermute column indexes)+atIndicesStable indexes (PackedText bm p) =+    PackedText+        ( fmap+            ( \bm0 ->+                buildBitmapFromValid $+                    VU.map (\i -> if bitmapTestBit bm0 i then 1 else 0) indexes+            )+            bm+        )+        (packedGather indexes p)+{-# INLINE atIndicesStable #-}++{- | Like 'atIndicesStable' but treats negative indices as null.+Keeps the index vector fully unboxed (no @VB.Vector (Maybe Int)@).+-}+gatherWithSentinel :: VU.Vector Int -> Column -> Column+gatherWithSentinel indices col =+    let !n = VU.length indices+        newBm = buildBitmapFromValid $ VU.generate n $ \i ->+            if VU.unsafeIndex indices i < 0 then 0 else 1+     in case col of+            PackedText srcBm p ->+                let bm = case srcBm of+                        Nothing -> Just newBm+                        Just sb ->+                            Just+                                ( mergeBitmaps+                                    newBm+                                    ( buildBitmapFromValid $ VU.generate n $ \i ->+                                        let idx = VU.unsafeIndex indices i+                                         in if idx >= 0 && bitmapTestBit sb idx then 1 else 0+                                    )+                                )+                 in PackedText bm (packedGather indices p)+            BoxedColumn srcBm v ->+                let dat = VB.generate n $ \i ->+                        let !idx = VU.unsafeIndex indices i+                         in if idx < 0 then VB.unsafeIndex v 0 else VB.unsafeIndex v idx+                    bm = case srcBm of+                        Nothing -> Just newBm+                        Just sb ->+                            Just+                                ( mergeBitmaps+                                    newBm+                                    ( buildBitmapFromValid $ VU.generate n $ \i ->+                                        let idx = VU.unsafeIndex indices i+                                         in if idx >= 0 && bitmapTestBit sb idx then 1 else 0+                                    )+                                )+                 in BoxedColumn bm dat+            UnboxedColumn srcBm v ->+                let dat = runST $ do+                        mv <- VUM.new n+                        VG.iforM_ indices $ \i idx ->+                            when (idx >= 0) $ VUM.unsafeWrite mv i (VU.unsafeIndex v idx)+                        VU.unsafeFreeze mv+                    bm = case srcBm of+                        Nothing -> Just newBm+                        Just sb ->+                            Just+                                ( mergeBitmaps+                                    newBm+                                    ( buildBitmapFromValid $ VU.generate n $ \i ->+                                        let idx = VU.unsafeIndex indices i+                                         in if idx >= 0 && bitmapTestBit sb idx then 1 else 0+                                    )+                                )+                 in UnboxedColumn bm dat+{-# INLINE gatherWithSentinel #-}++-- | Internal helper to get indices in a boxed vector.+getIndices :: VU.Vector Int -> VB.Vector a -> VB.Vector a+getIndices indices xs = VB.generate (VU.length indices) (\i -> xs VB.! (indices VU.! i))+{-# INLINE getIndices #-}++-- | Internal helper to get indices in an unboxed vector.+getIndicesUnboxed :: (VU.Unbox a) => VU.Vector Int -> VU.Vector a -> VU.Vector a+getIndicesUnboxed indices xs = VU.generate (VU.length indices) (\i -> xs VU.! (indices VU.! i))+{-# INLINE getIndicesUnboxed #-}++findIndices ::+    forall a.+    (Columnable a) =>+    (a -> Bool) ->+    Column ->+    Either DataFrameException (VU.Vector Int)+findIndices predicate = \case+    BoxedColumn _ (v :: VB.Vector b) -> run v VG.convert+    UnboxedColumn _ (v :: VU.Vector b) -> run v id+    c@(PackedText _ _) -> findIndices predicate (materializePacked c)+  where+    run ::+        forall b v.+        (Typeable b, VG.Vector v b, VG.Vector v Int) =>+        v b ->+        (v Int -> VU.Vector Int) ->+        Either DataFrameException (VU.Vector Int)+    run column finalize = case testEquality (typeRep @a) (typeRep @b) of+        Just Refl -> Right . finalize $ VG.findIndices predicate column+        Nothing ->+            Left $+                TypeMismatchException+                    MkTypeErrorContext+                        { userType = Right (typeRep @a)+                        , expectedType = Right (typeRep @b)+                        , callingFunctionName = Just "findIndices"+                        , errorColumnName = Nothing+                        }++-- | Fold (right) column with index.+ifoldrColumn ::+    forall a b.+    (Columnable a, Columnable b) =>+    (Int -> a -> b -> b) -> b -> Column -> Either DataFrameException b+ifoldrColumn f acc = \case+    BoxedColumn _ column -> foldrWorker column+    UnboxedColumn _ column -> foldrWorker column+    c@(PackedText _ _) -> ifoldrColumn f acc (materializePacked c)+  where+    foldrWorker ::+        forall c v.+        (Typeable c, VG.Vector v c) =>+        v c ->+        Either DataFrameException b+    foldrWorker vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl -> pure $ VG.ifoldr f acc vec+        Nothing ->+            Left $+                TypeMismatchException+                    ( MkTypeErrorContext+                        { userType = Right (typeRep @a)+                        , expectedType = Right (typeRep @c)+                        , callingFunctionName = Just "ifoldrColumn"+                        , errorColumnName = Nothing+                        }+                    )++foldlColumn ::+    forall a b.+    (Columnable a, Columnable b) =>+    (b -> a -> b) -> b -> Column -> Either DataFrameException b+foldlColumn f acc = \case+    BoxedColumn _ column -> foldlWorker column+    UnboxedColumn _ column -> foldlWorker column+    c@(PackedText _ _) -> foldlColumn f acc (materializePacked c)+  where+    foldlWorker ::+        forall c v.+        (Typeable c, VG.Vector v c) =>+        v c ->+        Either DataFrameException b+    foldlWorker vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl -> pure $ VG.foldl' f acc vec+        Nothing ->+            Left $+                TypeMismatchException+                    ( MkTypeErrorContext+                        { userType = Right (typeRep @a)+                        , expectedType = Right (typeRep @c)+                        , callingFunctionName = Just "ifoldrColumn"+                        , errorColumnName = Nothing+                        }+                    )++foldl1Column ::+    forall a.+    (Columnable a) =>+    (a -> a -> a) -> Column -> Either DataFrameException a+foldl1Column f = \case+    BoxedColumn _ column -> foldl1Worker column+    UnboxedColumn _ column -> foldl1Worker column+    c@(PackedText _ _) -> foldl1Column f (materializePacked c)+  where+    foldl1Worker ::+        forall c v.+        (Typeable c, VG.Vector v c) =>+        v c ->+        Either DataFrameException a+    foldl1Worker vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl -> pure $ VG.foldl1' f vec+        Nothing ->+            Left $+                TypeMismatchException+                    ( MkTypeErrorContext+                        { userType = Right (typeRep @a)+                        , expectedType = Right (typeRep @c)+                        , callingFunctionName = Just "foldl1Column"+                        , errorColumnName = Nothing+                        }+                    )++{- | 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 Column+foldl1DirectGroups f col valueIndices offsets+    | VU.length offsets <= 1 = pure $ fromVector @a VB.empty+    | otherwise = case col of+        UnboxedColumn _ (vec :: VU.Vector d) -> UnboxedColumn Nothing <$> foldl1Worker vec+        BoxedColumn _ (vec :: VB.Vector d) -> BoxedColumn Nothing <$> foldl1Worker vec+        PackedText _ _ -> foldl1DirectGroups f (materializePacked col) valueIndices offsets+  where+    foldl1Worker ::+        forall c v.+        (Typeable c, VG.Vector v c) =>+        v c ->+        Either DataFrameException (v c)+    foldl1Worker vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl ->+            Right $+                VG.generate (VU.length offsets - 1) foldGroup+          where+            foldGroup k =+                let !s = VU.unsafeIndex offsets k+                    !e = VU.unsafeIndex offsets (k + 1)+                    !seed = VG.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 (VG.unsafeIndex vec (VU.unsafeIndex valueIndices i))+        Nothing ->+            Left $+                TypeMismatchException+                    MkTypeErrorContext+                        { userType = Right (typeRep @a)+                        , expectedType = Right (typeRep @c)+                        , callingFunctionName = Just "foldl1DirectGroups"+                        , errorColumnName = Nothing+                        }+{-# INLINEABLE foldl1DirectGroups #-}++{- | O(n) fold over groups by scanning the column linearly (rowToGroup[i] = group+of row i). Random writes hit the small per-group accumulator array; when @acc@ is+unboxable that array is unboxed, avoiding pointer indirection.+-}+foldLinearGroups ::+    forall b acc.+    (Columnable b, Columnable acc) =>+    (acc -> b -> acc) ->+    acc ->+    Column ->+    VU.Vector Int ->+    Int ->+    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) -> foldLinearWorker vec+        BoxedColumn _ (vec :: VB.Vector d) -> foldLinearWorker vec+        PackedText _ _ ->+            foldLinearGroups f seed (materializePacked col) rowToGroup nGroups+  where+    foldLinearWorker ::+        forall c v.+        (Typeable c, VG.Vector v c) =>+        v c ->+        Either DataFrameException Column+    foldLinearWorker vec = case testEquality (typeRep @b) (typeRep @c) of+        Just Refl ->+            Right $+                unsafePerformIO $+                    runWith+                        ( \readAt writeAt ->+                            VG.iforM_ vec $ \row x -> do+                                let !k = VG.unsafeIndex rowToGroup row+                                cur <- readAt k+                                writeAt k $! f cur x+                        )+        Nothing ->+            Left $+                TypeMismatchException+                    MkTypeErrorContext+                        { userType = Right (typeRep @b)+                        , expectedType = Right (typeRep @c)+                        , callingFunctionName = Just "foldLinearGroups"+                        , errorColumnName = Nothing+                        }++    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 Nothing <$> 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 = \case+    BoxedColumn _ col -> headWorker col+    UnboxedColumn _ col -> headWorker col+    c@(PackedText _ _) -> headColumn (materializePacked c)+  where+    headWorker ::+        forall c v.+        (Typeable c, VG.Vector v c) =>+        v c ->+        Either DataFrameException a+    headWorker vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl ->+            if VG.null vec+                then Left (EmptyDataSetException "headColumn")+                else pure (VG.head vec)+        Nothing ->+            Left $+                TypeMismatchException+                    ( MkTypeErrorContext+                        { userType = Right (typeRep @a)+                        , expectedType = Right (typeRep @c)+                        , callingFunctionName = Just "headColumn"+                        , errorColumnName = Nothing+                        }+                    )++-- | An internal, column version of zip.+zipColumns :: Column -> Column -> Column+zipColumns l@(PackedText _ _) r = zipColumns (materializePacked l) r+zipColumns l r@(PackedText _ _) = zipColumns l (materializePacked r)+zipColumns (BoxedColumn _ column) (BoxedColumn _ other) = BoxedColumn Nothing (VG.zip column other)+zipColumns (BoxedColumn _ column) (UnboxedColumn _ other) =+    BoxedColumn+        Nothing+        ( VB.generate+            (min (VG.length column) (VG.length other))+            (\i -> (column VG.! i, other VG.! i))+        )+zipColumns (UnboxedColumn _ column) (BoxedColumn _ other) =+    BoxedColumn+        Nothing+        ( VB.generate+            (min (VG.length column) (VG.length other))+            (\i -> (column VG.! i, other VG.! i))+        )+zipColumns (UnboxedColumn _ column) (UnboxedColumn _ other) = UnboxedColumn Nothing (VG.zip column other)+{-# INLINE zipColumns #-}++-- | Merge two columns using `These`.+mergeColumns :: Column -> Column -> Column+mergeColumns colA colB = case (colA, colB) of+    (PackedText _ _, _) -> mergeColumns (materializePacked colA) colB+    (_, PackedText _ _) -> mergeColumns colA (materializePacked colB)+    (BoxedColumn bmA c1, BoxedColumn bmB c2) -> case (bmA, bmB) of+        (Just ba, Just bb) ->+            BoxedColumn Nothing $ mkVec c1 c2 $ \i v1 v2 ->+                let nullA = not (bitmapTestBit ba i)+                    nullB = not (bitmapTestBit bb i)+                 in case (nullA, nullB) of+                        (True, True) -> error "mergeColumns: both null"+                        (False, True) -> This v1+                        (True, False) -> That v2+                        (False, False) -> These v1 v2+        (Just ba, Nothing) ->+            BoxedColumn Nothing $ mkVec c1 c2 $ \i v1 v2 ->+                if not (bitmapTestBit ba i) then That v2 else These v1 v2+        (Nothing, Just bb) ->+            BoxedColumn Nothing $ mkVec c1 c2 $ \i v1 v2 ->+                if not (bitmapTestBit bb i) then This v1 else These v1 v2+        (Nothing, Nothing) ->+            BoxedColumn Nothing $ mkVecSimple c1 c2 These+    (BoxedColumn _ c1, UnboxedColumn _ c2) ->+        BoxedColumn Nothing $ mkVecSimple c1 c2 These+    (UnboxedColumn _ c1, BoxedColumn _ c2) ->+        BoxedColumn Nothing $ mkVecSimple c1 c2 These+    (UnboxedColumn _ c1, UnboxedColumn _ c2) ->+        BoxedColumn Nothing $ mkVecSimple c1 c2 These+  where+    mkVec c1 c2 combineElements =+        VB.generate+            (min (VG.length c1) (VG.length c2))+            (\i -> combineElements i (c1 VG.! i) (c2 VG.! i))+    {-# INLINE mkVec #-}++    mkVecSimple c1 c2 f =+        VB.generate+            (min (VG.length c1) (VG.length c2))+            (\i -> f (c1 VG.! i) (c2 VG.! i))+    {-# INLINE mkVecSimple #-}+{-# INLINE mergeColumns #-}++-- | An internal, column version of zipWith.+zipWithColumns ::+    forall a b c.+    (Columnable a, Columnable b, Columnable c) =>+    (a -> b -> c) -> Column -> Column -> Either DataFrameException Column+zipWithColumns f (UnboxedColumn bmL (column :: VU.Vector d)) (UnboxedColumn bmR (other :: VU.Vector e)) = case testEquality (typeRep @a) (typeRep @d) of+    Just Refl -> case testEquality (typeRep @b) (typeRep @e) of+        Just Refl+            | isNothing bmL+            , isNothing bmR ->+                pure $ case sUnbox @c of+                    STrue -> UnboxedColumn Nothing (VU.zipWith f column other)+                    SFalse -> fromVector $ VB.zipWith f (VG.convert column) (VG.convert other)+        _ -> zipWithColumnsGeneral f (UnboxedColumn bmL column) (UnboxedColumn bmR other)+    Nothing -> zipWithColumnsGeneral f (UnboxedColumn bmL column) (UnboxedColumn bmR other)+-- TODO: mchavinda - reuse pattern from interpret where we augment the+-- error at the end.+zipWithColumns f left right = zipWithColumnsGeneral f left right++zipWithColumnsGeneral ::+    forall a b c.+    (Columnable a, Columnable b, Columnable c) =>+    (a -> b -> c) -> Column -> Column -> Either DataFrameException Column+zipWithColumnsGeneral f left right = case toVector @a left of+    Left (TypeMismatchException context) ->+        Left $+            TypeMismatchException (context{callingFunctionName = Just "zipWithColumns"})+    Left e -> Left e+    Right left' -> case toVector @b right of+        Left (TypeMismatchException context) ->+            Left $+                TypeMismatchException (context{callingFunctionName = Just "zipWithColumns"})+        Left e -> Left e+        Right right' -> pure $ fromVector $ VB.zipWith f left' right'+{-# INLINE zipWithColumnsGeneral #-}+{-# INLINE zipWithColumns #-}++-- writeColumn and freezeColumn' (CSV-ingest helpers) moved to+-- DataFrame.IO.Internal.MutableColumn so the core column module does not+-- need to depend on DataFrame.Internal.Parsing.++{- | Freeze a mutable column into an @Either Text a@ column: every recorded+null position becomes @Left rawText@ (preserving the original input), every+other position becomes @Right v@. Used by CSV readers under 'EitherRead' mode.+-}+freezeColumnEither :: [(Int, T.Text)] -> MutableColumn -> IO Column+freezeColumnEither nulls (MBoxedColumn col) = do+    frozen <- VB.unsafeFreeze col+    let nullMap = nulls+    pure $+        BoxedColumn Nothing $+            VB.imap+                ( \i v -> case lookup i nullMap of+                    Just t -> Left t+                    Nothing -> Right v+                )+                frozen+freezeColumnEither nulls (MUnboxedColumn col) = do+    c <- VU.unsafeFreeze col+    let nullMap = nulls+    pure $+        BoxedColumn Nothing $+            VB.generate (VU.length c) $ \i ->+                case lookup i nullMap of+                    Just t -> Left t+                    Nothing -> Right (c VU.! i)+{-# INLINE freezeColumnEither #-}++{- | Promote a non-nullable column to a nullable one (add an all-valid bitmap).+No-op when already nullable.+-}+ensureOptional :: Column -> Column+ensureOptional c@(BoxedColumn (Just _) _) = c+ensureOptional (BoxedColumn Nothing col) =+    BoxedColumn (Just (allValidBitmap (VB.length col))) col+ensureOptional c@(UnboxedColumn (Just _) _) = c+ensureOptional (UnboxedColumn Nothing col) =+    UnboxedColumn (Just (allValidBitmap (VU.length col))) col+ensureOptional c@(PackedText (Just _) _) = c+ensureOptional (PackedText Nothing p) =+    PackedText (Just (allValidBitmap (packedLength p))) p++-- | Fills the end of a column, up to n, with null rows. Does nothing if column has length >= n.+expandColumn :: Int -> Column -> Column+expandColumn n c@(PackedText _ p)+    | n <= packedLength p = c+    | otherwise = expandColumn n (materializePacked c)+expandColumn n column@(BoxedColumn bm col)+    | n <= VG.length col = column+    | otherwise =+        let extra = n - VG.length col+            newBm = case bm of+                Nothing -> Just (buildBitmapFromNulls n [VG.length col .. n - 1])+                Just b ->+                    Just+                        (bitmapConcat (VG.length col) b extra (VU.replicate ((extra + 7) `shiftR` 3) 0))+            newCol = col <> VB.replicate extra (errorWithoutStackTrace "expandColumn: null slot")+         in BoxedColumn newBm newCol+expandColumn n column@(UnboxedColumn bm col)+    | n <= VG.length col = column+    | otherwise =+        let extra = n - VG.length col+            newBm = case bm of+                Nothing -> Just (buildBitmapFromNulls n [VG.length col .. n - 1])+                Just b ->+                    Just+                        (bitmapConcat (VG.length col) b extra (VU.replicate ((extra + 7) `shiftR` 3) 0))+            newCol = runST $ do+                mv <- VUM.new n+                VU.imapM_ (VUM.unsafeWrite mv) col+                VU.unsafeFreeze mv+         in UnboxedColumn newBm newCol++-- | Fills the beginning of a column, up to n, with null rows. Does nothing if column has length >= n.+leftExpandColumn :: Int -> Column -> Column+leftExpandColumn n c@(PackedText _ p)+    | n <= packedLength p = c+    | otherwise = leftExpandColumn n (materializePacked c)+leftExpandColumn n column@(BoxedColumn bm col)+    | n <= VG.length col = column+    | otherwise =+        let extra = n - VG.length col+            origLen = VG.length col+            newBm = case bm of+                Nothing -> Just (buildBitmapFromNulls n [0 .. extra - 1])+                Just b ->+                    let nullPart = VU.replicate ((extra + 7) `shiftR` 3) 0+                     in Just (bitmapConcat extra nullPart origLen b)+            newCol =+                VB.replicate extra (errorWithoutStackTrace "leftExpandColumn: null slot") <> col+         in BoxedColumn newBm newCol+leftExpandColumn n column@(UnboxedColumn bm col)+    | n <= VG.length col = column+    | otherwise =+        let extra = n - VG.length col+            origLen = VG.length col+            newBm = case bm of+                Nothing -> Just (buildBitmapFromNulls n [0 .. extra - 1])+                Just b ->+                    let nullPart = VU.replicate ((extra + 7) `shiftR` 3) 0+                     in Just (bitmapConcat extra nullPart origLen b)+            newCol = runST $ do+                mv <- VUM.new n+                VU.imapM_ (\i x -> VUM.unsafeWrite mv (extra + i) x) col+                VU.unsafeFreeze mv+         in UnboxedColumn newBm newCol++{- | Concatenates two columns.+Returns Nothing if the columns are of different types.+-}+concatColumns :: Column -> Column -> Either DataFrameException Column+concatColumns left right = case (left, right) of+    (PackedText _ _, _) -> concatColumns (materializePacked left) right+    (_, PackedText _ _) -> concatColumns left (materializePacked right)+    (BoxedColumn bmL l, BoxedColumn bmR r) -> case testEquality (typeOf l) (typeOf r) of+        Just Refl ->+            let newBm = case (bmL, bmR) of+                    (Nothing, Nothing) -> Nothing+                    (Just bl, Nothing) ->+                        Just+                            (bitmapConcat (VB.length l) bl (VB.length r) (allValidBitmap (VB.length r)))+                    (Nothing, Just br) ->+                        Just+                            (bitmapConcat (VB.length l) (allValidBitmap (VB.length l)) (VB.length r) br)+                    (Just bl, Just br) -> Just (bitmapConcat (VB.length l) bl (VB.length r) br)+             in pure (BoxedColumn newBm (l <> r))+        Nothing -> Left (mismatchErr (typeOf r) (typeOf l))+    (UnboxedColumn bmL l, UnboxedColumn bmR r) -> case testEquality (typeOf l) (typeOf r) of+        Just Refl ->+            let newBm = case (bmL, bmR) of+                    (Nothing, Nothing) -> Nothing+                    (Just bl, Nothing) ->+                        Just+                            (bitmapConcat (VU.length l) bl (VU.length r) (allValidBitmap (VU.length r)))+                    (Nothing, Just br) ->+                        Just+                            (bitmapConcat (VU.length l) (allValidBitmap (VU.length l)) (VU.length r) br)+                    (Just bl, Just br) -> Just (bitmapConcat (VU.length l) bl (VU.length r) br)+             in pure (UnboxedColumn newBm (l <> r))+        Nothing -> Left (mismatchErr (typeOf r) (typeOf l))+    _ -> Left (mismatchErr (typeOf right) (typeOf left))+  where+    mismatchErr ::+        forall (x :: Type) (y :: Type). TypeRep x -> TypeRep y -> DataFrameException+    mismatchErr ta tb =+        withTypeable ta $+            withTypeable tb $+                TypeMismatchException+                    ( MkTypeErrorContext+                        { userType = Right ta+                        , expectedType = Right tb+                        , callingFunctionName = Just "concatColumns"+                        , errorColumnName = Nothing+                        }+                    )++{- | Like 'concatColumns' but also combines columns of different types by wrapping+values in 'Either' (e.g. @[1,2]@ and @["a","b"]@ become+@[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 all'+    | any isPackedText all' =+        concatManyColumns (map materializePacked all')+concatManyColumns (c0 : cs) = case c0 of+    BoxedColumn bm0 v0 ->+        let getCol (BoxedColumn bm v) = case testEquality (typeOf v0) (typeOf v) of+                Just Refl -> (bm, v)+                Nothing -> error "concatManyColumns: BoxedColumn type mismatch"+            getCol _ = error "concatManyColumns: column constructor mismatch"+            rest = map getCol cs+            allVecs = v0 : map snd rest+            allBms = bm0 : map fst rest+            newBm+                | all isNothing allBms = Nothing+                | otherwise =+                    let pairs = zip allVecs allBms+                        expandedBms = map (\(v, mb) -> fromMaybe (allValidBitmap (VB.length v)) mb) pairs+                        go b1 n1 b2 n2 = bitmapConcat n1 b1 n2 b2+                        concatBms [] = VU.empty+                        concatBms [(b, _v)] = b+                        concatBms ((b1, v1) : (b2, v2) : rest') =+                            let merged = go b1 (VB.length v1) b2 (VB.length v2)+                             in concatBms ((merged, v1 <> v2) : rest')+                     in Just $ concatBms (zip expandedBms allVecs)+         in BoxedColumn newBm (VB.concat allVecs)+    UnboxedColumn bm0 v0 ->+        let getCol (UnboxedColumn bm v) = case testEquality (typeOf v0) (typeOf v) of+                Just Refl -> (bm, v)+                Nothing -> error "concatManyColumns: UnboxedColumn type mismatch"+            getCol _ = error "concatManyColumns: column constructor mismatch"+            rest = map getCol cs+            allVecs = v0 : map snd rest+            allBms = bm0 : map fst rest+            newBm+                | all isNothing allBms = Nothing+                | otherwise =+                    let pairs = zip allVecs allBms+                        expandedBms = map (\(v, mb) -> fromMaybe (allValidBitmap (VU.length v)) mb) pairs+                        go b1 n1 b2 n2 = bitmapConcat n1 b1 n2 b2+                        concatBms [] = VU.empty+                        concatBms [(b, _)] = b+                        concatBms ((b1, v1) : (b2, v2) : rest') =+                            let merged = go b1 (VU.length v1) b2 (VU.length v2)+                             in concatBms ((merged, v1 <> v2) : rest')+                     in Just $ concatBms (zip expandedBms allVecs)+         in UnboxedColumn newBm (VU.concat allVecs)+    PackedText _ _ -> concatManyColumns (map materializePacked (c0 : cs))++concatColumnsEither :: Column -> Column -> Column+concatColumnsEither l@(PackedText _ _) r = concatColumnsEither (materializePacked l) r+concatColumnsEither l r@(PackedText _ _) = concatColumnsEither l (materializePacked r)+concatColumnsEither (BoxedColumn bmL left) (BoxedColumn bmR right) = case testEquality (typeOf left) (typeOf right) of+    Nothing ->+        BoxedColumn Nothing $ fmap Left left <> fmap Right right+    Just Refl ->+        let newBm = case (bmL, bmR) of+                (Nothing, Nothing) -> Nothing+                (Just bl, Nothing) ->+                    Just+                        ( bitmapConcat+                            (VB.length left)+                            bl+                            (VB.length right)+                            (allValidBitmap (VB.length right))+                        )+                (Nothing, Just br) ->+                    Just+                        ( bitmapConcat+                            (VB.length left)+                            (allValidBitmap (VB.length left))+                            (VB.length right)+                            br+                        )+                (Just bl, Just br) -> Just (bitmapConcat (VB.length left) bl (VB.length right) br)+         in BoxedColumn newBm $ left <> right+concatColumnsEither (UnboxedColumn bmL left) (UnboxedColumn bmR right) = case testEquality (typeOf left) (typeOf right) of+    Nothing ->+        BoxedColumn Nothing $+            fmap Left (VG.convert left) <> fmap Right (VG.convert right)+    Just Refl ->+        let newBm = case (bmL, bmR) of+                (Nothing, Nothing) -> Nothing+                (Just bl, Nothing) ->+                    Just+                        ( bitmapConcat+                            (VU.length left)+                            bl+                            (VU.length right)+                            (allValidBitmap (VU.length right))+                        )+                (Nothing, Just br) ->+                    Just+                        ( bitmapConcat+                            (VU.length left)+                            (allValidBitmap (VU.length left))+                            (VU.length right)+                            br+                        )+                (Just bl, Just br) -> Just (bitmapConcat (VU.length left) bl (VU.length right) br)+         in UnboxedColumn newBm $ left <> right+concatColumnsEither (BoxedColumn _ left) (UnboxedColumn _ right) =+    BoxedColumn Nothing $ fmap Left left <> fmap Right (VG.convert right)+concatColumnsEither (UnboxedColumn _ left) (BoxedColumn _ right) =+    BoxedColumn Nothing $ fmap Left (VG.convert left) <> 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 (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))+newMutableColumn n c@(PackedText _ _) = newMutableColumn n (materializePacked c)++-- | Copy a column chunk into a mutable column starting at offset @off@.+copyIntoMutableColumn :: MutableColumn -> Int -> Column -> IO ()+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 mc off c@(PackedText _ _) =+    copyIntoMutableColumn mc off (materializePacked c)+copyIntoMutableColumn _ _ _ =+    error "copyIntoMutableColumn: constructor mismatch"++-- | Freeze a mutable column into an immutable column.+freezeMutableColumn :: MutableColumn -> IO Column+freezeMutableColumn (MBoxedColumn mv) = BoxedColumn Nothing <$> VB.unsafeFreeze mv+freezeMutableColumn (MUnboxedColumn mv) = UnboxedColumn Nothing <$> VU.unsafeFreeze mv++{- | O(n) Converts a column to a list. Throws an exception if the wrong type is specified.++__Examples:__++@+> column = fromList [(1 :: Int), 2, 3, 4]+> toList @Int column+[1,2,3,4]+> toList @Double column+exception: ...+@+-}+toList :: forall a. (Columnable a) => Column -> [a]+toList xs = case toVector @a xs of+    Left err -> throw err+    Right val -> VB.toList val++{- | Type-safe conversion of a column to a vector of element type @a@ (specify via+type application); 'Left' 'TypeMismatchException' when the column's type differs.++>>> toVector @Int @VU.Vector column+Right (unboxed vector of Ints)++>>> toVector @Text @VB.Vector column+Right (boxed vector of Text)+-}+toVector ::+    forall a v.+    (VG.Vector v a, Columnable a) => Column -> Either DataFrameException (v a)+toVector col = case col of+    PackedText _ _ -> toVector (materializePacked col)+    BoxedColumn bm (inner :: VB.Vector c) ->+        -- Check if user wants Maybe c (nullable) or c directly+        case testEquality (typeRep @a) (typeRep @c) of+            Just Refl -> Right $ VG.convert inner+            Nothing ->+                -- Try: a = Maybe c+                case testEquality (typeRep @a) (typeRep @(Maybe c)) of+                    Just Refl ->+                        -- Use VB.generate to avoid fusion forcing null slots+                        let !n = VB.length inner+                            maybeVec = case bm of+                                Nothing -> VB.generate n (Just . VB.unsafeIndex inner)+                                Just bitmap -> VB.generate n $ \i ->+                                    if bitmapTestBit bitmap i then Just (VB.unsafeIndex inner i) else Nothing+                         in Right $ VG.convert maybeVec+                    Nothing ->+                        Left $+                            TypeMismatchException+                                ( MkTypeErrorContext+                                    { userType = Right (typeRep @a)+                                    , expectedType = Right (typeRep @c)+                                    , callingFunctionName = Just "toVector"+                                    , errorColumnName = Nothing+                                    }+                                )+    UnboxedColumn bm (inner :: VU.Vector c) ->+        case testEquality (typeRep @a) (typeRep @c) of+            Just Refl -> Right $ VG.convert inner+            Nothing ->+                case testEquality (typeRep @a) (typeRep @(Maybe c)) of+                    Just Refl ->+                        let maybeVec = case bm of+                                Nothing -> VB.generate (VU.length inner) (Just . VU.unsafeIndex inner)+                                Just bitmap -> VB.generate (VU.length inner) $ \i ->+                                    if bitmapTestBit bitmap i then Just (VU.unsafeIndex inner i) else Nothing+                         in Right $ VG.convert maybeVec+                    Nothing ->+                        Left $+                            TypeMismatchException+                                ( MkTypeErrorContext+                                    { userType = Right (typeRep @a)+                                    , expectedType = Right (typeRep @c)+                                    , callingFunctionName = Just "toVector"+                                    , errorColumnName = Nothing+                                    }+                                )++-- Some common types we will use for numerical computing.++{- | Convert a column to an unboxed 'Double' vector, coercing numeric types+('realToFrac' for floats, 'fromIntegral' for integrals; nulls become @NaN@).+'Left' 'TypeMismatchException' when the column is not numeric.+-}+toDoubleVector :: Column -> Either DataFrameException (VU.Vector Double)+toDoubleVector column =+    case column of+        PackedText _ _ -> toDoubleVector (materializePacked column)+        UnboxedColumn bm (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Double) of+            Just Refl -> case bm of+                Nothing -> Right f+                Just bitmap -> Right $ VU.imap (\i x -> if bitmapTestBit bitmap i then x else read "NaN") f+            Nothing -> case sFloating @a of+                STrue ->+                    Right+                        ( VU.imap+                            ( \i x -> case bm of+                                Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"+                                _ -> realToFrac x+                            )+                            f+                        )+                SFalse -> case sIntegral @a of+                    STrue ->+                        Right+                            ( VU.imap+                                ( \i x -> case bm of+                                    Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"+                                    _ -> fromIntegral x+                                )+                                f+                            )+                    SFalse ->+                        Left $+                            TypeMismatchException+                                ( MkTypeErrorContext+                                    { userType = Right (typeRep @Double)+                                    , expectedType = Right (typeRep @a)+                                    , callingFunctionName = Just "toDoubleVector"+                                    , errorColumnName = Nothing+                                    }+                                )+        BoxedColumn bm (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of+            Just Refl ->+                Right+                    ( VB.convert $+                        VB.imap+                            ( \i x -> case bm of+                                Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"+                                _ -> fromIntegral x+                            )+                            f+                    )+            Nothing ->+                Left $+                    TypeMismatchException+                        ( MkTypeErrorContext+                            { userType = Right (typeRep @Double)+                            , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())+                            , callingFunctionName = Just "toDoubleVector"+                            , errorColumnName = Nothing+                            }+                        )++{- | Convert a column to an unboxed 'Float' vector, coercing numeric types (nulls+become @NaN@); 'Left' 'TypeMismatchException' when not numeric. Converting from+'Double' may lose precision.+-}+toFloatVector :: Column -> Either DataFrameException (VU.Vector Float)+toFloatVector column =+    case column of+        PackedText _ _ -> toFloatVector (materializePacked column)+        UnboxedColumn bm (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Float) of+            Just Refl -> case bm of+                Nothing -> Right f+                Just bitmap -> Right $ VU.imap (\i x -> if bitmapTestBit bitmap i then x else read "NaN") f+            Nothing -> case sFloating @a of+                STrue ->+                    Right+                        ( VU.imap+                            ( \i x -> case bm of+                                Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"+                                _ -> realToFrac x+                            )+                            f+                        )+                SFalse -> case sIntegral @a of+                    STrue ->+                        Right+                            ( VU.imap+                                ( \i x -> case bm of+                                    Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"+                                    _ -> fromIntegral x+                                )+                                f+                            )+                    SFalse ->+                        Left $+                            TypeMismatchException+                                ( MkTypeErrorContext+                                    { userType = Right (typeRep @Float)+                                    , expectedType = Right (typeRep @a)+                                    , callingFunctionName = Just "toFloatVector"+                                    , errorColumnName = Nothing+                                    }+                                )+        BoxedColumn bm (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of+            Just Refl ->+                Right+                    ( VB.convert $+                        VB.imap+                            ( \i x -> case bm of+                                Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"+                                _ -> fromIntegral x+                            )+                            f+                    )+            Nothing ->+                Left $+                    TypeMismatchException+                        ( MkTypeErrorContext+                            { userType = Right (typeRep @Float)+                            , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())+                            , callingFunctionName = Just "toFloatVector"+                            , errorColumnName = Nothing+                            }+                        )++{- | Convert a column to an unboxed 'Int' vector, coercing numeric types+(floats are 'round'ed via banker's rounding); 'Left' 'TypeMismatchException'+when the column is not numeric. Does not support nullable columns.+-}+toIntVector :: Column -> Either DataFrameException (VU.Vector Int)+toIntVector column =+    case column of+        PackedText _ _ -> toIntVector (materializePacked column)+        UnboxedColumn _ (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Int) of+            Just Refl -> Right f+            Nothing -> case sFloating @a of+                STrue -> Right (VU.map (round . (realToFrac :: a -> Double)) f)+                SFalse -> case sIntegral @a of+                    STrue -> Right (VU.map fromIntegral f)+                    SFalse ->+                        Left $+                            TypeMismatchException+                                ( MkTypeErrorContext+                                    { userType = Right (typeRep @Int)+                                    , expectedType = Right (typeRep @a)+                                    , callingFunctionName = Just "toIntVector"+                                    , errorColumnName = Nothing+                                    }+                                )+        BoxedColumn _ (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of+            Just Refl -> Right (VB.convert $ VB.map fromIntegral f)+            Nothing ->+                Left $+                    TypeMismatchException+                        ( MkTypeErrorContext+                            { userType = Right (typeRep @Int)+                            , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())+                            , callingFunctionName = Just "toIntVector"+                            , errorColumnName = Nothing+                            }+                        )++toUnboxedVector ::+    forall a.+    (Columnable a, VU.Unbox a) => Column -> Either DataFrameException (VU.Vector a)+toUnboxedVector column =+    case column of+        UnboxedColumn _ (f :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of+            Just Refl -> Right f+            Nothing ->+                Left $+                    TypeMismatchException+                        ( MkTypeErrorContext+                            { userType = Right (typeRep @a)+                            , expectedType = Right (typeRep @b)+                            , callingFunctionName = Just "toUnboxedVector"+                            , errorColumnName = Nothing+                            }+                        )+        _ ->+            Left $+                TypeMismatchException+                    ( MkTypeErrorContext+                        { userType = Right (typeRep @a)+                        , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())+                        , callingFunctionName = Just "toUnboxedVector"+                        , errorColumnName = Nothing+                        }+                    )+{-# INLINE toUnboxedVector #-}++-- Shared finaliser for the two parseUnboxedColumn* helpers.  Freezes+-- the mutable data vector, and only materialises the bitmap when the+-- column actually had nulls.+{-# INLINE finalizeParseResult #-}+finalizeParseResult ::+    (VU.Unbox a) =>+    VUM.STVector s a ->+    VUM.STVector s Word8 ->+    Bool ->+    ST s (Maybe (Maybe Bitmap, VU.Vector a))+finalizeParseResult values vmask anyNull+    | anyNull = do+        vs <- VU.unsafeFreeze values+        vm <- VU.unsafeFreeze vmask+        return (Just (Just (buildBitmapFromValid vm), vs))+    | otherwise = do+        vs <- VU.unsafeFreeze values+        return (Just (Nothing, vs))

@@ -0,0 +1,299 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE UndecidableInstances #-}++{- | Mutable, growable column builders for high-throughput ingest. No+per-append @IORef@ traffic: hot counters live in an unboxed vector, payloads+double on demand, and validity is only materialized once a null is seen.+-}+module DataFrame.Internal.ColumnBuilder (+    ColumnBuilder (..),+    NumBuilder,+    IntBuilder,+    DoubleBuilder,+    TextBuilder,+    TextChunk (..),+    newIntBuilder,+    newDoubleBuilder,+    newNumBuilder,+    newTextBuilder,+    appendInt,+    appendDouble,+    appendNum,+    appendText,+    appendTextSlice,+    appendTextSliceFromPtr,+    freezeTextChunk,+    mergeColumns,+    mergeTextChunks,+) where++import qualified Data.Text as T+import qualified Data.Text.Array as A+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM++import Control.Monad (when)+import Control.Monad.ST (ST)+import Data.Bits (shiftR)+import Data.STRef+import Data.Text.Internal (Text (..))+import Data.Word (Word8)+import DataFrame.Internal.Column hiding (mergeColumns)+import DataFrame.Internal.ColumnMerge (+    TextChunk (..),+    mergeColumns,+    mergeTextChunks,+    packValidity,+ )+import Foreign.Ptr (Ptr)++{- | Operations shared by all column builders. A builder must not be used+again after 'freezeBuilder' (its storage is frozen in place, not copied).+-}+class ColumnBuilder b where+    -- | Append a null row (sentinel payload + invalid bit).+    appendNull :: b s -> ST s ()++    -- | Rows appended so far.+    builderLength :: b s -> ST s Int++    -- | Freeze into a fully-forced 'Column'; bitmap only when a null was seen.+    freezeBuilder :: b s -> ST s Column++-- Counter slots shared by the builders: rows, any-null flag, text bytes used.+cRows, cAnyNull, cBytes :: Int+cRows = 0+cAnyNull = 1+cBytes = 2++{- | Builder for unboxed numeric payloads ('Int', 'Double', ...). 'nbNull'+is the sentinel written into null slots (protected by the bitmap).+-}+data NumBuilder a s = NumBuilder+    { nbNull :: !a+    , nbCounters :: !(VUM.MVector s Int)+    , nbArrays :: !(STRef s (NumArrays a s))+    }++data NumArrays a s = NumArrays+    { naData :: !(VUM.MVector s a)+    , naValid :: !(VUM.MVector s Word8)+    }++type IntBuilder = NumBuilder Int++type DoubleBuilder = NumBuilder Double++-- | New numeric builder with a row-capacity hint and a null sentinel.+newNumBuilder :: (VU.Unbox a) => a -> Int -> ST s (NumBuilder a s)+newNumBuilder nullValue hint = do+    let cap = max 16 hint+    counters <- VUM.replicate 2 0+    dat <- VUM.unsafeNew cap+    val <- VUM.unsafeNew cap+    NumBuilder nullValue counters <$> newSTRef (NumArrays dat val)++newIntBuilder :: Int -> ST s (IntBuilder s)+newIntBuilder = newNumBuilder 0++newDoubleBuilder :: Int -> ST s (DoubleBuilder s)+newDoubleBuilder = newNumBuilder 0++appendNum :: (VU.Unbox a) => NumBuilder a s -> a -> ST s ()+appendNum b !x = do+    n <- VUM.unsafeRead (nbCounters b) cRows+    anyNull <- VUM.unsafeRead (nbCounters b) cAnyNull+    NumArrays dat val <- reserveNum b n+    VUM.unsafeWrite dat n x+    when (anyNull /= 0) $ VUM.unsafeWrite val n 1+    VUM.unsafeWrite (nbCounters b) cRows (n + 1)+{-# INLINE appendNum #-}++appendInt :: IntBuilder s -> Int -> ST s ()+appendInt = appendNum+{-# INLINE appendInt #-}++appendDouble :: DoubleBuilder s -> Double -> ST s ()+appendDouble = appendNum+{-# INLINE appendDouble #-}++-- Fetch the arrays, growing (doubling) first if row @n@ would not fit.+reserveNum :: (VU.Unbox a) => NumBuilder a s -> Int -> ST s (NumArrays a s)+reserveNum b n = do+    arrs <- readSTRef (nbArrays b)+    if n < VUM.length (naData arrs) then pure arrs else growNum b arrs+{-# INLINE reserveNum #-}++growNum ::+    (VU.Unbox a) => NumBuilder a s -> NumArrays a s -> ST s (NumArrays a s)+growNum b (NumArrays dat val) = do+    let cap = VUM.length dat+    dat' <- VUM.unsafeGrow dat cap+    val' <- VUM.unsafeGrow val cap+    let arrs = NumArrays dat' val'+    writeSTRef (nbArrays b) arrs+    pure arrs++instance (Columnable a, VU.Unbox a) => ColumnBuilder (NumBuilder a) where+    appendNull b = do+        n <- VUM.unsafeRead (nbCounters b) cRows+        anyNull <- VUM.unsafeRead (nbCounters b) cAnyNull+        NumArrays dat val <- reserveNum b n+        VUM.unsafeWrite dat n (nbNull b)+        when (anyNull == 0) $ do+            VUM.set (VUM.slice 0 n val) 1+            VUM.unsafeWrite (nbCounters b) cAnyNull 1+        VUM.unsafeWrite val n 0+        VUM.unsafeWrite (nbCounters b) cRows (n + 1)+    {-# INLINE appendNull #-}++    builderLength b = VUM.unsafeRead (nbCounters b) cRows++    freezeBuilder b = do+        n <- VUM.unsafeRead (nbCounters b) cRows+        anyNull <- VUM.unsafeRead (nbCounters b) cAnyNull+        NumArrays dat val <- readSTRef (nbArrays b)+        !vs <- freezeTrimmed n dat+        if anyNull /= 0+            then do+                !bm <- packValidity n val+                pure $! UnboxedColumn (Just bm) vs+            else pure $! UnboxedColumn Nothing vs++-- Zero-copy freeze; copies to exact size when slack exceeds a quarter of n.+freezeTrimmed :: (VU.Unbox a) => Int -> VUM.MVector s a -> ST s (VU.Vector a)+freezeTrimmed n mv+    | VUM.length mv - n <= n `shiftR` 2 = VU.unsafeFreeze (VUM.slice 0 n mv)+    | otherwise = VU.freeze (VUM.slice 0 n mv)++{- | Builder for 'Text' columns. All field bytes go into one exponentially+grown byte array; rows are recorded as offsets, so an append is a memcpy+and freezing slices 'Text' values off the shared array without copying.+-}+data TextBuilder s = TextBuilder+    { tbCounters :: !(VUM.MVector s Int)+    , tbArrays :: !(STRef s (TextArrays s))+    }++data TextArrays s = TextArrays+    { taBytes :: !(A.MArray s)+    , taByteCap :: !Int+    , taOffsets :: !(VUM.MVector s Int)+    -- ^ Row @i@ spans bytes @[offsets!i, offsets!(i+1))@.+    , taValid :: !(VUM.MVector s Word8)+    }++-- | New text builder with row-count and total-byte capacity hints.+newTextBuilder :: Int -> Int -> ST s (TextBuilder s)+newTextBuilder rowHint byteHint = do+    let rcap = max 16 rowHint+        bcap = max 64 byteHint+    counters <- VUM.replicate 3 0+    bytes <- A.new bcap+    offsets <- VUM.unsafeNew (rcap + 1)+    VUM.unsafeWrite offsets 0 0+    val <- VUM.unsafeNew rcap+    TextBuilder counters <$> newSTRef (TextArrays bytes bcap offsets val)++-- | Append @len@ raw bytes at @off@ in @src@ as one field (one memcpy).+appendTextSlice :: TextBuilder s -> A.Array -> Int -> Int -> ST s ()+appendTextSlice b src off len = do+    (n, pos, arrs) <- reserveText b len+    A.copyI len (taBytes arrs) pos src off+    finishTextAppend b arrs n (pos + len)+{-# INLINE appendTextSlice #-}++-- | 'appendTextSlice' from foreign memory (e.g. an mmapped file buffer).+appendTextSliceFromPtr :: TextBuilder s -> Ptr Word8 -> Int -> ST s ()+appendTextSliceFromPtr b ptr len = do+    (n, pos, arrs) <- reserveText b len+    A.copyFromPointer (taBytes arrs) pos ptr len+    finishTextAppend b arrs n (pos + len)+{-# INLINE appendTextSliceFromPtr #-}++-- | Append an already-decoded 'Text' (its bytes are UTF-8 already).+appendText :: TextBuilder s -> T.Text -> ST s ()+appendText b (Text src off len) = appendTextSlice b src off len+{-# INLINE appendText #-}++finishTextAppend :: TextBuilder s -> TextArrays s -> Int -> Int -> ST s ()+finishTextAppend b arrs n endPos = do+    anyNull <- VUM.unsafeRead (tbCounters b) cAnyNull+    when (anyNull /= 0) $ VUM.unsafeWrite (taValid arrs) n 1+    VUM.unsafeWrite (taOffsets arrs) (n + 1) endPos+    VUM.unsafeWrite (tbCounters b) cRows (n + 1)+    VUM.unsafeWrite (tbCounters b) cBytes endPos+{-# INLINE finishTextAppend #-}++reserveText :: TextBuilder s -> Int -> ST s (Int, Int, TextArrays s)+reserveText b extra = do+    n <- VUM.unsafeRead (tbCounters b) cRows+    pos <- VUM.unsafeRead (tbCounters b) cBytes+    arrs <- readSTRef (tbArrays b)+    arrs' <-+        if n < VUM.length (taValid arrs) && pos + extra <= taByteCap arrs+            then pure arrs+            else growText b arrs (n + 1) (pos + extra)+    pure (n, pos, arrs')+{-# INLINE reserveText #-}++growText :: TextBuilder s -> TextArrays s -> Int -> Int -> ST s (TextArrays s)+growText b (TextArrays bytes bcap offsets val) needRows needBytes = do+    let rcap = VUM.length val+    (offsets', val') <-+        if needRows > rcap+            then do+                let rcap' = max (2 * rcap) needRows+                o <- VUM.unsafeGrow offsets (rcap' - rcap)+                v <- VUM.unsafeGrow val (rcap' - rcap)+                pure (o, v)+            else pure (offsets, val)+    (bytes', bcap') <-+        if needBytes > bcap+            then do+                let cap' = max (2 * bcap) needBytes+                bs <- A.resizeM bytes cap'+                pure (bs, cap')+            else pure (bytes, bcap)+    let arrs = TextArrays bytes' bcap' offsets' val'+    writeSTRef (tbArrays b) arrs+    pure arrs++{- | Freeze a 'TextBuilder' into a raw 'TextChunk' for byte-level merging+('mergeTextChunks'): no 'T.Text' values are created until chunks merge.+-}+freezeTextChunk :: TextBuilder s -> ST s TextChunk+freezeTextChunk b = do+    n <- VUM.unsafeRead (tbCounters b) cRows+    anyNull <- VUM.unsafeRead (tbCounters b) cAnyNull+    used <- VUM.unsafeRead (tbCounters b) cBytes+    TextArrays bytes bcap offsets val <- readSTRef (tbArrays b)+    when (used < bcap) (A.shrinkM bytes used)+    arr <- A.unsafeFreeze bytes+    offs <- VU.unsafeFreeze (VUM.slice 0 (n + 1) offsets)+    bm <-+        if anyNull /= 0+            then Just <$> packValidity n val+            else pure Nothing+    pure (TextChunk arr used offs bm)++instance ColumnBuilder TextBuilder where+    appendNull b = do+        (n, pos, arrs) <- reserveText b 0+        anyNull <- VUM.unsafeRead (tbCounters b) cAnyNull+        when (anyNull == 0) $ do+            VUM.set (VUM.slice 0 n (taValid arrs)) 1+            VUM.unsafeWrite (tbCounters b) cAnyNull 1+        VUM.unsafeWrite (taValid arrs) n 0+        VUM.unsafeWrite (taOffsets arrs) (n + 1) pos+        VUM.unsafeWrite (tbCounters b) cRows (n + 1)+        VUM.unsafeWrite (tbCounters b) cBytes pos+    {-# INLINE appendNull #-}++    builderLength b = VUM.unsafeRead (tbCounters b) cRows++    freezeBuilder b = do+        chunk <- freezeTextChunk b+        pure $! mergeTextChunks [chunk]

@@ -0,0 +1,176 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | Concatenation of per-chunk 'Column's (e.g. from parallel CSV chunks). Text+columns merge at the byte level via 'TextChunk' \/ 'mergeTextChunks', so no+per-chunk 'Data.Text.Text' values are ever materialized.+-}+module DataFrame.Internal.ColumnMerge (+    TextChunk (..),+    mergeColumns,+    mergeTextChunks,+    packedFromTextChunk,+    packValidity,+    spliceBitmaps,+    tcRows,+) where++import qualified Data.Text.Array as A+import qualified Data.Vector as VB+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM++import Control.Monad (foldM_, forM_, when)+import Control.Monad.ST (ST, runST)+import Data.Bits (shiftL, shiftR, (.&.), (.|.))+import Data.Maybe (fromMaybe, isNothing)+import Data.Type.Equality (testEquality, (:~:) (Refl))+import Data.Word (Word8)+import DataFrame.Internal.Column (+    Bitmap,+    Column (..),+    Columnable,+    allValidBitmap,+    materializePacked,+ )+import DataFrame.Internal.PackedText (mkPackedContiguous)+import Type.Reflection (typeRep)++{- | A frozen text-builder chunk: raw UTF-8 bytes plus row offsets (row @i@+spans bytes @[offsets!i, offsets!(i+1))@) and an optional validity bitmap.+'Data.Text.Text' values are only created when chunks merge into a 'Column'.+-}+data TextChunk = TextChunk+    { tcBytes :: !A.Array+    , tcUsed :: !Int+    , tcOffsets :: !(VU.Vector Int)+    , tcBitmap :: !(Maybe Bitmap)+    }++tcRows :: TextChunk -> Int+tcRows c = VU.length (tcOffsets c) - 1++{- | Freeze a builder chunk directly into a packed-text column: no+'Data.Text.Text' materialization, no UTF-8 validation pass (deferred to decode).+Not yet called by any reader.+-}+packedFromTextChunk :: TextChunk -> Column+packedFromTextChunk (TextChunk arr _used offs bm) =+    PackedText bm (mkPackedContiguous arr offs)++{- | Merge text chunks into one packed-text 'Column': one byte-array copy per+chunk, one offset rebase, then wrap the shared buffer + offsets as 'PackedText'+(no per-row header, decode deferred).+-}+mergeTextChunks :: [TextChunk] -> Column+mergeTextChunks [] = error "DataFrame.Internal.ColumnMerge.mergeTextChunks: empty list"+mergeTextChunks [c] = packedFromTextChunk c+mergeTextChunks cs = runST $ do+    let totalBytes = sum (map tcUsed cs)+        totalRows = sum (map tcRows cs)+    arr <- A.new (max 1 totalBytes)+    offs <- VUM.unsafeNew (totalRows + 1)+    VUM.unsafeWrite offs 0 0+    let splice !byteBase !rowBase c = do+            let n = tcRows c+                co = tcOffsets c+            A.copyI (tcUsed c) arr byteBase (tcBytes c) 0+            forM_ [1 .. n] $ \i ->+                VUM.unsafeWrite offs (rowBase + i) (byteBase + VU.unsafeIndex co i)+            pure (byteBase + tcUsed c, rowBase + n)+    foldM_ (\(b, r) c -> splice b r c) (0, 0) cs+    farr <- A.unsafeFreeze arr+    foffs <- VU.unsafeFreeze offs+    let !bm = spliceBitmaps [(tcBitmap c, tcRows c) | c <- cs]+    pure (PackedText bm (mkPackedContiguous farr foffs))++{- | Merge per-chunk columns into one column: one allocation + memcpy per+payload, with bitmaps spliced across non-byte-aligned chunk boundaries.+All chunks must have the same element type.+-}+mergeColumns :: [Column] -> Column+mergeColumns [] = error "DataFrame.Internal.ColumnBuilder.mergeColumns: empty list"+mergeColumns [c] = c+mergeColumns cols@(c0 : _) = case c0 of+    PackedText _ _ -> mergeColumns (map materializePacked cols)+    UnboxedColumn _ (_ :: VU.Vector a) ->+        let parts = map (unboxedPart @a) cols+            !merged = VU.concat (map snd parts)+            !bm = spliceBitmaps [(mb, VU.length v) | (mb, v) <- parts]+         in UnboxedColumn bm merged+    BoxedColumn _ (_ :: VB.Vector a) ->+        let parts = map (boxedPart @a) cols+            !merged = VB.concat (map snd parts)+            !bm = spliceBitmaps [(mb, VB.length v) | (mb, v) <- parts]+         in BoxedColumn bm merged++unboxedPart ::+    forall a. (Columnable a, VU.Unbox a) => Column -> (Maybe Bitmap, VU.Vector a)+unboxedPart (UnboxedColumn mb (v :: VU.Vector b)) =+    case testEquality (typeRep @a) (typeRep @b) of+        Just Refl -> (mb, v)+        Nothing -> mergeMismatch+unboxedPart _ = mergeMismatch++boxedPart ::+    forall a. (Columnable a) => Column -> (Maybe Bitmap, VB.Vector a)+boxedPart (BoxedColumn mb (v :: VB.Vector b)) =+    case testEquality (typeRep @a) (typeRep @b) of+        Just Refl -> (mb, v)+        Nothing -> mergeMismatch+boxedPart _ = mergeMismatch++mergeMismatch :: a+mergeMismatch =+    error "DataFrame.Internal.ColumnBuilder.mergeColumns: chunk column types differ"++{- | Splice chunk bitmaps end to end at the bit level. 'Nothing' if no chunk+carries a bitmap; chunks without one count as all-valid otherwise.+-}+spliceBitmaps :: [(Maybe Bitmap, Int)] -> Maybe Bitmap+spliceBitmaps parts+    | all (isNothing . fst) parts = Nothing+    | otherwise = Just $ VU.create $ do+        let total = sum (map snd parts)+            outBytes = (total + 7) `shiftR` 3+        mv <- VUM.replicate outBytes 0+        let orInto i w =+                when (i < outBytes && w /= 0) $ do+                    old <- VUM.unsafeRead mv i+                    VUM.unsafeWrite mv i (old .|. w)+            splice !bitPos (mb, len) = do+                let bm = fromMaybe (allValidBitmap len) mb+                    sh = bitPos .&. 7+                    byte0 = bitPos `shiftR` 3+                    lastIdx = ((len + 7) `shiftR` 3) - 1+                    tailBits = len .&. 7+                    lastMask =+                        if tailBits == 0 then 0xFF else (1 `shiftL` tailBits) - 1+                forM_ [0 .. lastIdx] $ \k -> do+                    let raw = VU.unsafeIndex bm k+                        masked = if k == lastIdx then raw .&. lastMask else raw+                        w = fromIntegral masked :: Word+                    orInto (byte0 + k) (fromIntegral (w `shiftL` sh))+                    when (sh /= 0) $+                        orInto (byte0 + k + 1) (fromIntegral (w `shiftR` (8 - sh)))+                pure (bitPos + len)+        foldM_ splice 0 parts+        pure mv++-- | Pack a 0\/1 byte-per-row validity prefix into a bit-packed 'Bitmap'.+packValidity :: Int -> VUM.MVector s Word8 -> ST s Bitmap+packValidity n val = do+    bytes <- VU.unsafeFreeze (VUM.slice 0 n val)+    let assemble b =+            let base = b `shiftL` 3+                m = min 8 (n - base)+                go !acc !k+                    | k >= m = acc+                    | VU.unsafeIndex bytes (base + k) /= 0 =+                        go (acc .|. (1 `shiftL` k)) (k + 1)+                    | otherwise = go acc (k + 1)+             in go (0 :: Word8) 0+    pure $! VU.generate ((n + 7) `shiftR` 3) assemble

@@ -0,0 +1,362 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module DataFrame.Internal.DataFrame where++import qualified Data.Map as M+import qualified Data.Text as T+import qualified Data.Vector as V+import qualified Data.Vector.Unboxed as VU++import Control.Exception (throw)+import Data.Function (on)+import Data.List (sortBy, (\\))+import Data.Maybe (fromMaybe)+import Data.Type.Equality (+    TestEquality (testEquality),+    type (:~:) (Refl),+    type (:~~:) (HRefl),+ )+import DataFrame.Display.Terminal.PrettyPrint+import DataFrame.Errors+import DataFrame.Internal.Column+import DataFrame.Internal.Expression+import DataFrame.Internal.PackedText (packedIndexText)+import Text.Printf+import Type.Reflection (Typeable, eqTypeRep, typeRep, pattern App)+import Prelude hiding (null)++data DataFrame = DataFrame+    { columns :: V.Vector Column+    -- ^ Column-oriented storage: the frame is a vector of columns.+    , columnIndices :: M.Map T.Text Int+    -- ^ Keeps the column names in the order they were inserted in.+    , dataframeDimensions :: (Int, Int)+    -- ^ (rows, columns)+    , derivingExpressions :: M.Map T.Text UExpr+    }++{- | Force evaluation of all columns in a DataFrame. Replacement for the removed+@instance NFData DataFrame@; used by the IO and lazy-executor strict paths.+-}+forceDataFrame :: DataFrame -> DataFrame+forceDataFrame df@(DataFrame cols idx dims _exprs) =+    V.foldl' (\() c -> forceColumn c) () cols `seq` idx `seq` dims `seq` df++{- | A record that contains information about how and what+rows are grouped in the dataframe. This can only be used with+`aggregate`.+-}+data GroupedDataFrame = Grouped+    { fullDataframe :: DataFrame+    , 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 _rtg) =+        printf+            "{ keyColumns: %s groupedColumns: %s }"+            (show cols)+            (show (M.keys (columnIndices df) \\ cols))++instance Eq GroupedDataFrame where+    (==) (Grouped df cols _indices _os _rtg) (Grouped df' cols' _indices' _os' _rtg') = (df == df') && (cols == cols')++instance Eq DataFrame where+    (==) :: DataFrame -> DataFrame -> Bool+    a == b =+        M.keys (columnIndices a) == M.keys (columnIndices b)+            && foldr+                ( \(name, index) acc -> acc && (columns a V.!? index == (columns b V.!? (columnIndices b M.! name)))+                )+                True+                (M.toList $ columnIndices a)++instance Show DataFrame where+    show :: DataFrame -> String+    show d =+        let (r, _) = dataframeDimensions d+            cfg = defaultTruncateConfig+            shown = if maxRows cfg > 0 then min (maxRows cfg) r else r+            body = asTextWith Plain (Just cfg) d+            footer+                | shown < r =+                    "\nShowing "+                        <> T.pack (show shown)+                        <> " rows out of "+                        <> T.pack (show r)+                | otherwise = T.empty+         in T.unpack (body <> footer)++{- | Configures how a 'DataFrame' is rendered as text: 'maxRows' caps rendered+rows, 'maxColumns' collapses middle columns past the limit into an ellipsis, and+'maxCellWidth' truncates long cells. A non-positive field means \"no limit\".+-}+data TruncateConfig = TruncateConfig+    { maxRows :: Int+    , maxColumns :: Int+    , maxCellWidth :: Int+    }+    deriving (Show, Eq)++-- | Sensible defaults for GHCi: 20 rows, 10 columns, 30 characters per cell.+defaultTruncateConfig :: TruncateConfig+defaultTruncateConfig =+    TruncateConfig{maxRows = 20, maxColumns = 10, maxCellWidth = 30}++-- | Ellipsis character used to mark elided columns and clipped cells.+ellipsisText :: T.Text+ellipsisText = "\x2026"++-- | For showing the dataframe as markdown in notebooks.+toMarkdown :: DataFrame -> T.Text+toMarkdown = asText Markdown++-- | For showing the dataframe as a string markdown in notebooks.+toMarkdown' :: DataFrame -> String+toMarkdown' = T.unpack . toMarkdown++asText :: RenderFormat -> DataFrame -> T.Text+asText fmt = asTextWith fmt Nothing++asTextWith :: RenderFormat -> Maybe TruncateConfig -> DataFrame -> T.Text+asTextWith fmt mTrunc d =+    let allHeaders =+            map fst (sortBy (compare `on` snd) (M.toList (columnIndices d)))+        nCols = length allHeaders+        (totalRows, _) = dataframeDimensions d++        rowCap = case mTrunc of+            Just cfg | maxRows cfg > 0 -> min totalRows (maxRows cfg)+            _ -> totalRows++        (visibleHeaders, ellipsisAt) = pickColumns mTrunc nCols allHeaders++        lookupCol name =+            fmap+                (takeColumn rowCap)+                ((V.!?) (columns d) ((M.!) (columnIndices d) name))+        survivingCols = map lookupCol visibleHeaders+        survivingTypes = map (maybe "" getType) survivingCols+        survivingData = map get survivingCols++        clipCell = case mTrunc of+            Just cfg | maxCellWidth cfg > 0 -> truncateCell (maxCellWidth cfg)+            _ -> id++        (finalHeaders, finalTypes, finalCols) = case ellipsisAt of+            Nothing -> (visibleHeaders, survivingTypes, survivingData)+            Just i ->+                let ellipsisCol = V.replicate rowCap ellipsisText+                 in ( insertAt i ellipsisText visibleHeaders+                    , insertAt i ellipsisText survivingTypes+                    , insertAt i ellipsisCol survivingData+                    )++        getType :: Column -> T.Text+        showMaybeType :: forall a. (Typeable a) => String+        showMaybeType =+            let s = show (typeRep @a)+             in "Maybe " <> if ' ' `elem` s then "(" <> s <> ")" else s+        getType (BoxedColumn Nothing (_ :: V.Vector a)) = T.pack $ show (typeRep @a)+        getType (BoxedColumn (Just _) (_ :: V.Vector a)) = T.pack $ showMaybeType @a+        getType (UnboxedColumn Nothing (_ :: VU.Vector a)) = T.pack $ show (typeRep @a)+        getType (UnboxedColumn (Just _) (_ :: VU.Vector a)) = T.pack $ showMaybeType @a+        getType (PackedText Nothing _) = T.pack $ show (typeRep @T.Text)+        getType (PackedText (Just _) _) = T.pack $ showMaybeType @T.Text++        get :: Maybe Column -> V.Vector T.Text+        get (Just (BoxedColumn (Just bm) (column :: V.Vector a))) =+            V.generate (V.length column) $ \i ->+                if bitmapTestBit bm i+                    then T.pack (show (Just (V.unsafeIndex column i)))+                    else "Nothing"+        get (Just (BoxedColumn Nothing (column :: V.Vector a))) =+            case testEquality (typeRep @a) (typeRep @T.Text) of+                Just Refl -> column+                Nothing -> case testEquality (typeRep @a) (typeRep @String) of+                    Just Refl -> V.map T.pack column+                    Nothing -> V.map (T.pack . show) column+        get (Just (UnboxedColumn (Just bm) column)) =+            V.generate (VU.length column) $ \i ->+                if bitmapTestBit bm i+                    then T.pack (show (Just (VU.unsafeIndex column i)))+                    else "Nothing"+        get (Just (UnboxedColumn Nothing column)) =+            V.generate (VU.length column) (T.pack . show . VU.unsafeIndex column)+        get (Just c@(PackedText _ _)) = get (Just (materializePacked c))+        get Nothing = V.empty+     in showTable+            fmt+            (map clipCell finalHeaders)+            (map clipCell finalTypes)+            (map (V.map clipCell) finalCols)++{- | Decide which columns survive horizontal truncation and where to splice the+ellipsis column. Splits with the extra column on the left for odd 'maxColumns';+inserts the ellipsis only when it actually saves space.+-}+pickColumns ::+    Maybe TruncateConfig ->+    Int ->+    [a] ->+    ([a], Maybe Int)+pickColumns mTrunc nCols xs = case mTrunc of+    Just cfg+        | let c = maxColumns cfg+        , c > 0+        , nCols > c + 1 ->+            let leftN = (c + 1) `div` 2+                rightN = c - leftN+             in ( Prelude.take leftN xs ++ Prelude.drop (nCols - rightN) xs+                , Just leftN+                )+    _ -> (xs, Nothing)++-- | Splice @x@ into @xs@ at index @i@ (0-based), shifting later elements right.+insertAt :: Int -> a -> [a] -> [a]+insertAt i x xs = let (l, r) = splitAt i xs in l ++ x : r++-- | Cap a single cell's rendered length, appending an ellipsis when shortened.+truncateCell :: Int -> T.Text -> T.Text+truncateCell n t+    | n <= 0 = t+    | T.compareLength t n /= GT = t+    | n == 1 = ellipsisText+    | otherwise = T.take (n - 1) t <> ellipsisText++-- | O(1) Creates an empty dataframe+empty :: DataFrame+empty =+    DataFrame+        { columns = V.empty+        , columnIndices = M.empty+        , dataframeDimensions = (0, 0)+        , derivingExpressions = M.empty+        }++-- | O(k) Get column names of the DataFrame in order of insertion.+columnNames :: DataFrame -> [T.Text]+columnNames = map fst . sortBy (compare `on` snd) . M.toList . columnIndices+{-# INLINE columnNames #-}++{- | Insert a column into a DataFrame. If a column with the same name already+exists it is replaced in-place; otherwise the column is appended at the end.+Other columns are expanded (padded with nulls) to match the new row count.+-}+insertColumn :: T.Text -> Column -> DataFrame -> DataFrame+insertColumn name column d =+    let+        (r, c) = dataframeDimensions d+        n = max (columnLength column) r+        exprs = M.delete name (derivingExpressions d)+     in+        case M.lookup name (columnIndices d) of+            Just i ->+                DataFrame+                    (V.map (expandColumn n) (columns d V.// [(i, column)]))+                    (columnIndices d)+                    (n, c)+                    exprs+            Nothing ->+                DataFrame+                    (V.map (expandColumn n) (columns d `V.snoc` column))+                    (M.insert name c (columnIndices d))+                    (n, c + 1)+                    exprs++-- | Build a DataFrame from a list of @(name, column)@ pairs using 'insertColumn'.+fromNamedColumns :: [(T.Text, Column)] -> DataFrame+fromNamedColumns = foldl (\df (name, column) -> insertColumn name column df) empty++{- | Safely retrieves a column by name from the dataframe.++Returns 'Nothing' if the column does not exist.++==== __Examples__++>>> getColumn "age" df+Just (UnboxedColumn ...)++>>> getColumn "nonexistent" df+Nothing+-}+getColumn :: T.Text -> DataFrame -> Maybe Column+getColumn name df+    | null df = Nothing+    | otherwise = do+        i <- columnIndices df M.!? name+        columns df V.!? i++{- | Retrieve a column by name, throwing 'ColumnsNotFoundException' if it does not+exist. Unsafe version of 'getColumn'; use when the column is certain to exist.+-}+unsafeGetColumn :: T.Text -> DataFrame -> Column+unsafeGetColumn name df = case getColumn name df of+    Nothing -> throw $ ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)+    Just col -> col++{- | Checks if the dataframe is empty (has no columns).++Returns 'True' if the dataframe has no columns, 'False' otherwise.+Note that a dataframe with columns but no rows is not considered null.+-}+null :: DataFrame -> Bool+null df = V.null (columns df)++-- | Convert a DataFrame to a CSV (comma-separated) text.+toCsv :: DataFrame -> T.Text+toCsv = toSeparated ','++-- | Convert a DataFrame to a CSV (comma-separated) string.+toCsv' :: DataFrame -> String+toCsv' = T.unpack . toSeparated ','++-- | Convert a DataFrame to a text representation with a custom separator.+toSeparated :: Char -> DataFrame -> T.Text+toSeparated sep df+    | null df = T.empty+    | otherwise =+        let (rows, _) = dataframeDimensions df+            headers = map fst (sortBy (compare `on` snd) (M.toList (columnIndices df)))+            sepText = T.singleton sep+            headerLine = T.intercalate sepText headers+            dataLines = map (T.intercalate sepText . getRowAsText df) [0 .. rows - 1]+         in T.unlines (headerLine : dataLines)++getRowAsText :: DataFrame -> Int -> [T.Text]+getRowAsText df i = map (`showElement` i) (V.toList (columns df))++showElement :: Column -> Int -> T.Text+showElement (BoxedColumn _ (c :: V.Vector a)) i = case c V.!? i of+    Nothing -> error $ "Column index out of bounds at row " ++ show i+    Just e+        | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) -> e+        | App t1 t2 <- typeRep @a+        , Just HRefl <- eqTypeRep t1 (typeRep @Maybe) ->+            case testEquality t2 (typeRep @T.Text) of+                Just Refl -> fromMaybe "null" e+                Nothing -> stripJust (T.pack (show e))+        | otherwise -> T.pack (show e)+showElement (UnboxedColumn _ c) i = case c VU.!? i of+    Nothing -> error $ "Column index out of bounds at row " ++ show i+    Just e -> T.pack (show e)+showElement (PackedText bm p) i = case bm of+    Just b | not (bitmapTestBit b i) -> "null"+    _ -> packedIndexText p i++stripJust :: T.Text -> T.Text+stripJust = fromMaybe "null" . T.stripPrefix "Just "

@@ -0,0 +1,129 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | Dictionary-encode a text (or factor) group key to dense @Int@ codes: each row+gets a first-appearance code @0..card-1@ (NULL reserved) plus the cardinality. A+tested building block; profiled slower than the hash group-by, so unused for now.+-}+module DataFrame.Internal.DictEncode (+    dictEncodeColumn,+    dictEncodeColumnUpTo,+    dictMaxCardinality,+) where++import Control.Monad.ST (runST)+import qualified Data.Text as T+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))+import qualified Data.Vector as V+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import Type.Reflection (typeRep)++import DataFrame.Internal.Column (Bitmap, Column (..), bitmapTestBit)+import DataFrame.Internal.Hash (fnvOffset, mixBytes, mixText, nullSalt)+import DataFrame.Internal.HashTable (htInsert, newHashTable)+import DataFrame.Internal.PackedText (+    PackedTextData,+    packedLength,+    packedSlice,+    sliceEqBytes,+ )++{- | Largest distinct-value count we will dictionary-encode. Above this the codes+no longer index a reasonable direct accumulator and the encode pass is pure+overhead, so the caller keeps the plain hash group-by.+-}+dictMaxCardinality :: Int+dictMaxCardinality = 1048576++{- | Dictionary-encode a text-like column to dense first-appearance @Int@ codes,+returning @Just (codes, cardinality)@ (a NULL row gets its own reserved code).+'Nothing' for non-text columns or cardinality above 'dictMaxCardinality'.+-}+dictEncodeColumn :: Column -> Maybe (VU.Vector Int, Int)+dictEncodeColumn = dictEncodeColumnUpTo dictMaxCardinality++{- | Dictionary-encode like 'dictEncodeColumn' but bail to 'Nothing' as soon as+the distinct count would exceed @maxCard@, letting a low-cardinality probe avoid+a full high-cardinality pass.+-}+dictEncodeColumnUpTo :: Int -> Column -> Maybe (VU.Vector Int, Int)+dictEncodeColumnUpTo maxCard (PackedText bm p) = encodePacked maxCard bm p+dictEncodeColumnUpTo maxCard (BoxedColumn bm (v :: V.Vector a)) =+    case testEquality (typeRep @a) (typeRep @T.Text) of+        Just Refl -> encodeBoxedText maxCard bm v+        Nothing -> Nothing+dictEncodeColumnUpTo _ _ = Nothing++{- | Encode a packed-text column: hash each row's raw UTF-8 bytes (the grouping+'mixBytes'), re-verify byte equality on collisions, assign dense codes in+first-appearance order. A null row hashes 'nullSalt'.+-}+encodePacked ::+    Int -> Maybe Bitmap -> PackedTextData -> Maybe (VU.Vector Int, Int)+encodePacked maxCard bm p =+    let !n = packedLength p+        valid i = case bm of+            Just b -> bitmapTestBit b i+            Nothing -> True+        hashAt i =+            if valid i+                then let (arr, o, l) = packedSlice p i in mixBytes fnvOffset arr o l+                else nullSalt+        eqAt a b =+            case (valid a, valid b) of+                (True, True) ->+                    let (arrA, oA, lA) = packedSlice p a+                        (arrB, oB, lB) = packedSlice p b+                     in sliceEqBytes arrA oA lA arrB oB lB+                (False, False) -> True+                _ -> False+     in buildCodes maxCard n hashAt eqAt++{- | Encode a boxed 'Data.Text.Text' column, mirroring 'encodePacked' but over+boxed values (used when a user-built Text column is grouped).+-}+encodeBoxedText ::+    Int -> Maybe Bitmap -> V.Vector T.Text -> Maybe (VU.Vector Int, Int)+encodeBoxedText maxCard bm v =+    let !n = V.length v+        valid i = case bm of+            Just b -> bitmapTestBit b i+            Nothing -> True+        hashAt i =+            if valid i then mixText fnvOffset (V.unsafeIndex v i) else nullSalt+        eqAt a b =+            case (valid a, valid b) of+                (True, True) -> V.unsafeIndex v a == V.unsafeIndex v b+                (False, False) -> True+                _ -> False+     in buildCodes maxCard n hashAt eqAt++{- | The shared code-assignment loop: bucket every row through an open-addressing+table on its precomputed hash, re-verify with @eqAt@ on a hit, assign dense+first-appearance codes. Bails to 'Nothing' once the distinct count exceeds @maxCard@.+-}+buildCodes ::+    Int -> Int -> (Int -> Int) -> (Int -> Int -> Bool) -> Maybe (VU.Vector Int, Int)+buildCodes maxCard n hashAt eqAt+    | n == 0 = Just (VU.empty, 0)+    | otherwise = runST $ do+        ht <- newHashTable (min n (maxCard + 1))+        codes <- VUM.new n+        let go !i !next+                | i >= n = pure (Just next)+                | next > maxCard = pure Nothing+                | otherwise = do+                    let !h = hashAt i+                    (code, isNew) <- htInsert ht eqAt next i h+                    VUM.unsafeWrite codes i code+                    go (i + 1) (if isNew then next + 1 else next)+        mres <- go 0 0+        case mres of+            Nothing -> pure Nothing+            Just card -> do+                frozen <- VU.unsafeFreeze codes+                pure (Just (frozen, card))

@@ -0,0 +1,511 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DisambiguateRecordFields #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE NoFieldSelectors #-}++module DataFrame.Internal.Expression where++import qualified Data.Map.Strict as M+import Data.Maybe (fromMaybe)+import Data.String+import qualified Data.Text as T+import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))+import qualified Data.Vector.Generic as VG+import DataFrame.Internal.Column+import qualified DataFrame.Internal.Pretty as P+import Type.Reflection (Typeable, typeOf, typeRep)++{- | Operators are an open typeclass: built-ins get their own 'Typeable' type so the+simplifier can match them by 'cast', and users can add instances. The generic+'UnUDF'/'BinUDF' carriers cover UDFs, dynamic-named, and arithmetic ops.+-}+class (Typeable op) => UnaryOp op where+    unaryFn :: op a b -> a -> b+    unaryName :: op a b -> T.Text+    unarySymbol :: op a b -> Maybe T.Text+    unarySymbol _ = Nothing++class (Typeable op) => BinaryOp op where+    binaryFn :: op a b c -> a -> b -> c+    binaryName :: op a b c -> T.Text+    binarySymbol :: op a b c -> Maybe T.Text+    binarySymbol _ = Nothing+    binaryCommutative :: op a b c -> Bool+    binaryCommutative _ = False+    binaryPrecedence :: op a b c -> Int+    binaryPrecedence _ = 9++data UnUDF a b = MkUnaryOp+    { unaryFn :: a -> b+    , unaryName :: T.Text+    , unarySymbol :: Maybe T.Text+    }++data BinUDF a b c = MkBinaryOp+    { binaryFn :: a -> b -> c+    , binaryName :: T.Text+    , binarySymbol :: Maybe T.Text+    , binaryCommutative :: Bool+    , binaryPrecedence :: Int+    }++instance UnaryOp UnUDF where+    unaryFn (MkUnaryOp{unaryFn = f}) = f+    unaryName (MkUnaryOp{unaryName = n}) = n+    unarySymbol (MkUnaryOp{unarySymbol = s}) = s++instance BinaryOp BinUDF where+    binaryFn (MkBinaryOp{binaryFn = f}) = f+    binaryName (MkBinaryOp{binaryName = n}) = n+    binarySymbol (MkBinaryOp{binarySymbol = s}) = s+    binaryCommutative (MkBinaryOp{binaryCommutative = c}) = c+    binaryPrecedence (MkBinaryOp{binaryPrecedence = p}) = p++data MeanAcc = MeanAcc {-# UNPACK #-} !Double {-# UNPACK #-} !Int+    deriving (Show, Eq, Ord, Read)++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+    CastWith ::+        (Columnable a, Columnable b, Read a) =>+        T.Text ->+        T.Text ->+        (Either String a -> b) ->+        Expr b+    CastExprWith ::+        (Columnable a, Columnable b, Columnable src, Read a) =>+        T.Text ->+        (Either String a -> b) ->+        Expr src ->+        Expr b+    Lit :: (Columnable a) => a -> Expr a+    Unary ::+        (UnaryOp op, Columnable a, Columnable b) => op b a -> Expr b -> Expr a+    Binary ::+        (BinaryOp op, Columnable c, Columnable b, Columnable a) =>+        op c b a -> Expr c -> Expr b -> Expr a+    If :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a+    Agg :: (Columnable a, Columnable b) => AggStrategy a b -> Expr b -> Expr a+    Over :: (Columnable a) => [T.Text] -> Expr a -> Expr a++data UExpr where+    UExpr :: (Columnable a) => Expr a -> UExpr++instance Show UExpr where+    show :: UExpr -> String+    show (UExpr expr) = show expr++type NamedExpr = (T.Text, UExpr)++instance (Num a, Columnable a) => Num (Expr a) where+    (+) :: Expr a -> Expr a -> Expr a+    (+) =+        Binary+            ( MkBinaryOp+                { binaryFn = (+)+                , binaryName = "add"+                , binarySymbol = Just "+"+                , binaryCommutative = True+                , binaryPrecedence = 6+                }+            )++    (-) :: Expr a -> Expr a -> Expr a+    (-) =+        Binary+            ( MkBinaryOp+                { binaryFn = (-)+                , binaryName = "sub"+                , binarySymbol = Just "-"+                , binaryCommutative = False+                , binaryPrecedence = 6+                }+            )++    (*) :: Expr a -> Expr a -> Expr a+    (*) =+        Binary+            ( MkBinaryOp+                { binaryFn = (*)+                , binaryName = "mult"+                , binarySymbol = Just "*"+                , binaryCommutative = True+                , binaryPrecedence = 7+                }+            )++    fromInteger :: Integer -> Expr a+    fromInteger = Lit . fromInteger++    negate :: Expr a -> Expr a+    negate =+        Unary+            (MkUnaryOp{unaryFn = negate, unaryName = "negate", unarySymbol = Nothing})++    abs :: (Num a) => Expr a -> Expr a+    abs = Unary (MkUnaryOp{unaryFn = abs, unaryName = "abs", unarySymbol = Nothing})++    signum :: (Num a) => Expr a -> Expr a+    signum =+        Unary+            (MkUnaryOp{unaryFn = signum, unaryName = "signum", unarySymbol = Nothing})++add :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a+add = (+)++sub :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a+sub = (-)++mult :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a+mult = (*)++instance (Fractional a, Columnable a) => Fractional (Expr a) where+    fromRational :: (Fractional a, Columnable a) => Rational -> Expr a+    fromRational = Lit . fromRational++    (/) :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a+    (/) =+        Binary+            ( MkBinaryOp+                { binaryFn = (/)+                , binaryName = "divide"+                , binarySymbol = Just "/"+                , binaryCommutative = False+                , binaryPrecedence = 7+                }+            )++divide :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a+divide = (/)++instance (IsString a, Columnable a) => IsString (Expr a) where+    fromString :: String -> Expr a+    fromString s = Lit (fromString s)++instance (Floating a, Columnable a) => Floating (Expr a) where+    pi :: (Floating a, Columnable a) => Expr a+    pi = Lit pi+    exp :: (Floating a, Columnable a) => Expr a -> Expr a+    exp = Unary (MkUnaryOp{unaryFn = exp, unaryName = "exp", unarySymbol = Nothing})+    sqrt :: (Floating a, Columnable a) => Expr a -> Expr a+    sqrt =+        Unary (MkUnaryOp{unaryFn = sqrt, unaryName = "sqrt", unarySymbol = Nothing})+    (**) :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a+    (**) =+        Binary+            ( MkBinaryOp+                { binaryFn = (**)+                , binaryName = "exponentiate"+                , binarySymbol = Just "**"+                , binaryCommutative = False+                , binaryPrecedence = 8+                }+            )+    log :: (Floating a, Columnable a) => Expr a -> Expr a+    log = Unary (MkUnaryOp{unaryFn = log, unaryName = "log", unarySymbol = Nothing})+    logBase :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a+    logBase =+        Binary+            ( MkBinaryOp+                { binaryFn = logBase+                , binaryName = "logBase"+                , binarySymbol = Nothing+                , binaryCommutative = False+                , binaryPrecedence = 1+                }+            )+    sin :: (Floating a, Columnable a) => Expr a -> Expr a+    sin = Unary (MkUnaryOp{unaryFn = sin, unaryName = "sin", unarySymbol = Nothing})+    cos :: (Floating a, Columnable a) => Expr a -> Expr a+    cos = Unary (MkUnaryOp{unaryFn = cos, unaryName = "cos", unarySymbol = Nothing})+    tan :: (Floating a, Columnable a) => Expr a -> Expr a+    tan = Unary (MkUnaryOp{unaryFn = tan, unaryName = "tan", unarySymbol = Nothing})+    asin :: (Floating a, Columnable a) => Expr a -> Expr a+    asin =+        Unary (MkUnaryOp{unaryFn = asin, unaryName = "asin", unarySymbol = Nothing})+    acos :: (Floating a, Columnable a) => Expr a -> Expr a+    acos =+        Unary (MkUnaryOp{unaryFn = acos, unaryName = "acos", unarySymbol = Nothing})+    atan :: (Floating a, Columnable a) => Expr a -> Expr a+    atan =+        Unary (MkUnaryOp{unaryFn = atan, unaryName = "atan", unarySymbol = Nothing})+    sinh :: (Floating a, Columnable a) => Expr a -> Expr a+    sinh =+        Unary (MkUnaryOp{unaryFn = sinh, unaryName = "sinh", unarySymbol = Nothing})+    cosh :: (Floating a, Columnable a) => Expr a -> Expr a+    cosh =+        Unary (MkUnaryOp{unaryFn = cosh, unaryName = "cosh", unarySymbol = Nothing})+    asinh :: (Floating a, Columnable a) => Expr a -> Expr a+    asinh =+        Unary+            (MkUnaryOp{unaryFn = asinh, unaryName = "asinh", unarySymbol = Nothing})+    acosh :: (Floating a, Columnable a) => Expr a -> Expr a+    acosh =+        Unary+            (MkUnaryOp{unaryFn = acosh, unaryName = "acosh", unarySymbol = Nothing})+    atanh :: (Floating a, Columnable a) => Expr a -> Expr a+    atanh =+        Unary+            (MkUnaryOp{unaryFn = atanh, unaryName = "atanh", unarySymbol = Nothing})++instance (Show a) => Show (Expr a) where+    show :: Expr a -> String+    show (Col name) = "(col @" ++ show (typeRep @a) ++ " " ++ show name ++ ")"+    show (CastWith name tag _) = "(castWith " ++ show tag ++ " " ++ show name ++ ")"+    show (CastExprWith tag _ inner) = "(castExprWith " ++ show tag ++ " " ++ show inner ++ ")"+    show (Lit value) = "(lit (" ++ show value ++ "))"+    show (If cond l r) = "(ifThenElse " ++ show cond ++ " " ++ show l ++ " " ++ show r ++ ")"+    show (Unary op value) = "(" ++ T.unpack (unaryName op) ++ " " ++ show value ++ ")"+    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 ++ ")"+    show (Over keys inner) = "(over " ++ show keys ++ " " ++ show inner ++ ")"++normalize :: (Show a, Typeable a) => Expr a -> Expr a+normalize expr = case expr of+    Col name -> Col name+    CastWith n t f -> CastWith n t f+    CastExprWith t f e -> CastExprWith t f (normalize e)+    Lit val -> Lit val+    If cond th el -> If (normalize cond) (normalize th) (normalize el)+    Unary op e -> Unary op (normalize e)+    Binary op e1 e2+        | binaryCommutative op ->+            let n1 = normalize e1+                n2 = normalize e2+             in case testEquality (typeOf n1) (typeOf n2) of+                    Nothing -> expr+                    Just Refl ->+                        if compareExpr n1 n2 == GT+                            then Binary op n2 n1+                            else Binary op n1 n2+        | otherwise -> Binary op (normalize e1) (normalize e2)+    Agg strat e -> Agg strat (normalize e)+    Over keys inner -> Over keys (normalize inner)++-- Compare expressions for ordering (used in normalization)+compareExpr :: Expr a -> Expr a -> Ordering+compareExpr e1 e2 = compare (exprKey e1) (exprKey e2)+  where+    exprKey :: Expr a -> String+    exprKey (Col name) = "0:" ++ T.unpack name+    exprKey (CastWith name tag _) = "0CW:" ++ T.unpack name ++ ":" ++ T.unpack tag+    exprKey (CastExprWith tag _ _) = "0CE:" ++ T.unpack tag+    exprKey (Lit val) = "1:" ++ show val+    exprKey (If c t e) = "2:" ++ exprKey c ++ exprKey t ++ exprKey e+    exprKey (Unary op e) = "3:" ++ T.unpack (unaryName op) ++ exprKey e+    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+    exprKey (Over keys e) = "6:over:" ++ show keys ++ exprKey e++eqExpr :: forall a. (Columnable a) => Expr a -> Expr a -> Bool+eqExpr l r = eqNormalized (normalize l) (normalize r)+  where+    exprEq :: (Columnable b, Columnable c) => Expr b -> Expr c -> Bool+    exprEq e1 e2 = case testEquality (typeOf e1) (typeOf e2) of+        Just Refl -> eqExpr e1 e2+        Nothing -> False+    eqNormalized :: Expr a -> Expr a -> Bool+    eqNormalized (Col n1) (Col n2) = n1 == n2+    eqNormalized (CastWith n1 t1 _) (CastWith n2 t2 _) = n1 == n2 && t1 == t2+    eqNormalized (CastExprWith t1 _ e1) (CastExprWith t2 _ e2) = t1 == t2 && e1 `exprEq` e2+    eqNormalized (Lit v1) (Lit v2) = v1 == v2+    eqNormalized (If c1 t1 e1) (If c2 t2 e2) =+        eqExpr c1 c2 && t1 `exprEq` t2 && e1 `exprEq` e2+    eqNormalized (Unary op1 e1) (Unary op2 e2) = unaryName op1 == unaryName op2 && e1 `exprEq` e2+    eqNormalized (Binary op1 e1a e1b) (Binary op2 e2a e2b) = binaryName op1 == binaryName op2 && e1a `exprEq` e2a && e1b `exprEq` e2b+    eqNormalized (Agg (CollectAgg n1 _) e1) (Agg (CollectAgg n2 _) e2) =+        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 (Over k1 e1) (Over k2 e2) = k1 == k2 && e1 `exprEq` e2+    eqNormalized _ _ = False++replaceExpr ::+    forall a b c.+    (Columnable a, Columnable b, Columnable c) =>+    Expr a -> Expr b -> Expr c -> Expr c+replaceExpr new old expr = case testEquality (typeRep @b) (typeRep @c) of+    Just Refl -> case testEquality (typeRep @a) (typeRep @c) of+        Just Refl -> if eqExpr old expr then new else replace'+        Nothing -> expr+    Nothing -> replace'+  where+    replace' = case expr of+        (Col _) -> expr+        (CastWith{}) -> expr+        (CastExprWith t f e) -> CastExprWith t f (replaceExpr new old e)+        (Lit _) -> expr+        (If cond l r) ->+            If (replaceExpr new old cond) (replaceExpr new old l) (replaceExpr new old r)+        (Unary op value) -> Unary op (replaceExpr new old value)+        (Binary op l r) -> Binary op (replaceExpr new old l) (replaceExpr new old r)+        (Agg op inner) -> Agg op (replaceExpr new old inner)+        (Over keys inner) -> Over keys (replaceExpr new old inner)++{- | Simultaneously substitute 'Col' references from a name→expression map in a+single parallel pass, so a swap like @{a ↦ col b, b ↦ col a}@ works. Raw-text+references (in 'CastWith', 'Over' keys) are left untouched; type mismatch raises.+-}+substituteColumns ::+    forall a. (Columnable a) => M.Map T.Text UExpr -> Expr a -> Expr a+substituteColumns subs = go+  where+    go :: forall b. (Columnable b) => Expr b -> Expr b+    go e@(Col name) = case M.lookup name subs of+        Nothing -> e+        Just (UExpr (repl :: Expr c)) -> case testEquality (typeRep @b) (typeRep @c) of+            Just Refl -> repl+            Nothing ->+                error $+                    "substituteColumns: type mismatch for column "+                        ++ show name+                        ++ "; column has type "+                        ++ show (typeRep @b)+                        ++ " but replacement has type "+                        ++ show (typeRep @c)+    go e@(CastWith{}) = e+    go (CastExprWith t f e) = CastExprWith t f (go e)+    go e@(Lit _) = e+    go (If cond l r) = If (go cond) (go l) (go r)+    go (Unary op value) = Unary op (go value)+    go (Binary op l r) = Binary op (go l) (go r)+    go (Agg op inner) = Agg op (go inner)+    go (Over keys inner) = Over keys (go inner)++eSize :: Expr a -> Int+eSize (Col _) = 1+eSize (CastWith{}) = 1+eSize (CastExprWith _ _ e) = 1 + eSize e+eSize (Lit _) = 1+eSize (If c l r) = 1 + eSize c + eSize l + eSize r+eSize (Unary _ e) = 1 + eSize e+eSize (Binary _ l r) = 1 + eSize l + eSize r+eSize (Agg _strategy expr) = eSize expr + 1+eSize (Over _ inner) = 1 + eSize inner++getColumns :: Expr a -> [T.Text]+getColumns (Col cName) = [cName]+getColumns (CastWith name _ _) = [name]+getColumns (CastExprWith _ _ e) = getColumns e+getColumns _expr@(Lit _) = []+getColumns (If cond l r) = getColumns cond <> getColumns l <> getColumns r+getColumns (Unary _op value) = getColumns value+getColumns (Binary _op l r) = getColumns l <> getColumns r+getColumns (Agg _strategy expr) = getColumns expr+getColumns (Over keys inner) = keys <> getColumns inner++{- | Render an expression as readable, width-aware pseudo-code at the default+width ('P.defaultWidth'). See 'prettyPrintWidth' to control wrapping.+-}+prettyPrint :: Expr a -> String+prettyPrint = prettyPrintWidth P.defaultWidth++{- | Render an expression as readable, width-aware pseudo-code: long binary chains+wrap onto aligned continuation lines, @if@/@then@/@else@ break onto their own lines+(nested @else if@ form a flat ladder), and sub-exprs are parenthesized by precedence.+-}+prettyPrintWidth :: Int -> Expr a -> String+prettyPrintWidth width = P.render width . toDoc 0+  where+    toDoc :: Int -> Expr x -> P.Doc+    toDoc prec expr = case expr of+        Col name -> P.text (T.unpack name)+        CastWith name _ _ -> P.text (T.unpack name)+        CastExprWith tag _ inner -> P.text (T.unpack tag) <> P.parens (toDoc 0 inner)+        Lit value -> P.text (show value)+        If{} -> renderIf prec expr+        Unary op arg ->+            let fn = fromMaybe (unaryName op) (unarySymbol op)+             in P.text (T.unpack fn) <> P.parens (toDoc 0 arg)+        Binary op l r -> case binarySymbol op of+            Just sym -> renderBinary prec op (T.unpack sym) l r+            Nothing ->+                P.text (T.unpack (binaryName op))+                    <> P.parens (toDoc 0 l <> P.text ", " <> toDoc 0 r)+        Agg (CollectAgg op _) arg -> P.text (T.unpack op) <> P.parens (toDoc 0 arg)+        Agg (FoldAgg op _ _) arg -> P.text (T.unpack op) <> P.parens (toDoc 0 arg)+        Agg (MergeAgg op _ _ _ _) arg -> P.text (T.unpack op) <> P.parens (toDoc 0 arg)+        Over keys inner ->+            toDoc 0 inner <> P.text (".over(" ++ show (map T.unpack keys) ++ ")")++    renderBinary ::+        (BinaryOp op) => Int -> op c b a -> String -> Expr c -> Expr b -> P.Doc+    renderBinary prec op sym l r =+        let p = binaryPrecedence op+            body+                | binaryCommutative op =+                    let operands =+                            flattenChain (binaryName op) p l+                                ++ flattenChain (binaryName op) p r+                     in case operands of+                            (o0 : os@(_ : _)) ->+                                P.group+                                    ( o0+                                        <> P.nest+                                            2+                                            (P.hcat [P.line <> P.text sym P.<+> o | o <- os])+                                    )+                            _ -> toDoc p l P.<+> P.text sym P.<+> toDoc p r+                | otherwise =+                    P.group (toDoc p l <> P.nest 2 (P.line <> P.text sym P.<+> toDoc p r))+         in if prec > p+                then P.parens body+                else if prec >= 1 && prec < p then P.parensWhenBroken body else body++    flattenChain :: T.Text -> Int -> Expr x -> [P.Doc]+    flattenChain name p e = case e of+        Binary op' l' r'+            | binaryName op' == name+            , binaryPrecedence op' == p+            , binaryCommutative op' ->+                flattenChain name p l' ++ flattenChain name p r'+        _ -> [toDoc p e]++    renderIf :: Int -> Expr x -> P.Doc+    renderIf prec (If c t e) =+        let blk =+                P.text "if" P.<+> P.nest 3 (P.group (toDoc 0 c))+                    <> P.hardline+                    <> P.text "then" P.<+> toDoc 0 t+                    <> P.hardline+                    <> renderElse e+         in if prec > 0 then P.parens (P.nest 2 blk) else blk+    renderIf _ _ = mempty++    renderElse :: Expr x -> P.Doc+    renderElse (If c t e) =+        P.text "else if" P.<+> P.nest 8 (P.group (toDoc 0 c))+            <> P.hardline+            <> P.text "then" P.<+> toDoc 0 t+            <> P.hardline+            <> renderElse e+    renderElse other = P.text "else" P.<+> toDoc 0 other

@@ -0,0 +1,410 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE Strict #-}+{-# LANGUAGE TypeApplications #-}++module DataFrame.Internal.Grouping (+    groupBy,+    groupBySeq,+    groupByPar,+    buildRowToGroup,+    changingPoints,+) where++import qualified Data.List as L+import qualified Data.Map as M+import qualified Data.Text as T+import qualified Data.Vector as V+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM++import Control.Exception (throw)+import Control.Monad+import Control.Monad.ST (ST, runST)+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))+import DataFrame.Errors+import DataFrame.Internal.Column (+    Bitmap,+    Column (..),+    bitmapTestBit,+ )+import DataFrame.Internal.DataFrame (DataFrame (..), GroupedDataFrame (..))+import DataFrame.Internal.DictEncode (dictEncodeColumnUpTo)+import DataFrame.Internal.GroupingDirect (+    DirectGrouping (..),+    tryDirectGroupColumn,+ )+import DataFrame.Internal.GroupingPar (parallelAssignGroups, shouldParallelize)+import DataFrame.Internal.Hash+import DataFrame.Internal.HashTable (htInsert, newHashTable)+import DataFrame.Internal.PackedText (+    PackedTextData,+    packedLength,+    packedSlice,+    sliceEqBytes,+ )+import DataFrame.Internal.RadixRank (rankByHash)+import DataFrame.Internal.Types+import System.IO.Unsafe (unsafePerformIO)+import Type.Reflection (typeRep)++{- | O(k * n) group the dataframe by the given key columns, bucketing rows with an+open-addressing hash table that re-verifies keys on each hash hit. Groups are+numbered in first-appearance order; 'valueIndices'/'offsets' follow by counting sort.+-}+groupBy ::+    [T.Text] ->+    DataFrame ->+    GroupedDataFrame+groupBy names df+    | any (`notElem` columnNames df) names =+        throw $+            ColumnsNotFoundException+                (names L.\\ columnNames df)+                "groupBy"+                (columnNames df)+    | nRows df == 0 =+        Grouped+            df+            names+            VU.empty+            (VU.fromList [0])+            VU.empty+    | Just dg <- tryDirectGroup names df = dg+    | shouldParallelize n = groupByPar names df+    | otherwise = groupBySeq names df+  where+    !n = nRows df++{- | Low-cardinality direct-indexed grouping fast path+('DataFrame.Internal.GroupingDirect'): fires only for a single clean small-range+@Int@ key. Returns 'Nothing' on any other key shape, falling back to the hash path.+-}+tryDirectGroup :: [T.Text] -> DataFrame -> Maybe GroupedDataFrame+tryDirectGroup [name] df = do+    col <- M.lookup name (columnIndices df) >>= \i -> columns df V.!? i+    case tryDirectGroupColumn col of+        Just dg ->+            Just (Grouped df [name] (dgValueIndices dg) (dgOffsets dg) (dgRowToGroup dg))+        Nothing -> tryDictGroup (nRows df) df [name] col+tryDirectGroup _ _ = Nothing++{- | Dictionary-encode a single text key to dense int codes, then derive+@valueIndices@/@offsets@ by counting sort. Profiled slower than the fused hash+group-by on every db-benchmark question, so it always falls back ('dictGroupEnabled').+-}+tryDictGroup ::+    Int -> DataFrame -> [T.Text] -> Column -> Maybe GroupedDataFrame+tryDictGroup n df names col+    | dictGroupEnabled && not (shouldParallelize n) = do+        (codes, card) <- dictEncodeColumnUpTo dictSingleThreshold col+        let (vis, os) = indicesFromGroups codes card+        Just (Grouped df names vis os codes)+    | otherwise = Nothing++{- | Master switch for the single-key dict-encode grouping path. 'False' because+it profiled slower than the hash group-by on every db-benchmark group-by question+(see 'tryDictGroup'); the path is kept compiled and tested but not taken.+-}+dictGroupEnabled :: Bool+dictGroupEnabled = False++{- | Cardinality ceiling for the single-key dict-encode probe: it bails to 'Nothing'+once the distinct count passes this. Only consulted when 'dictGroupEnabled' is 'True'.+-}+dictSingleThreshold :: Int+dictSingleThreshold = 4096++{- | The sequential grouping path: a single open-addressing table over all rows,+canonically remapped. Always available regardless of capabilities; the parallel+path is verified equal to it by a property test.+-}+groupBySeq :: [T.Text] -> DataFrame -> GroupedDataFrame+groupBySeq names df =+    let !n = nRows df+        indicesToGroup = keyColIndices names df+        (rtg0, repHash, repRow) = assignGroups df indicesToGroup n+        !nGroups = VU.length repHash+        !remap = canonicalRemap repHash repRow+        !rtg = VU.map (VU.unsafeIndex remap) rtg0+        (vis, os) = indicesFromGroups rtg nGroups+     in Grouped df names vis os rtg++{- | The parallel partitioned grouping path (see 'DataFrame.Internal.GroupingPar'):+forks one task per capability, producing output bit-for-bit identical to+'groupBySeq'. Pure via 'unsafePerformIO' (deterministic thread fan-out only).+-}+groupByPar :: [T.Text] -> DataFrame -> GroupedDataFrame+groupByPar names df =+    let !n = nRows df+        indicesToGroup = keyColIndices names df+        !hashes = runST (computeHashes df indicesToGroup n)+        !eqRow = eqKeyRow df indicesToGroup+        (rtg, vis, os) = unsafePerformIO (parallelAssignGroups n hashes eqRow)+     in Grouped df names vis os rtg+{-# NOINLINE groupByPar #-}++-- | Column indices of the requested key columns, in column order.+keyColIndices :: [T.Text] -> DataFrame -> [Int]+keyColIndices names df =+    M.elems $ M.filterWithKey (\k _ -> k `elem` names) (columnIndices df)++{- | Assign every row to a dense group id in first-appearance order. Returns+@(rowToGroup, repHash, repRow)@ — the hash and representative row of each group.+'eqKeyRow' re-verifies the real key on each hash hit so colliding keys stay apart.+-}+assignGroups ::+    DataFrame -> [Int] -> Int -> (VU.Vector Int, VU.Vector Int, VU.Vector Int)+assignGroups df indicesToGroup n = runST $ do+    hashes <- computeHashes df indicesToGroup n+    let !eqRow = eqKeyRow df indicesToGroup+    ht <- newHashTable n+    rtg <- VUM.new n+    repHashM <- VUM.new n+    repRowM <- VUM.new n+    let go !i !next+            | i >= n = pure next+            | otherwise = do+                let !h = VU.unsafeIndex hashes i+                (gid, isNew) <- htInsert ht eqRow next i h+                VUM.unsafeWrite rtg i gid+                when isNew $ do+                    VUM.unsafeWrite repHashM next h+                    VUM.unsafeWrite repRowM next i+                go (i + 1) (if isNew then next + 1 else next)+    !nGroups <- go 0 0+    frozen <- VU.unsafeFreeze rtg+    repHash <- VU.unsafeFreeze (VUM.slice 0 nGroups repHashM)+    repRow <- VU.unsafeFreeze (VUM.slice 0 nGroups repRowM)+    pure (frozen, repHash, repRow)++{- | Map each first-appearance group id to its canonical id: groups ordered by+ascending representative hash (tie-broken by representative row), making group+order a deterministic function of the key set so set ops commute. O(g), no sort.+-}+canonicalRemap :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int+canonicalRemap repHash _repRow =+    runST (rankByHash (pure . VU.unsafeIndex repHash) (VU.length repHash))++{- | Compute the FNV row-hash of the key columns into a fresh unboxed vector,+mixing 'nullSalt' for null slots so a missing value never collides with a+present one of the same bits.+-}+computeHashes :: DataFrame -> [Int] -> Int -> ST s (VU.Vector Int)+computeHashes df indicesToGroup n = do+    mh <- VUM.replicate n fnvOffset+    let selectedCols = map (columns df V.!) indicesToGroup+    forM_ selectedCols $ \case+        UnboxedColumn ubm (v :: VU.Vector a) ->+            case testEquality (typeRep @a) (typeRep @Int) of+                Just Refl -> hashUnboxed mh ubm mixInt v+                Nothing ->+                    case testEquality (typeRep @a) (typeRep @Double) of+                        Just Refl -> hashUnboxed mh ubm mixDouble v+                        Nothing ->+                            case sIntegral @a of+                                STrue ->+                                    hashUnboxed mh ubm (\h d -> mixInt h (fromIntegral @a @Int d)) v+                                SFalse ->+                                    case sFloating @a of+                                        STrue ->+                                            hashUnboxed mh ubm (\h d -> mixDouble h (realToFrac d :: Double)) v+                                        SFalse ->+                                            hashUnboxed mh ubm mixShow v+        BoxedColumn bm (v :: V.Vector a) ->+            case testEquality (typeRep @a) (typeRep @T.Text) of+                Just Refl ->+                    V.imapM_+                        ( \i t -> do+                            !h <- VUM.unsafeRead mh i+                            let h' = case bm of+                                    Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt+                                    _ -> mixText h t+                            VUM.unsafeWrite mh i h'+                        )+                        v+                Nothing ->+                    V.imapM_+                        ( \i d -> do+                            !h <- VUM.unsafeRead mh i+                            let h' = case bm of+                                    Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt+                                    _ -> mixShow h d+                            VUM.unsafeWrite mh i h'+                        )+                        v+        PackedText bm p -> hashPacked mh bm p+    VU.unsafeFreeze mh++{- | Build the row-key equality predicate over the selected key columns.+@eqKeyRow df idxs a b@ is 'True' iff rows @a@ and @b@ agree on all key columns+(validity first, a null equals only a null). Used to reject hash collisions.+-}+eqKeyRow :: DataFrame -> [Int] -> Int -> Int -> Bool+eqKeyRow df indicesToGroup =+    let !preds = map (colEqRow . (columns df V.!)) indicesToGroup+        go [] _ _ = True+        go (p : ps) a b = p a b && go ps a b+     in go preds++{- | Per-column row equality respecting nulls. Two rows are equal at a column+when both are null, or both are valid and their values compare equal.+-}+colEqRow :: Column -> (Int -> Int -> Bool)+colEqRow (UnboxedColumn bm v) =+    let eqV a b = VU.unsafeIndex v a == VU.unsafeIndex v b+     in withNulls bm eqV+colEqRow (BoxedColumn bm v) =+    let eqV a b = V.unsafeIndex v a == V.unsafeIndex v b+     in withNulls bm eqV+colEqRow (PackedText bm p) =+    let eqV a b =+            let (arrA, oA, lA) = packedSlice p a+                (arrB, oB, lB) = packedSlice p b+             in sliceEqBytes arrA oA lA arrB oB lB+     in withNulls bm eqV+{-# INLINE colEqRow #-}++{- | Wrap a value-equality with null handling: equal iff both valid and the+values agree, or both null.+-}+withNulls :: Maybe Bitmap -> (Int -> Int -> Bool) -> (Int -> Int -> Bool)+withNulls Nothing eqV = eqV+withNulls (Just bm) eqV = \a b ->+    case (bitmapTestBit bm a, bitmapTestBit bm b) of+        (True, True) -> eqV a b+        (False, False) -> True+        _ -> False+{-# INLINE withNulls #-}++{- | Derive @(valueIndices, offsets)@ from @rowToGroup@ via a stable counting+sort on the group id: a per-group count, a prefix-sum into group offsets, then a+single placement pass keeps rows in original order within each group.+-}+indicesFromGroups :: VU.Vector Int -> Int -> (VU.Vector Int, VU.Vector Int)+indicesFromGroups rtg nGroups = runST $ do+    let !n = VU.length rtg+    counts <- VUM.replicate (nGroups + 1) 0+    let countLoop !i+            | i >= n = pure ()+            | otherwise = do+                let !g = VU.unsafeIndex rtg i+                c <- VUM.unsafeRead counts g+                VUM.unsafeWrite counts g (c + 1)+                countLoop (i + 1)+    countLoop 0+    offsM <- VUM.new (nGroups + 1)+    let scan !k !acc+            | k > nGroups = pure ()+            | otherwise = do+                VUM.unsafeWrite offsM k acc+                c <- VUM.unsafeRead counts k+                scan (k + 1) (acc + c)+    scan 0 0+    let seed !k+            | k > nGroups = pure ()+            | otherwise = do+                s <- VUM.unsafeRead offsM k+                VUM.unsafeWrite counts k s+                seed (k + 1)+    seed 0+    vis <- VUM.new n+    let place !i+            | i >= n = pure ()+            | otherwise = do+                let !g = VU.unsafeIndex rtg i+                pos <- VUM.unsafeRead counts g+                VUM.unsafeWrite vis pos i+                VUM.unsafeWrite counts g (pos + 1)+                place (i + 1)+    place 0+    offs <- VU.unsafeFreeze offsM+    frozenVis <- VU.unsafeFreeze vis+    pure (frozenVis, offs)++{- | Fold a value-mix over an unboxed column into the running hash vector,+respecting the null bitmap: a null slot mixes a fixed 'nullSalt' sentinel.+-}+hashUnboxed ::+    (VU.Unbox a) =>+    VUM.MVector s Int ->+    Maybe Bitmap ->+    (Int -> a -> Int) ->+    VU.Vector a ->+    ST s ()+hashUnboxed mh ubm mix v = case ubm of+    Nothing ->+        VU.imapM_+            ( \i x -> do+                !h <- VUM.unsafeRead mh i+                VUM.unsafeWrite mh i (mix h x)+            )+            v+    Just bm ->+        VU.imapM_+            ( \i x -> do+                !h <- VUM.unsafeRead mh i+                VUM.unsafeWrite+                    mh+                    i+                    (if bitmapTestBit bm i then mix h x else mixInt h nullSalt)+            )+            v+{-# INLINE hashUnboxed #-}++{- | Hash a packed-text column over its raw UTF-8 byte slices (no per-row+'Data.Text.Text'), mixing 'nullSalt' for null rows. Shares 'mixBytes' with+'mixText' so packed and boxed Text columns hash identically.+-}+hashPacked ::+    VUM.MVector s Int -> Maybe Bitmap -> PackedTextData -> ST s ()+hashPacked mh bm p = go 0+  where+    !n = packedLength p+    go !i+        | i >= n = pure ()+        | otherwise = do+            !h <- VUM.unsafeRead mh i+            let h' = case bm of+                    Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt+                    _ -> let (arr, o, l) = packedSlice p i in mixBytes h arr o l+            VUM.unsafeWrite mh i h'+            go (i + 1)+{-# INLINE hashPacked #-}++-- Inline accessors to avoid depending on Operations.Core++columnNames :: DataFrame -> [T.Text]+columnNames = M.keys . columnIndices++nRows :: DataFrame -> Int+nRows = fst . dataframeDimensions++{- | 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+        (VU.fromList (VU.length vs : fst (VU.ifoldl' findChangePoints initialState vs)))+  where+    initialState = ([0], snd (VU.head vs))+    findChangePoints (!offs, !currentVal) index (_, !newVal)+        | currentVal == newVal = (offs, currentVal)+        | otherwise = (index : offs, newVal)

@@ -0,0 +1,234 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | Low-cardinality direct-indexed grouping fast path: when the key is a single+clean unboxed @Int@ column of small value range, the value itself indexes a dense+accumulator (no hashing/probing). Emits groups in ascending value order.+-}+module DataFrame.Internal.GroupingDirect (+    directGroupThreshold,+    tryDirectGroupColumn,+    DirectGrouping (..),+) where++import Control.Concurrent (forkIO, getNumCapabilities)+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)+import Control.Exception (SomeException, throwIO, try)+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import System.IO.Unsafe (unsafePerformIO)+import Type.Reflection (typeRep)++import DataFrame.Internal.Column (Column (..))++{- | Largest key value RANGE (max - min + 1) the direct grouping path accepts. A+@2^20@-slot histogram is 8MB; the low-cardinality questions sit far below it+(id4 range 100, id6 range 1e5). Wider ranges fall back to the hash group-by.+-}+directGroupThreshold :: Int+directGroupThreshold = 1048576++{- | The grouping layout the hash path also produces: @rowToGroup@, the+group-sorted @valueIndices@, the @offsets@ prefix array, and the group count.+-}+data DirectGrouping = DirectGrouping+    { dgRowToGroup :: !(VU.Vector Int)+    , dgValueIndices :: !(VU.Vector Int)+    , dgOffsets :: !(VU.Vector Int)+    , dgNGroups :: !Int+    }++capabilities :: Int+capabilities = unsafePerformIO getNumCapabilities+{-# NOINLINE capabilities #-}++parThreshold :: Int+parThreshold = 200000++{- | Take the direct path if the (single) key column is a clean non-null unboxed+@Int@ column with a small value range. Returns 'Nothing' to fall back to the+hash group-by on anything else (boxed/text keys, nullable, wide ranges, empty).+-}+tryDirectGroupColumn :: Column -> Maybe DirectGrouping+tryDirectGroupColumn (UnboxedColumn Nothing (v :: VU.Vector a))+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int)+    , not (VU.null v) =+        let (!mn, !mx) = rangeOf v+            !range = mx - mn + 1+         in if range >= 1 && range <= directGroupThreshold+                then Just (directGroup v mn range)+                else Nothing+tryDirectGroupColumn _ = Nothing++-- | Parallel min/max reduce (order-independent).+rangeOf :: VU.Vector Int -> (Int, Int)+rangeOf v+    | not (shouldPar n) = rangeChunk v 0 n+    | otherwise = unsafePerformIO $ do+        let !caps = capabilities+            !per = (n + caps - 1) `div` caps+            spawn w = do+                var <- newEmptyMVar+                let !lo = min n (w * per)+                    !hi = min n (lo + per)+                _ <- forkIO (try (pure $! rangeChunk v lo hi) >>= putMVar var)+                pure var+        vars <- mapM spawn [0 .. caps - 1]+        rs <- mapM takeMVar vars+        rs' <- mapM (either (throwIO @SomeException) pure) rs+        pure (combineRanges (filter (\(a, _) -> a /= maxBound) rs'))+  where+    !n = VU.length v+{-# NOINLINE rangeOf #-}++rangeChunk :: VU.Vector Int -> Int -> Int -> (Int, Int)+rangeChunk v lo hi = go lo maxBound minBound+  where+    go !i !mn !mx+        | i >= hi = (mn, mx)+        | otherwise =+            let !x = VU.unsafeIndex v i+             in go (i + 1) (min mn x) (max mx x)++combineRanges :: [(Int, Int)] -> (Int, Int)+combineRanges [] = (0, 0)+combineRanges ((a0, b0) : rest) = foldr (\(a, b) (ma, mb) -> (min ma a, max mb b)) (a0, b0) rest++shouldPar :: Int -> Bool+shouldPar n = n >= parThreshold && capabilities > 1++{- | Build the grouping by counting sort on @value - min@: a (parallel) per-value+histogram, compaction of non-empty values into ascending dense ids, a scan into+offsets, then a stable placement pass building @valueIndices@ and @rowToGroup@.+-}+directGroup :: VU.Vector Int -> Int -> Int -> DirectGrouping+directGroup v mn range = unsafePerformIO $ do+    let !n = VU.length v+    hist <- buildHistogram v mn range n+    valToGroup <- VUM.replicate range (-1 :: Int)+    grpCount <- VUM.new range+    nGroups <- compact hist range valToGroup grpCount+    offsM <- VUM.new (nGroups + 1)+    cursor <- VUM.new nGroups+    scanOffsets grpCount nGroups offsM cursor+    rtg <- VUM.new n+    vis <- VUM.new n+    place v mn n valToGroup cursor rtg vis+    frozenRtg <- VU.unsafeFreeze rtg+    frozenVis <- VU.unsafeFreeze vis+    frozenOffs <- VU.unsafeFreeze offsM+    pure (DirectGrouping frozenRtg frozenVis frozenOffs nGroups)+{-# NOINLINE directGroup #-}++{- | Parallel per-value histogram: each worker fills a private @range@-slot+count over its row chunk, then the partials are summed (exact integers, so the+merge order is irrelevant). Sequential single pass below 'parThreshold'.+-}+buildHistogram :: VU.Vector Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)+buildHistogram v mn range n+    | not (shouldPar n) = histChunk v mn range 0 n+    | otherwise = do+        let !caps = capabilities+            !per = (n + caps - 1) `div` caps+            spawn w = do+                var <- newEmptyMVar+                let !lo = min n (w * per)+                    !hi = min n (lo + per)+                _ <- forkIO (try (histChunk v mn range lo hi) >>= putMVar var)+                pure var+        vars <- mapM spawn [0 .. caps - 1]+        rs <- mapM takeMVar vars+        parts <- mapM (either (throwIO @SomeException) pure) rs+        case parts of+            [] -> VUM.replicate range 0+            (p0 : rest) -> do+                mapM_ (addInto p0 range) rest+                pure p0++histChunk :: VU.Vector Int -> Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)+histChunk v mn range lo hi = do+    acc <- VUM.replicate range (0 :: Int)+    let go !i+            | i >= hi = pure ()+            | otherwise = do+                let !k = VU.unsafeIndex v i - mn+                c <- VUM.unsafeRead acc k+                VUM.unsafeWrite acc k (c + 1)+                go (i + 1)+    go lo+    pure acc++addInto :: VUM.IOVector Int -> Int -> VUM.IOVector Int -> IO ()+addInto dst range src = go 0+  where+    go !k+        | k >= range = pure ()+        | otherwise = do+            a <- VUM.unsafeRead dst k+            b <- VUM.unsafeRead src k+            VUM.unsafeWrite dst k (a + b)+            go (k + 1)++{- | Walk the histogram in ascending value order, assigning a dense group id to+each non-empty value and copying its count into @grpCount@ at that id. Returns+the group count.+-}+compact ::+    VUM.IOVector Int -> Int -> VUM.IOVector Int -> VUM.IOVector Int -> IO Int+compact hist range valToGroup grpCount = go 0 0+  where+    go !val !next+        | val >= range = pure next+        | otherwise = do+            c <- VUM.unsafeRead hist val+            if c == 0+                then go (val + 1) next+                else do+                    VUM.unsafeWrite valToGroup val next+                    VUM.unsafeWrite grpCount next c+                    go (val + 1) (next + 1)++{- | Exclusive prefix scan of group counts into @offsM@ (length nGroups+1) and+seed the per-group write @cursor@ at each group's start offset.+-}+scanOffsets ::+    VUM.IOVector Int -> Int -> VUM.IOVector Int -> VUM.IOVector Int -> IO ()+scanOffsets grpCount nGroups offsM cursor = go 0 0+  where+    go !g !acc+        | g >= nGroups = VUM.unsafeWrite offsM nGroups acc+        | otherwise = do+            VUM.unsafeWrite offsM g acc+            VUM.unsafeWrite cursor g acc+            c <- VUM.unsafeRead grpCount g+            go (g + 1) (acc + c)++{- | Stable placement pass: for each row in original order, look up its group id+through the value map, write @rowToGroup@, and append the row to its group's run+in @valueIndices@ via the advancing cursor (rows keep original order per group).+-}+place ::+    VU.Vector Int ->+    Int ->+    Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    IO ()+place v mn n valToGroup cursor rtg vis = go 0+  where+    go !i+        | i >= n = pure ()+        | otherwise = do+            let !val = VU.unsafeIndex v i - mn+            g <- VUM.unsafeRead valToGroup val+            VUM.unsafeWrite rtg i g+            pos <- VUM.unsafeRead cursor g+            VUM.unsafeWrite vis pos i+            VUM.unsafeWrite cursor g (pos + 1)+            go (i + 1)

@@ -0,0 +1,308 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE Strict #-}++{- | Parallel partitioned group-by: rows are counting-sorted into partitions by the+top hash bits, then one task per capability groups its partitions independently.+Output is bit-for-bit identical to the sequential 'DataFrame.Internal.Grouping.groupBy'.+-}+module DataFrame.Internal.GroupingPar (+    parallelAssignGroups,+    shouldParallelize,+    parThreshold,+    numPartitionsFor,+) where++import Control.Concurrent (forkIO, getNumCapabilities)+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)+import Control.Exception (SomeException, throwIO, try)+import Control.Monad (forM_, when)+import Data.Bits (countLeadingZeros, unsafeShiftR)+import Data.IORef (atomicModifyIORef', newIORef)+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 Data.Word (Word64)+import DataFrame.Internal.HashTable (+    htInsert,+    newHashTable,+ )+import DataFrame.Internal.RadixRank (rankByHash)+import System.IO.Unsafe (unsafePerformIO)++{- | Below this many rows the partition/fork overhead is not worth it; 'groupBy'+uses its sequential 'ST' path instead.+-}+parThreshold :: Int+parThreshold = 200000++{- | Whether 'groupBy' should take the parallel path: more than one capability+and at least 'parThreshold' rows.+-}+shouldParallelize :: Int -> Bool+shouldParallelize n = n >= parThreshold && capabilities > 1+{-# NOINLINE shouldParallelize #-}++capabilities :: Int+capabilities = unsafePerformIO getNumCapabilities+{-# NOINLINE capabilities #-}++{- | Sign-preserving unsigned remap: ascending 'Word64' order of @key h@ equals+ascending signed-'Int' order of @h@, so partitioning and sorting on it reproduce+the sequential @compare \`on\` repHash@ ordering exactly.+-}+key :: Int -> Word64+key h = fromIntegral h + 0x8000000000000000+{-# INLINE key #-}++-- | Partition index of a hash: the top @log2 p@ bits of its unsigned key.+partIx :: Int -> Int -> Int+partIx shift h = fromIntegral (key h `unsafeShiftR` shift)+{-# INLINE partIx #-}++{- | Number of partitions: a power of two, at least @4 * caps@ (P >> cores for+skew tolerance), floored at 256.+-}+numPartitionsFor :: Int -> Int+numPartitionsFor caps = go 1+  where+    target = max 256 (4 * caps)+    go p+        | p >= target = p+        | otherwise = go (p * 2)++-- | @floor (log2 x)@ for a power-of-two @x@.+intLog2 :: Int -> Int+intLog2 x = 63 - countLeadingZeros x+{-# INLINE intLog2 #-}++{- | Parallel group assignment. @parallelAssignGroups n hashes eqRow@ returns+@(rowToGroup, valueIndices, offsets)@ in canonical group order. @eqRow a b@ must+report whether rows @a@ and @b@ share all key columns (null-aware).+-}+parallelAssignGroups ::+    Int ->+    VU.Vector Int ->+    (Int -> Int -> Bool) ->+    IO (VU.Vector Int, VU.Vector Int, VU.Vector Int)+parallelAssignGroups n hashes eqRow = do+    caps <- getNumCapabilities+    let !p = numPartitionsFor caps+        !shift = 64 - intLog2 p+    (partStart, sortedRows) <- partitionRows n hashes p shift+    localGid <- VUM.new (max 1 n)+    canonBoxes <- VM.replicate p (VU.empty :: VU.Vector Int)+    nLocalGroups <- VUM.replicate p (0 :: Int)+    runPartitions+        caps+        p+        partStart+        sortedRows+        hashes+        eqRow+        localGid+        canonBoxes+        nLocalGroups+    (globalBase, canonOf, nGroups) <- canonicalize p canonBoxes nLocalGroups+    assemble n p partStart sortedRows localGid globalBase canonOf nGroups++-------------------------------------------------------------------------------+-- Phase 1: counting sort by partition+-------------------------------------------------------------------------------++{- | Bucket every row index into its partition by a counting sort. Returns the+exclusive prefix-sum @partStart@ (length @p+1@, @partStart[p] == n@) and the row+indices laid out partition-by-partition in @sortedRows@.+-}+partitionRows ::+    Int -> VU.Vector Int -> Int -> Int -> IO (VU.Vector Int, VU.Vector Int)+partitionRows n hashes p shift = do+    counts <- VUM.replicate (p + 1) (0 :: Int)+    let countLoop !i+            | i >= n = pure ()+            | otherwise = do+                let !pp = partIx shift (VU.unsafeIndex hashes i)+                c <- VUM.unsafeRead counts pp+                VUM.unsafeWrite counts pp (c + 1)+                countLoop (i + 1)+    countLoop 0+    partStartM <- VUM.new (p + 1)+    let scan !k !acc+            | k > p = pure ()+            | otherwise = do+                VUM.unsafeWrite partStartM k acc+                c <- if k < p then VUM.unsafeRead counts k else pure 0+                scan (k + 1) (acc + c)+    scan 0 0+    cursor <- VUM.new p+    forM_ [0 .. p - 1] $ \k -> VUM.unsafeRead partStartM k >>= VUM.unsafeWrite cursor k+    sortedM <- VUM.new (max 1 n)+    let place !i+            | i >= n = pure ()+            | otherwise = do+                let !pp = partIx shift (VU.unsafeIndex hashes i)+                pos <- VUM.unsafeRead cursor pp+                VUM.unsafeWrite sortedM pos i+                VUM.unsafeWrite cursor pp (pos + 1)+                place (i + 1)+    place 0+    partStart <- VU.unsafeFreeze partStartM+    sortedRows <- VU.unsafeFreeze sortedM+    pure (partStart, sortedRows)++-------------------------------------------------------------------------------+-- Phase 2: per-partition grouping (parallel)+-------------------------------------------------------------------------------++{- | Group each partition with its own hash table, then rank its local groups into+canonical order — all inside the parallel worker. Forks @caps@ workers pulling+partition indices off a shared counter; disjoint keys mean no cross-partition merge.+-}+runPartitions ::+    Int ->+    Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    (Int -> Int -> Bool) ->+    VUM.IOVector Int ->+    VM.IOVector (VU.Vector Int) ->+    VUM.IOVector Int ->+    IO ()+runPartitions caps p partStart sortedRows hashes eqRow localGid canonBoxes nLocalGroups = do+    next <- newIORef 0+    let groupPartition !pp = do+            let !s = VU.unsafeIndex partStart pp+                !e = VU.unsafeIndex partStart (pp + 1)+                !sz = e - s+            when (sz > 0) $ do+                ht <- newHashTable sz+                repHashM <- VUM.new sz+                let loop !pos !nextGid+                        | pos >= e = pure nextGid+                        | otherwise = do+                            let !row = VU.unsafeIndex sortedRows pos+                                !h = VU.unsafeIndex hashes row+                            (gid, isNew) <- htInsert ht eqRow nextGid row h+                            VUM.unsafeWrite localGid pos gid+                            if isNew+                                then do+                                    VUM.unsafeWrite repHashM nextGid h+                                    loop (pos + 1) (nextGid + 1)+                                else loop (pos + 1) nextGid+                ng <- loop s 0+                VUM.unsafeWrite nLocalGroups pp ng+                canon <- rankByHash (VUM.unsafeRead repHashM) ng+                VM.unsafeWrite canonBoxes pp canon+        worker = do+            i <- atomicModifyIORef' next (\j -> (j + 1, j))+            when (i < p) $ groupPartition i >> worker+    forkJoin_ (replicate caps worker)++-------------------------------------------------------------------------------+-- Phase 3: global base ids + assembly+-------------------------------------------------------------------------------++{- | Exclusive prefix sum of the per-partition group counts into @globalBase@+(@globalBase[pp]@ = first global id of partition @pp@). Ranks were computed in+'runPartitions'; prepending the base to each yields the sequential order.+-}+canonicalize ::+    Int ->+    VM.IOVector (VU.Vector Int) ->+    VUM.IOVector Int ->+    IO (VU.Vector Int, V.Vector (VU.Vector Int), Int)+canonicalize p canonBoxes nLocalGroups = do+    globalBaseM <- VUM.new (p + 1)+    let go !pp !base+            | pp >= p = VUM.unsafeWrite globalBaseM p base >> pure base+            | otherwise = do+                VUM.unsafeWrite globalBaseM pp base+                ng <- VUM.unsafeRead nLocalGroups pp+                go (pp + 1) (base + ng)+    total <- go 0 0+    globalBase <- VU.unsafeFreeze globalBaseM+    canonOf <- V.unsafeFreeze canonBoxes+    pure (globalBase, canonOf, total)++{- | Build the final @(rowToGroup, valueIndices, offsets)@: the global group id of a+sorted position is @globalBase[pp] + canonOf[pp][localGid]@. @valueIndices@ orders+rows by group, @offsets@ the boundaries, @rowToGroup@ the inverse per original row.+-}+assemble ::+    Int ->+    Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    VUM.IOVector Int ->+    VU.Vector Int ->+    V.Vector (VU.Vector Int) ->+    Int ->+    IO (VU.Vector Int, VU.Vector Int, VU.Vector Int)+assemble n p partStart sortedRows localGid globalBase canonOf nGroups = do+    rtgM <- VUM.new (max 1 n)+    counts <- VUM.replicate (nGroups + 1) (0 :: Int)+    gidAt <- VUM.new (max 1 n)+    let scanPos !pp+            | pp >= p = pure ()+            | otherwise = do+                let !s = VU.unsafeIndex partStart pp+                    !e = VU.unsafeIndex partStart (pp + 1)+                    !base = VU.unsafeIndex globalBase pp+                    !canon = V.unsafeIndex canonOf pp+                let inner !pos+                        | pos >= e = pure ()+                        | otherwise = do+                            lg <- VUM.unsafeRead localGid pos+                            let !g = base + VU.unsafeIndex canon lg+                                !row = VU.unsafeIndex sortedRows pos+                            VUM.unsafeWrite gidAt pos g+                            VUM.unsafeWrite rtgM row g+                            c <- VUM.unsafeRead counts g+                            VUM.unsafeWrite counts g (c + 1)+                            inner (pos + 1)+                inner s+                scanPos (pp + 1)+    scanPos 0+    offsM <- VUM.new (nGroups + 1)+    let scan !k !acc+            | k > nGroups = pure ()+            | otherwise = do+                VUM.unsafeWrite offsM k acc+                c <- if k < nGroups then VUM.unsafeRead counts k else pure 0+                scan (k + 1) (acc + c)+    scan 0 0+    cursor <- VUM.new (max 1 nGroups)+    forM_ [0 .. nGroups - 1] $ \k -> VUM.unsafeRead offsM k >>= VUM.unsafeWrite cursor k+    visM <- VUM.new (max 1 n)+    let placeVis !pos+            | pos >= n = pure ()+            | otherwise = do+                g <- VUM.unsafeRead gidAt pos+                let !row = VU.unsafeIndex sortedRows pos+                c <- VUM.unsafeRead cursor g+                VUM.unsafeWrite visM c row+                VUM.unsafeWrite cursor g (c + 1)+                placeVis (pos + 1)+    placeVis 0+    rtg <- VU.unsafeFreeze rtgM+    offs <- VU.unsafeFreeze offsM+    vis <- VU.unsafeFreeze visM+    pure (rtg, vis, offs)++-------------------------------------------------------------------------------+-- Thread fan-out (plain forkIO + MVar join, no sparks)+-------------------------------------------------------------------------------++-- | Run each action on its own thread; rethrow the first failure (in order).+forkJoin_ :: [IO ()] -> IO ()+forkJoin_ actions = do+    vars <- mapM spawn actions+    results <- mapM takeMVar vars+    mapM_ (either (throwIO :: SomeException -> IO ()) pure) results+  where+    spawn act = do+        var <- newEmptyMVar+        _ <- forkIO (try act >>= putMVar var)+        pure var

@@ -0,0 +1,113 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE MagicHash #-}++{- | A poor-man's hash used by 'DataFrame.Internal.Grouping' to bucket rows+without depending on @hashable@. Each value is folded into an 'Int' with an+FxHash-style step (rotate, xor, multiply); small and not cryptographic.+-}+module DataFrame.Internal.Hash (+    fnvOffset,+    nullSalt,+    mixInt,+    mixDouble,+    mixBool,+    mixChar,+    mixText,+    mixBytes,+    mixShow,+) where++import Data.Bits (rotateL, unsafeShiftL, unsafeShiftR, xor)+import Data.Char (ord)+import qualified Data.Text as T+import qualified Data.Text.Array as A+#if MIN_VERSION_text(2,1,0)+import Data.Array.Byte (ByteArray (ByteArray))+#else+import Data.Text.Array (Array (ByteArray))+#endif+import Data.Text.Internal (Text (Text))+import GHC.Exts (Int (I#), indexWord8Array#, indexWord8ArrayAsWord64#)+import GHC.Word (Word64 (W64#), Word8 (W8#))++{- | FNV-1a 64-bit offset basis (used as the initial accumulator).+The literal is unsigned and exceeds 'Int' range, so we round-trip through+'Word64' to get the well-defined two's-complement bit pattern.+-}+fnvOffset :: Int+fnvOffset = fromIntegral (0xcbf29ce484222325 :: Word64)++-- | FNV-1a 64-bit prime.+fnvPrime :: Int+fnvPrime = 0x00000100000001b3++{- | Sentinel mixed in for a /null/ slot, so @Nothing@ does not hash the same as+a present value with equal bits (e.g. @Just 0@). A fixed distinctive constant+keeps null hashing deterministic; a real value equal to it collides only rarely.+-}+nullSalt :: Int+nullSalt = fromIntegral (0x9E3779B97F4A7C15 :: Word64)++{- | Mix an 'Int' into the accumulator with an FxHash-style step. The rotate+diffuses each value's bits before the next is folded in, avoiding the structured+collisions a plain xor-then-multiply produces on small/adjacent group keys.+-}+mixInt :: Int -> Int -> Int+mixInt acc x = (rotateL acc 13 `xor` x) * fnvPrime+{-# INLINE mixInt #-}++{- | Mix a 'Double' into the accumulator. Loses sub-millisecond precision+but matches the bucketing the old hashable-based code used.+-}+mixDouble :: Int -> Double -> Int+mixDouble acc d = mixInt acc (floor (d * 1000))+{-# INLINE mixDouble #-}++mixBool :: Int -> Bool -> Int+mixBool acc b = mixInt acc (if b then 1 else 0)+{-# INLINE mixBool #-}++mixChar :: Int -> Char -> Int+mixChar acc = mixInt acc . ord+{-# INLINE mixChar #-}++{- | Mix a 'T.Text' value into the accumulator over its raw UTF-8 bytes, eight at+a time. Reading a whole 'Word64' per step cuts the multiply count ~8x on long+keys while staying collision-equivalent (UTF-8 is injective).+-}+mixText :: Int -> T.Text -> Int+mixText !acc (Text arr off len) = mixBytes acc arr off len+{-# INLINE mixText #-}++{- | Mix a raw UTF-8 byte slice @[off, off+len)@ of a 'Data.Text.Array.Array'+into the accumulator, eight bytes at a time. The shared kernel behind+'mixText' and the packed-text hash path, so the two never drift.+-}+mixBytes :: Int -> A.Array -> Int -> Int -> Int+mixBytes !acc arr off len = goBytes (goWords acc off) wordsEnd+  where+    !(ByteArray ba) = arr+    !nWords = len `unsafeShiftR` 3+    !wordsEnd = off + (nWords `unsafeShiftL` 3)+    !end = off + len+    goWords !h !i+        | i >= wordsEnd = h+        | otherwise =+            let !(I# i#) = i+                !w = fromIntegral (W64# (indexWord8ArrayAsWord64# ba i#)) :: Int+             in goWords (mixInt h w) (i + 8)+    goBytes !h !i+        | i >= end = h+        | otherwise =+            let !(I# i#) = i+                !b = fromIntegral (W8# (indexWord8Array# ba i#)) :: Int+             in goBytes (mixInt h b) (i + 1)+{-# INLINE mixBytes #-}++{- | Fallback for arbitrary 'Show'-able values. Slower but covers types+without a dedicated combinator (e.g. 'Day', 'UTCTime').+-}+mixShow :: (Show a) => Int -> a -> Int+mixShow acc = mixText acc . T.pack . show+{-# INLINE mixShow #-}

@@ -0,0 +1,94 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}++{- | A flat, unboxed, open-addressing (linear-probe) hash table mapping a row's+key-hash to a dense group id, re-verifying the real key on every hash hit to+reject collisions. Runs in any 'PrimMonad' ('ST' for grouping, 'IO' per worker).+-}+module DataFrame.Internal.HashTable (+    HashTable (..),+    newHashTable,+    htInsert,+    nextPow2Above,+) where++import Control.Monad.Primitive (PrimMonad, PrimState)+import Data.Bits ((.&.))+import qualified Data.Vector.Unboxed.Mutable as VUM++{- | An open-addressing linear-probe table. @htMask@ is @capacity - 1@ (capacity+is a power of two) and maps a hash to its home slot.+-}+data HashTable s = HashTable+    { htHash :: !(VUM.MVector s Int)+    , htGroup :: !(VUM.MVector s Int)+    , htRep :: !(VUM.MVector s Int)+    , htMask :: !Int+    }++{- | Smallest power of two strictly greater than @n@, at least 2. Sizes the+table so the load factor stays below ~0.5 even when every row is a distinct+group.+-}+nextPow2Above :: Int -> Int+nextPow2Above n = go 2+  where+    go !p+        | p > n = p+        | otherwise = go (p * 2)+{-# INLINE nextPow2Above #-}++{- | Allocate an empty table able to hold up to @n@ distinct groups while+keeping the load factor under ~0.5 (capacity @= nextPow2Above (2*n)@). All+group slots start empty (@-1@).+-}+newHashTable :: (PrimMonad m) => Int -> m (HashTable (PrimState m))+newHashTable n = do+    let !cap = nextPow2Above (2 * max 1 n)+    h <- VUM.unsafeNew cap+    g <- VUM.replicate cap (-1)+    r <- VUM.unsafeNew cap+    pure (HashTable h g r (cap - 1))+{-# INLINE newHashTable #-}++{- | Look up @row@ (with precomputed @hash@) and return its dense group id: an+empty slot starts a new group via @nextGroup@, a stored-hash match is re-verified+with @eqRow@ before reuse. The 'Bool' is 'True' when a new group was created.+-}+htInsert ::+    (PrimMonad m) =>+    HashTable (PrimState m) ->+    -- | @eqRow a b@: do rows @a@ and @b@ have equal key columns?+    (Int -> Int -> Bool) ->+    -- | Next dense group id to assign if this row starts a new group.+    Int ->+    -- | Row index being inserted.+    Int ->+    -- | Precomputed hash of the row's key.+    Int ->+    m (Int, Bool)+htInsert ht eqRow nextGroup row hash = go (hash .&. mask)+  where+    !mask = htMask ht+    !hs = htHash ht+    !gs = htGroup ht+    !rs = htRep ht+    go !slot = do+        g <- VUM.unsafeRead gs slot+        if g < 0+            then do+                VUM.unsafeWrite hs slot hash+                VUM.unsafeWrite gs slot nextGroup+                VUM.unsafeWrite rs slot row+                pure (nextGroup, True)+            else do+                h <- VUM.unsafeRead hs slot+                if h == hash+                    then do+                        rep <- VUM.unsafeRead rs slot+                        if eqRow rep row+                            then pure (g, False)+                            else go ((slot + 1) .&. mask)+                    else go ((slot + 1) .&. mask)+{-# INLINE htInsert #-}

@@ -0,0 +1,1051 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -Wno-orphans #-}++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 qualified Data.Vector.Unboxed.Mutable as VUM+import DataFrame.Errors+import DataFrame.Internal.Column+import DataFrame.Internal.DataFrame+import DataFrame.Internal.Expression+import qualified DataFrame.Internal.Grouping as G+import DataFrame.Internal.Types+import Type.Reflection (+    Typeable,+    typeRep,+ )++import Data.Int (Int16, Int32, Int64, Int8)++-- Specializations for common aggregation types to avoid dictionary overhead.+-- foldLinearGroups: mean accumulator+{-# SPECIALIZE foldLinearGroups ::+    (MeanAcc -> Double -> MeanAcc) ->+    MeanAcc ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (MeanAcc -> Float -> MeanAcc) ->+    MeanAcc ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (MeanAcc -> Int -> MeanAcc) ->+    MeanAcc ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (MeanAcc -> Int8 -> MeanAcc) ->+    MeanAcc ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (MeanAcc -> Int16 -> MeanAcc) ->+    MeanAcc ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (MeanAcc -> Int32 -> MeanAcc) ->+    MeanAcc ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (MeanAcc -> Int64 -> MeanAcc) ->+    MeanAcc ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+-- foldLinearGroups: count accumulator+{-# SPECIALIZE foldLinearGroups ::+    (Int -> Double -> Int) ->+    Int ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (Int -> Float -> Int) ->+    Int ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (Int -> Int -> Int) ->+    Int ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (Int -> Int8 -> Int) ->+    Int ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (Int -> Int16 -> Int) ->+    Int ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (Int -> Int32 -> Int) ->+    Int ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (Int -> Int64 -> Int) ->+    Int ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+-- foldLinearGroups: sum/min/max (acc == elem)+{-# SPECIALIZE foldLinearGroups ::+    (Double -> Double -> Double) ->+    Double ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (Float -> Float -> Float) ->+    Float ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (Int8 -> Int8 -> Int8) ->+    Int8 ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (Int16 -> Int16 -> Int16) ->+    Int16 ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (Int32 -> Int32 -> Int32) ->+    Int32 ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    (Int64 -> Int64 -> Int64) ->+    Int64 ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}++-- mapColumn: finalize+{-# SPECIALIZE mapColumn ::+    (MeanAcc -> Double) -> Column -> Either DataFrameException Column+    #-}+{-# SPECIALIZE mapColumn ::+    (Double -> Double) -> Column -> Either DataFrameException Column+    #-}+{-# SPECIALIZE mapColumn ::+    (Float -> Float) -> Column -> Either DataFrameException Column+    #-}+{-# SPECIALIZE mapColumn ::+    (Int -> Int) -> Column -> Either DataFrameException Column+    #-}++-- zipWithColumns: binary ops+{-# SPECIALIZE zipWithColumns ::+    (Double -> Double -> Double) ->+    Column ->+    Column ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE zipWithColumns ::+    (Float -> Float -> Float) ->+    Column ->+    Column ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE zipWithColumns ::+    (Int -> Int -> Int) -> Column -> Column -> Either DataFrameException Column+    #-}+{-# SPECIALIZE zipWithColumns ::+    (Int8 -> Int8 -> Int8) -> Column -> Column -> Either DataFrameException Column+    #-}+{-# SPECIALIZE zipWithColumns ::+    (Int16 -> Int16 -> Int16) ->+    Column ->+    Column ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE zipWithColumns ::+    (Int32 -> Int32 -> Int32) ->+    Column ->+    Column ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE zipWithColumns ::+    (Int64 -> Int64 -> Int64) ->+    Column ->+    Column ->+    Either DataFrameException Column+    #-}+-- Bool-returning binary comparators (hot path for Expr Bool used in+-- DecisionTree splits)+{-# SPECIALIZE zipWithColumns ::+    (Double -> Double -> Bool) ->+    Column ->+    Column ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE zipWithColumns ::+    (Float -> Float -> Bool) ->+    Column ->+    Column ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE zipWithColumns ::+    (Int -> Int -> Bool) ->+    Column ->+    Column ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE zipWithColumns ::+    (Bool -> Bool -> Bool) ->+    Column ->+    Column ->+    Either DataFrameException Column+    #-}++-- Bool-mapping unary ops (e.g. 'not')+{-# SPECIALIZE mapColumn ::+    (Bool -> Bool) -> Column -> Either DataFrameException Column+    #-}++-------------------------------------------------------------------------------+-- 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++instance (Show a) => Show (Value a) where+    show (Scalar v) = show v+    show (Flat v) = show v+    show (Group v) = show v++-- | 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+{-# INLINEABLE liftValue #-}++{- | Apply a binary function to two 'Value's. When one side is a 'Scalar' the+operation degenerates to 'liftValue', recovering the old @Binary op (Lit l) right@+special cases without explicit pattern matches.+-}+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+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"+{-# INLINEABLE liftValue2 #-}++-- | 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"+{-# INLINEABLE branchValue #-}++{- | 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+    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+    PackedText _ _ -> sliceGroups (materializePacked col) os indices+    BoxedColumn bm vec ->+        let !sorted =+                V.generate+                    (VU.length indices)+                    ((vec `V.unsafeIndex`) . (indices `VU.unsafeIndex`))+         in V.generate nGroups $ \i ->+                BoxedColumn+                    (fmap (bitmapSlice (start i) (len i)) bm)+                    (V.unsafeSlice (start i) (len i) sorted)+    UnboxedColumn bm vec ->+        let !sorted = VU.unsafeBackpermute vec indices+         in V.generate nGroups $ \i ->+                UnboxedColumn+                    (fmap (bitmapSlice (start i) (len i)) bm)+                    (VU.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++-- | Build the inverse of a permutation vector.+invertPermutation :: VU.Vector Int -> VU.Vector Int+invertPermutation perm = VU.create $ do+    let !n = VU.length perm+    inv <- VUM.new n+    VU.imapM_ (flip (VUM.unsafeWrite inv)) perm+    return inv+{-# INLINE invertPermutation #-}++-------------------------------------------------------------------------------+-- promoteColumnWith: unified numeric / text coercion for CastWith+-------------------------------------------------------------------------------++{- | Coerce a column to type @a@, then apply @onResult@ to each element; the handler+selects the mode (like @cast@, @castWithDefault@, or @castEither@). Handles Double/+Float/Int coercion and 'reads'-parses Text; other mismatches return 'Left'.+-}+promoteColumnWith ::+    forall a b.+    (Columnable a, Columnable b, Read a) =>+    (Either String a -> b) -> Column -> Either DataFrameException Column+promoteColumnWith onResult col+    | hasElemType @b col = Right col+    | hasElemType @a col = mapColumn @a (onResult . Right) col+    | Just result <- tryMaybeWrap @a @b onResult col = result+    | otherwise =+        case testEquality (typeRep @a) (typeRep @Double) of+            Just Refl -> promoteToDoubleWith onResult col+            Nothing ->+                case testEquality (typeRep @a) (typeRep @Float) of+                    Just Refl -> promoteToFloatWith onResult col+                    Nothing ->+                        case testEquality (typeRep @a) (typeRep @Int) of+                            Just Refl -> promoteToIntWith onResult col+                            Nothing -> tryParseWith @a onResult col++promoteToDoubleWith ::+    forall b.+    (Columnable b) =>+    (Either String Double -> b) -> Column -> Either DataFrameException Column+promoteToDoubleWith onResult col = case col of+    UnboxedColumn Nothing (v :: VU.Vector c) ->+        case sFloating @c of+            STrue ->+                Right $+                    fromVector @b+                        (V.map (onResult . Right . (realToFrac :: c -> Double)) (VG.convert v))+            SFalse -> case sIntegral @c of+                STrue ->+                    Right $+                        fromVector @b+                            (V.map (onResult . Right . (fromIntegral :: c -> Double)) (VG.convert v))+                SFalse -> castMismatch @c @b+    UnboxedColumn (Just bm) (v :: VU.Vector c) ->+        case sFloating @c of+            STrue ->+                Right $+                    fromVector @b+                        ( V.generate (VU.length v) $ \i ->+                            if bitmapTestBit bm i+                                then onResult (Right (realToFrac (VU.unsafeIndex v i) :: Double))+                                else onResult (Left "null")+                        )+            SFalse -> case sIntegral @c of+                STrue ->+                    Right $+                        fromVector @b+                            ( V.generate (VU.length v) $ \i ->+                                if bitmapTestBit bm i+                                    then onResult (Right (fromIntegral (VU.unsafeIndex v i) :: Double))+                                    else onResult (Left "null")+                            )+                SFalse -> castMismatch @c @b+    BoxedColumn _ _ -> tryParseWith @Double onResult col+    PackedText _ _ -> promoteToDoubleWith onResult (materializePacked col)++promoteToFloatWith ::+    forall b.+    (Columnable b) =>+    (Either String Float -> b) -> Column -> Either DataFrameException Column+promoteToFloatWith onResult col = case col of+    UnboxedColumn Nothing (v :: VU.Vector c) ->+        case sFloating @c of+            STrue ->+                Right $+                    fromVector @b+                        (V.map (onResult . Right . (realToFrac :: c -> Float)) (VG.convert v))+            SFalse -> case sIntegral @c of+                STrue ->+                    Right $+                        fromVector @b+                            (V.map (onResult . Right . (fromIntegral :: c -> Float)) (VG.convert v))+                SFalse -> castMismatch @c @b+    UnboxedColumn (Just bm) (v :: VU.Vector c) ->+        case sFloating @c of+            STrue ->+                Right $+                    fromVector @b+                        ( V.generate (VU.length v) $ \i ->+                            if bitmapTestBit bm i+                                then onResult (Right (realToFrac (VU.unsafeIndex v i) :: Float))+                                else onResult (Left "null")+                        )+            SFalse -> case sIntegral @c of+                STrue ->+                    Right $+                        fromVector @b+                            ( V.generate (VU.length v) $ \i ->+                                if bitmapTestBit bm i+                                    then onResult (Right (fromIntegral (VU.unsafeIndex v i) :: Float))+                                    else onResult (Left "null")+                            )+                SFalse -> castMismatch @c @b+    BoxedColumn _ _ -> tryParseWith @Float onResult col+    PackedText _ _ -> promoteToFloatWith onResult (materializePacked col)++promoteToIntWith ::+    forall b.+    (Columnable b) =>+    (Either String Int -> b) -> Column -> Either DataFrameException Column+promoteToIntWith onResult col = case col of+    UnboxedColumn Nothing (v :: VU.Vector c) ->+        case sFloating @c of+            STrue ->+                Right $+                    fromVector @b+                        (V.map (onResult . Right . (round . (realToFrac :: c -> Double))) (VG.convert v))+            SFalse -> case sIntegral @c of+                STrue ->+                    Right $+                        fromVector @b+                            (V.map (onResult . Right . (fromIntegral :: c -> Int)) (VG.convert v))+                SFalse -> castMismatch @c @b+    UnboxedColumn (Just bm) (v :: VU.Vector c) ->+        case sFloating @c of+            STrue ->+                Right $+                    fromVector @b+                        ( V.generate (VU.length v) $ \i ->+                            if bitmapTestBit bm i+                                then onResult (Right (round (realToFrac (VU.unsafeIndex v i) :: Double)))+                                else onResult (Left "null")+                        )+            SFalse -> case sIntegral @c of+                STrue ->+                    Right $+                        fromVector @b+                            ( V.generate (VU.length v) $ \i ->+                                if bitmapTestBit bm i+                                    then onResult (Right (fromIntegral (VU.unsafeIndex v i) :: Int))+                                    else onResult (Left "null")+                            )+                SFalse -> castMismatch @c @b+    BoxedColumn _ _ -> tryParseWith @Int onResult col+    PackedText _ _ -> promoteToIntWith onResult (materializePacked col)++-- | Single parse primitive: apply @onResult@ to the result of 'reads'.+parseWith :: (Read a) => (Either String a -> b) -> String -> b+parseWith f s = case reads s of+    [(x, "")] -> f (Right x)+    _ -> case reads (show s) of+        [(x, "")] -> f (Right x)+        _ -> f (Left s)++tryParseWith ::+    forall a b.+    (Columnable a, Columnable b, Read a) =>+    (Either String a -> b) -> Column -> Either DataFrameException Column+tryParseWith onResult col = case col of+    PackedText _ _ -> tryParseWith onResult (materializePacked col)+    BoxedColumn bm (v :: V.Vector c) ->+        case testEquality (typeRep @c) (typeRep @String) of+            Just Refl -> case bm of+                Nothing -> Right $ fromVector @b $ V.map (parseWith onResult) v+                Just bitmap ->+                    Right $+                        fromVector @b $+                            V.imap+                                ( \i x ->+                                    if bitmapTestBit bitmap i then parseWith onResult x else onResult (Left "null")+                                )+                                v+            Nothing ->+                case testEquality (typeRep @c) (typeRep @T.Text) of+                    Just Refl -> case bm of+                        Nothing -> Right $ fromVector @b $ V.map (parseWith onResult . T.unpack) v+                        Just bitmap ->+                            Right $+                                fromVector @b $+                                    V.imap+                                        ( \i x ->+                                            if bitmapTestBit bitmap i+                                                then parseWith onResult (T.unpack x)+                                                else onResult (Left "null")+                                        )+                                        v+                    Nothing -> castMismatch @c @b+    UnboxedColumn bm (v :: VU.Vector c) -> case bm of+        Nothing -> Right $ fromVector @b $ V.map (parseWith onResult . show) (V.convert v)+        Just bitmap ->+            Right $+                fromVector @b $+                    V.imap+                        ( \i x ->+                            if bitmapTestBit bitmap i+                                then parseWith onResult (show x)+                                else onResult (Left "null")+                        )+                        (V.convert v)++{- | When output type @b@ is @Maybe c@ (or @Maybe (Maybe c)@) and the column stores+plain @c@, wrap each element in 'Just' (the double-Maybe case collapses to a single+@Maybe c@). Returns 'Nothing' when neither condition holds.+-}+tryMaybeWrap ::+    forall a b.+    (Columnable a, Columnable b) =>+    (Either String a -> b) -> Column -> Maybe (Either DataFrameException Column)+tryMaybeWrap _onResult col = case col of+    UnboxedColumn Nothing (v :: VU.Vector c) ->+        let wrapped = V.map Just (VG.convert v) :: V.Vector (Maybe c)+         in case testEquality (typeRep @b) (typeRep @(Maybe c)) of+                Just Refl -> Just $ Right $ fromVector @b wrapped+                Nothing ->+                    case testEquality (typeRep @b) (typeRep @(Maybe (Maybe c))) of+                        Just _ -> Just $ Right $ fromVector @(Maybe c) wrapped+                        Nothing -> Nothing+    BoxedColumn Nothing (v :: V.Vector c) ->+        let wrapped = V.map Just v :: V.Vector (Maybe c)+         in case testEquality (typeRep @b) (typeRep @(Maybe c)) of+                Just Refl -> Just $ Right $ fromVector @b wrapped+                Nothing ->+                    case testEquality (typeRep @b) (typeRep @(Maybe (Maybe c))) of+                        Just _ -> Just $ Right $ fromVector @(Maybe c) wrapped+                        Nothing -> Nothing+    _ -> Nothing++castMismatch ::+    forall src tgt.+    (Typeable src, Typeable tgt) =>+    Either DataFrameException Column+castMismatch =+    Left $+        TypeMismatchException+            MkTypeErrorContext+                { userType = Right (typeRep @tgt)+                , expectedType = Right (typeRep @src)+                , callingFunctionName = Just "cast"+                , errorColumnName = Nothing+                }++-------------------------------------------------------------------------------+-- 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)+eval _ (Lit v) = Right (Scalar v)+eval (FlatCtx df) (Col name) =+    case getColumn name df of+        Nothing ->+            Left $ ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)+        Just c+            | hasElemType @a c -> Right (Flat c)+            | otherwise ->+                Left $+                    TypeMismatchException+                        ( MkTypeErrorContext+                            { userType = Right (typeRep @a)+                            , expectedType = Left (columnTypeString c)+                            , errorColumnName = Just (T.unpack name)+                            , callingFunctionName = Just "col"+                            } ::+                            TypeErrorContext a ()+                        )+eval (GroupCtx gdf) (Col name) =+    case getColumn name (fullDataframe gdf) of+        Nothing ->+            Left $+                ColumnsNotFoundException+                    [name]+                    ""+                    (M.keys $ columnIndices $ fullDataframe gdf)+        Just c+            | hasElemType @a c ->+                Right (Group (sliceGroups c (offsets gdf) (valueIndices gdf)))+            | otherwise ->+                Left $+                    TypeMismatchException+                        ( MkTypeErrorContext+                            { userType = Right (typeRep @a)+                            , expectedType = Left (columnTypeString c)+                            , errorColumnName = Just (T.unpack name)+                            , callingFunctionName = Just "col"+                            } ::+                            TypeErrorContext a ()+                        )+eval (FlatCtx df) (CastWith name _tag onResult) =+    case getColumn name df of+        Nothing ->+            Left $+                ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)+        Just c -> Flat <$> promoteColumnWith onResult c+eval (GroupCtx gdf) (CastWith name _tag onResult) =+    case getColumn name (fullDataframe gdf) of+        Nothing ->+            Left $+                ColumnsNotFoundException+                    [name]+                    ""+                    (M.keys $ columnIndices $ fullDataframe gdf)+        Just c -> do+            promoted <- promoteColumnWith onResult c+            Right $ Group (sliceGroups promoted (offsets gdf) (valueIndices gdf))+eval ctx (CastExprWith _tag onResult (inner :: Expr src)) = do+    v <- eval @src ctx inner+    case v of+        Scalar s ->+            Flat <$> promoteColumnWith onResult (fromList @src [s])+        Flat col ->+            Flat <$> promoteColumnWith onResult col+        Group gs ->+            Group <$> V.mapM (promoteColumnWith onResult) gs+eval ctx expr@(Unary op (inner :: Expr b)) = addContext expr $ do+    v <- eval @b ctx inner+    liftValue (unaryFn op) v+eval ctx expr@(Binary op (left :: Expr c) (right :: Expr b)) =+    addContext expr $ do+        l <- eval @c ctx left+        r <- eval @b ctx right+        liftValue2 (binaryFn op) l r+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+eval (FlatCtx df) expr@(Over keys inner) = addContext expr $ do+    let gdf = G.groupBy keys df+    v <- eval (GroupCtx gdf) inner+    case v of+        Scalar s ->+            Right (Scalar s)+        Flat groupCol ->+            Right (Flat (atIndicesStable (rowToGroup gdf) groupCol))+        Group groupCols -> do+            sorted <- V.fold1M' concatColumns groupCols+            let inv = invertPermutation (valueIndices gdf)+            Right (Flat (atIndicesStable inv sorted))+eval (GroupCtx _) expr@(Over _ _) =+    addContext expr $+        Left+            ( InternalException+                "Over (window function) is not supported inside a grouped context"+            )+-- 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 $+                    ColumnsNotFoundException+                        [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 $+                            ColumnsNotFoundException+                                [name]+                                ""+                                (M.keys $ columnIndices $ fullDataframe gdf)+                    Just col ->+                        Flat <$> 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 $+                        ColumnsNotFoundException+                            [name]+                            ""+                            (M.keys $ columnIndices $ fullDataframe gdf)+                Just col ->+                    Flat+                        <$> ( foldLinearGroups @b step seed col (rowToGroup gdf) (numGroups gdf)+                                >>= mapColumn finalize+                            )+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+eval ctx expr@(Agg (FoldAgg _ (Just seed) (f :: a -> b -> a)) inner) =+    addContext expr $ do+        v <- eval @b ctx inner+        case v of+            Scalar x -> Right (broadcastFold ctx seed f x)+            Flat col ->+                Scalar <$> foldlColumn @b @a f seed col+            Group gs ->+                Flat . fromVector+                    <$> V.mapM (foldlColumn @b @a f seed) gs+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 x -> case broadcastFold ctx seed step x of+                    Scalar acc -> Right (Scalar (finalize acc))+                    Flat col -> Flat <$> mapColumn @acc @a finalize col+                    Group _ ->+                        Left+                            ( InternalException+                                "broadcastFold unexpectedly produced a Group value"+                            )+                Flat col ->+                    Scalar . finalize <$> foldlColumn @b step seed col+                Group gs ->+                    Flat . fromVector+                        <$> V.mapM (fmap finalize . foldlColumn @b step seed) gs+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+                                "fold1 requires at least one element"+                    Flat col ->+                        Scalar <$> foldl1Column @a f col+                    Group gs ->+                        Flat . fromVector+                            <$> V.mapM (foldl1Column @a f) gs++broadcastFold ::+    forall acc b.+    (Columnable acc) =>+    Ctx -> acc -> (acc -> b -> acc) -> b -> Value acc+broadcastFold (FlatCtx df) seed step x =+    let n = fst (dataframeDimensions df)+     in Scalar (iterateStep n step seed x)+broadcastFold (GroupCtx gdf) seed step x =+    let offs = offsets gdf+        ng = VU.length offs - 1+        results =+            V.generate ng $ \i ->+                let sz = offs VU.! (i + 1) - offs VU.! i+                 in iterateStep sz step seed x+     in Flat (fromVector results)++iterateStep :: Int -> (acc -> b -> acc) -> acc -> b -> acc+iterateStep n step = go n+  where+    go 0 !acc _ = acc+    go k !acc x = go (k - 1) (step acc x) x++{- | 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++{- | 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.+Calls 'eval' then 'materialize'; 'Lit' values are broadcast here at the boundary+rather than eagerly.+-}+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 _ ->+            Right $ UnAggregated $ BoxedColumn @T.Text Nothing V.empty