packages feed

dataframe-core 1.1.0.4 → 2.5.0.1

raw patch · 75 files changed

Files

dataframe-core.cabal view
@@ -1,6 +1,6 @@-cabal-version:      2.4+cabal-version:      3.4 name:               dataframe-core-version:            1.1.0.4+version:            2.5.0.1 synopsis:           Core data structures for the dataframe library. description:     Minimal interchange-format types for the @dataframe@ ecosystem:@@ -31,46 +31,63 @@ library     import:             warnings     exposed-modules:+                        DataFrame.Core                         DataFrame.Errors-                        DataFrame.Operators+                        DataFrame.Expression.Operators+                        DataFrame.Typed.Freeze+                        DataFrame.Typed.Generic+                        DataFrame.Typed.Record+                        DataFrame.Typed.Schema+                        DataFrame.Typed.Types+                        DataFrame.Typed.Util                         DataFrame.Display.Terminal.Colours                         DataFrame.Display.Terminal.PrettyPrint-                        DataFrame.Internal.AggKernel-                        DataFrame.Internal.AggKernelDirect-                        DataFrame.Internal.AggKernelPar-                        DataFrame.Internal.AggPlan-                        DataFrame.Internal.GroupingDirect+                        -- Concurrency+                        DataFrame.Internal.Control.Concurrent+                        -- Aggregation kernels and planning.+                        DataFrame.Internal.Aggregation.Kernel.Dense+                        DataFrame.Internal.Aggregation.Kernel.Fused+                        DataFrame.Internal.Aggregation.Kernel.Moments+                        DataFrame.Internal.Aggregation.Kernel.Scatter+                        DataFrame.Internal.Aggregation.Plan+                        DataFrame.Internal.Aggregation.Reduction+                        -- Shared algorithms.+                        DataFrame.Internal.Algorithms.Hash+                        DataFrame.Internal.Algorithms.Rank.Radix+                        DataFrame.Internal.Algorithms.Sort.Radix.Parallel+                        -- The column representation.                         DataFrame.Internal.Column-                        DataFrame.Internal.ColumnBuilder-                        DataFrame.Internal.ColumnMerge+                        DataFrame.Internal.Column.Base+                        DataFrame.Internal.Column.Bitmap+                        DataFrame.Internal.Column.Builder+                        DataFrame.Internal.Column.Conversion+                        DataFrame.Internal.Column.Encode+                        DataFrame.Internal.Column.Merge+                        DataFrame.Internal.Column.Operations+                        DataFrame.Internal.Column.Properties+                        DataFrame.Internal.Column.Types+                        -- Backing data structures.+                        DataFrame.Internal.Data.HashTable+                        DataFrame.Internal.Data.PackedText+                        DataFrame.Internal.Data.PackedText.Utf8                         DataFrame.Internal.DataFrame-                        DataFrame.Internal.DictEncode+                        DataFrame.Internal.Display.Pretty+                        -- Expressions: syntax, operators, simplification.                         DataFrame.Internal.Expression+                        DataFrame.Internal.Expression.Operators+                        DataFrame.Internal.Expression.Operators.Nullable+                        DataFrame.Internal.Expression.Simplify                         DataFrame.Internal.Grouping-                        DataFrame.Internal.GroupingPar-                        DataFrame.Internal.Hash-                        DataFrame.Internal.HashTable+                        DataFrame.Internal.Grouping.Direct+                        DataFrame.Internal.Grouping.Partitioned                         DataFrame.Internal.Interpreter-                        DataFrame.Internal.Nullable-                        DataFrame.Internal.PackedText-                        DataFrame.Internal.ParRadixSort-                        DataFrame.Internal.RadixRank-                        DataFrame.Internal.RowHash                         DataFrame.Internal.Row-                        DataFrame.Internal.Simplify-                        DataFrame.Internal.Types-                        DataFrame.Internal.Utf8-                        DataFrame.Typed.Freeze-                        DataFrame.Typed.Generic-                        DataFrame.Typed.Record-                        DataFrame.Typed.Schema-                        DataFrame.Typed.Types-                        DataFrame.Typed.Util+                        DataFrame.Internal.Row.RowHash     build-depends:      base >= 4 && < 5,-                        containers >= 0.6.7 && < 0.9,-                        primitive >= 0.7 && < 0.10,+                        containers >= 0.6.7 && < 0.10,+                        primitive >= 0.7 && < 0.11,                         random >= 1 && < 2,                         text >= 2.1 && < 3,-                        vector ^>= 0.13-    hs-source-dirs:     src+                        vector >= 0.13 && < 0.15+    hs-source-dirs:     src, src-internal     default-language:   Haskell2010
+ src-internal/DataFrame/Display/Terminal/Colours.hs view
@@ -0,0 +1,7 @@+module DataFrame.Display.Terminal.Colours where++red, green, brightGreen, brightBlue :: String -> String+red = id+green = id+brightGreen = id+brightBlue = id
+ src-internal/DataFrame/Display/Terminal/PrettyPrint.hs view
@@ -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 "|" "\\|"
+ src-internal/DataFrame/Errors.hs view
@@ -0,0 +1,189 @@+{-# 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+    ExpectedNonNullableException :: 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 ExpectedNonNullableException = "Expected non-nullable column"+    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)+            ]
+ src-internal/DataFrame/Internal/Aggregation/Kernel/Dense.hs view
@@ -0,0 +1,832 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | The low-cardinality DENSE reduction kernel: when the group domain is small,+the grouping layer's @rowToGroup@ already maps row -> group, so the reduction+scatters straight off it with no @valueIndices@ gather.++Parallel by ROW range with a private per-worker accumulator of @nGroups@ slots,+merged afterwards — which is why it needs a small domain, and why it admits only+order-independent reductions: the merge must be exact for the result to stay+byte-identical to @-N1@. Anything it rejects falls back to+"DataFrame.Internal.Aggregation.Kernel.Scatter".++The caller decides whether the domain is small enough; see @denseThreshold@ in+the operations layer.+-}+module DataFrame.Internal.Aggregation.Kernel.Dense (+    denseReduce,+    denseMaxMinusMin,+) where++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.Aggregation.Reduction (Reduction (..))+import DataFrame.Internal.Column (+    Column (..),+    fromUnboxedVector,+    materializePacked,+ )+import DataFrame.Internal.Control.Concurrent (+    capabilities,+    parThreshold,+    parallelChunks,+    shouldParallelize,+ )++{- | Group count at or below which the float (Double sum/mean) direct+reductions run as ONE sequential row-order pass: the accumulators fit in L2, the+pass is memory-bandwidth bound, and adding each value to its group in ascending+row order is exactly the order the group-range gather kernel uses (the grouping+layer's @valueIndices@ is a stable counting sort), so the result is+byte-identical to it. Above this the parallel chunked variant runs instead (see+'sumDblDense'). Var/std are never taken directly: any chunked merge of+variance state changes the float recurrence, and the group-range gather kernel+already runs them in parallel while replaying the interpreter's per-group+update order bit-for-bit.+-}+seqFloatGroups :: Int+seqFloatGroups = 65536++{- | Run a recognised reduction through the direct-indexed path. 'Nothing' (so+the caller falls back to the order-preserving kernel) unless the reduction is+admitted at this element type AND the column is a clean unboxed Int/Double.+-}+denseReduce :: Reduction -> VU.Vector Int -> Int -> Column -> Maybe Column+denseReduce red g nGroups col = case col of+    UnboxedColumn Nothing (v :: VU.Vector a) ->+        case testEquality (typeRep @a) (typeRep @Int) of+            Just Refl -> denseInt red g nGroups v+            Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+                Just Refl -> denseDouble red g nGroups v+                Nothing -> Nothing+    p@(PackedText _ _) -> denseReduce red g nGroups (materializePacked p)+    _ -> Nothing+{-# INLINEABLE denseReduce #-}++{- | The reductions admitted over an Int column. Sum/min/max/mean/count are+exact in the Int domain (any merge order gives the same bits); top2sum selects+the two largest values (order-independent as a multiset selection) and only+adds them once at finalize. Var/std stay with the group-range gather kernel:+its per-group Welford recurrence replays the interpreter's update order+bit-for-bit, which no chunk-merged direct pass can.+-}+denseInt :: Reduction -> VU.Vector Int -> Int -> VU.Vector Int -> Maybe Column+denseInt red g nGroups v = case red of+    RCount -> Just (fromUnboxedVector (countDense g nGroups (VU.length v)))+    RSum -> Just (fromUnboxedVector (sumIntDense g nGroups v))+    RMin -> Just (fromUnboxedVector (extremaIntDense True g nGroups v))+    RMax -> Just (fromUnboxedVector (extremaIntDense False g nGroups v))+    RMean -> Just (fromUnboxedVector (meanIntDense g nGroups v))+    RTop2Sum -> Just (fromUnboxedVector (top2Dense g nGroups v))+    RTop2Snd -> Just (fromUnboxedVector (top2SndDense g nGroups v))+    _ -> Nothing++{- | The reductions admitted over a Double column. Count/min/max/top2sum are+order-independent (exact per-worker merge, byte-identical at any @-N@). The+float sum/mean run sequentially in row order below 'seqFloatGroups' (matching+the gather kernel's per-group addition order exactly) and as deterministic+chunked partials above it. Var/std keep the gather kernel (see 'denseInt').+-}+denseDouble ::+    Reduction -> VU.Vector Int -> Int -> VU.Vector Double -> Maybe Column+denseDouble red g nGroups v = case red of+    RCount -> Just (fromUnboxedVector (countDense g nGroups (VU.length v)))+    RSum -> Just (fromUnboxedVector (sumDblDense g nGroups v))+    RMean -> Just (fromUnboxedVector (meanDblDense g nGroups v))+    RMin -> Just (fromUnboxedVector (extremaDblDense True g nGroups v))+    RMax -> Just (fromUnboxedVector (extremaDblDense False g nGroups v))+    RTop2Sum -> Just (fromUnboxedVector (top2Dense g nGroups v))+    RTop2Snd -> Just (fromUnboxedVector (top2SndDense g nGroups v))+    _ -> Nothing++{- | The fused @max a - min b@ direct pass: BOTH extrema accumulate in one+streaming loop over the rows (min/max are order-independent, so the per-worker+merge is exact and the result byte-identical to the two gather passes it+replaces). 'Nothing' unless both columns are clean unboxed and same-typed+(Int/Int keeps the Int result of the interpreter; Double/Double the Double one);+mixed pairs keep the gather fallback.+-}+denseMaxMinusMin :: VU.Vector Int -> Int -> Column -> Column -> Maybe Column+denseMaxMinusMin g nGroups ca cb = case (ca, cb) of+    ( UnboxedColumn Nothing (va :: VU.Vector x)+        , UnboxedColumn Nothing (vb :: VU.Vector y)+        )+            | Just Refl <- testEquality (typeRep @x) (typeRep @Int)+            , Just Refl <- testEquality (typeRep @y) (typeRep @Int) ->+                Just (fromUnboxedVector (maxMinusMinDenseInt g nGroups va vb))+            | Just Refl <- testEquality (typeRep @x) (typeRep @Double)+            , Just Refl <- testEquality (typeRep @y) (typeRep @Double) ->+                Just (fromUnboxedVector (maxMinusMinDenseDbl g nGroups va vb))+    _ -> Nothing+{-# INLINEABLE denseMaxMinusMin #-}++{- | Monomorphic entry points: the 'testEquality' dispatch above only yields an+unsafe coercion, so a direct call to the polymorphic 'maxMinusMinDense' there+would stay at the abstract element type and never meet its SPECIALIZE rules+(measured ~3x on the whole pass); calling through these fixed-type wrappers+(the coercion lands on the argument) does.+-}+maxMinusMinDenseInt ::+    VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int+maxMinusMinDenseInt = maxMinusMinDense minBound maxBound+{-# NOINLINE maxMinusMinDenseInt #-}++maxMinusMinDenseDbl ::+    VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double -> VU.Vector Double+maxMinusMinDenseDbl = maxMinusMinDense (negate (1 / 0)) (1 / 0)+{-# NOINLINE maxMinusMinDenseDbl #-}++-- | Whether to fan out at this row count.+shouldPar :: Int -> Bool+shouldPar = shouldParallelize parThreshold++{- | 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.+The chunking is a fixed function of @n@ and @caps@, so any merge over the+partials is deterministic at a given @-N@.+-}+runPartialsOver ::+    Int -> Int -> (Int -> Int -> IO acc) -> IO [acc]+runPartialsOver n _caps = parallelChunks parThreshold n++-------------------------------------------------------------------------------+-- Count (order-independent: per-group row count)+-------------------------------------------------------------------------------++countDense :: VU.Vector Int -> Int -> Int -> VU.Vector Int+countDense g nGroups n = unsafePerformIO $ do+    parts <- runPartialsOver n capabilities (countChunk g nGroups)+    mergeIntSum nGroups parts+{-# NOINLINE countDense #-}++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)+-------------------------------------------------------------------------------++sumIntDense :: VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int+sumIntDense g nGroups v = unsafePerformIO $ do+    parts <- runPartialsOver (VU.length v) capabilities (sumIntChunk g v nGroups)+    mergeIntSum nGroups parts+{-# NOINLINE sumIntDense #-}++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++-------------------------------------------------------------------------------+-- Double sum / mean (streaming scatter; chunked partials above seqFloatGroups)+-------------------------------------------------------------------------------++{- | Double group sums. At or below 'seqFloatGroups' a single sequential pass in+ascending row order — each group's additions happen in exactly the order the+group-range gather kernel performs them, so the result is byte-identical to it.+Above that, per-worker chunk partials merged in worker order: still+deterministic at a fixed @-N@, but the float summation order differs from the+sequential pass.+-}+sumDblDense :: VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double+sumDblDense g nGroups v+    | nGroups <= seqFloatGroups || not (shouldPar n) =+        unsafePerformIO (sumDblChunk g v nGroups 0 n >>= VU.unsafeFreeze)+    | otherwise = unsafePerformIO $ do+        parts <- runPartialsOver n capabilities (sumDblChunk g v nGroups)+        mergeDblSum nGroups parts+  where+    !n = VU.length v+{-# NOINLINE sumDblDense #-}++sumDblChunk ::+    VU.Vector Int ->+    VU.Vector Double ->+    Int ->+    Int ->+    Int ->+    IO (VUM.IOVector Double)+sumDblChunk g v nGroups lo hi = do+    acc <- VUM.replicate nGroups (0 :: Double)+    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++{- | Double mean: fused (Double sum, count) per group, divided once at finalize.+Same order policy as 'sumDblDense' (sequential row order is byte-identical to+the gather kernel; the chunked variant changes the float summation order).+-}+meanDblDense :: VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double+meanDblDense g nGroups v+    | nGroups <= seqFloatGroups || not (shouldPar n) = unsafePerformIO $ do+        (s, c) <- meanDblChunk g v nGroups 0 n+        finalizeMeanDbl nGroups s c+    | otherwise = unsafePerformIO $ do+        parts <- runPartialsOver n capabilities (meanDblChunk g v nGroups)+        (s, c) <- mergeMeanDbl nGroups parts+        finalizeMeanDbl nGroups s c+  where+    !n = VU.length v+{-# NOINLINE meanDblDense #-}++meanDblChunk ::+    VU.Vector Int ->+    VU.Vector Double ->+    Int ->+    Int ->+    Int ->+    IO (VUM.IOVector Double, VUM.IOVector Int)+meanDblChunk g v nGroups lo hi = do+    s <- VUM.replicate nGroups (0 :: Double)+    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)++finalizeMeanDbl ::+    Int -> VUM.IOVector Double -> VUM.IOVector Int -> IO (VU.Vector Double)+finalizeMeanDbl 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 sv / fromIntegral cv)+                go (k + 1)+    go 0+    VU.unsafeFreeze out++-------------------------------------------------------------------------------+-- Integer min / max (order-independent)+-------------------------------------------------------------------------------++extremaIntDense ::+    Bool -> VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int+extremaIntDense isMin g nGroups v = unsafePerformIO $ do+    parts <-+        runPartialsOver (VU.length v) capabilities (extremaIntChunk isMin g v nGroups)+    mergeExtremaInt isMin nGroups parts+{-# NOINLINE extremaIntDense #-}++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++-------------------------------------------------------------------------------+-- Double min / max (order-independent: exact per-worker merge)+-------------------------------------------------------------------------------++extremaDblDense ::+    Bool -> VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double+extremaDblDense isMin g nGroups v+    | not (shouldPar n) =+        unsafePerformIO (extremaDblChunk isMin g v nGroups 0 n >>= VU.unsafeFreeze)+    | otherwise = unsafePerformIO $ do+        parts <- runPartialsOver n capabilities (extremaDblChunk isMin g v nGroups)+        mergeExtremaDbl isMin nGroups parts+  where+    !n = VU.length v+{-# NOINLINE extremaDblDense #-}++extremaDblChunk ::+    Bool ->+    VU.Vector Int ->+    VU.Vector Double ->+    Int ->+    Int ->+    Int ->+    IO (VUM.IOVector Double)+extremaDblChunk isMin g v nGroups lo hi = do+    let !seed = if isMin then 1 / 0 else negate (1 / 0)+        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++-------------------------------------------------------------------------------+-- Fused max(a) - min(b) (order-independent: exact per-worker merge)+-------------------------------------------------------------------------------++{- | One streaming pass accumulating @max a@ and @min b@ together. @maxSeed@ is+the identity of @max@ (the type's least value), @minSeed@ of @min@ (its+greatest).+-}+maxMinusMinDense ::+    (VU.Unbox a, Num a, Ord a) =>+    a ->+    a ->+    VU.Vector Int ->+    Int ->+    VU.Vector a ->+    VU.Vector a ->+    VU.Vector a+{-# SPECIALIZE maxMinusMinDense ::+    Int ->+    Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Int+    #-}+{-# SPECIALIZE maxMinusMinDense ::+    Double ->+    Double ->+    VU.Vector Int ->+    Int ->+    VU.Vector Double ->+    VU.Vector Double ->+    VU.Vector Double+    #-}+maxMinusMinDense maxSeed minSeed g nGroups va vb+    | not (shouldPar n) = unsafePerformIO $ do+        (mx, mn) <- maxMinusMinChunk maxSeed minSeed g va vb nGroups 0 n+        finalizeMaxMinusMin nGroups mx mn+    | otherwise = unsafePerformIO $ do+        parts <-+            runPartialsOver+                n+                capabilities+                (maxMinusMinChunk maxSeed minSeed g va vb nGroups)+        (mx, mn) <- mergeMaxMin nGroups parts+        finalizeMaxMinusMin nGroups mx mn+  where+    !n = VU.length va+{- INLINEABLE (not NOINLINE) so the SPECIALIZE pragmas above take effect; the+kernel is a pure function of its arguments, so the usual unsafePerformIO+sharing concern does not apply. -}+{-# INLINEABLE maxMinusMinDense #-}++maxMinusMinChunk ::+    (VU.Unbox a, Ord a) =>+    a ->+    a ->+    VU.Vector Int ->+    VU.Vector a ->+    VU.Vector a ->+    Int ->+    Int ->+    Int ->+    IO (VUM.IOVector a, VUM.IOVector a)+maxMinusMinChunk maxSeed minSeed g va vb nGroups lo hi = do+    mx <- VUM.replicate nGroups maxSeed+    mn <- VUM.replicate nGroups minSeed+    let go !i+            | i >= hi = pure ()+            | otherwise = do+                let !k = VU.unsafeIndex g i+                cx <- VUM.unsafeRead mx k+                VUM.unsafeWrite mx k (max cx (VU.unsafeIndex va i))+                cn <- VUM.unsafeRead mn k+                VUM.unsafeWrite mn k (min cn (VU.unsafeIndex vb i))+                go (i + 1)+    go lo+    pure (mx, mn)++mergeMaxMin ::+    (VU.Unbox a, Ord a) =>+    Int ->+    [(VUM.IOVector a, VUM.IOVector a)] ->+    IO (VUM.IOVector a, VUM.IOVector a)+mergeMaxMin nGroups parts = case parts of+    [] -> error "mergeMaxMin: no partials"+    ((mx0, mn0) : rest) -> do+        let add (mx, mn) = do+                let go !k+                        | k >= nGroups = pure ()+                        | otherwise = do+                            xa <- VUM.unsafeRead mx0 k+                            xb <- VUM.unsafeRead mx k+                            VUM.unsafeWrite mx0 k (max xa xb)+                            na <- VUM.unsafeRead mn0 k+                            nb <- VUM.unsafeRead mn k+                            VUM.unsafeWrite mn0 k (min na nb)+                            go (k + 1)+                go 0+        mapM_ add rest+        pure (mx0, mn0)++finalizeMaxMinusMin ::+    (VU.Unbox a, Num a) =>+    Int ->+    VUM.IOVector a ->+    VUM.IOVector a ->+    IO (VU.Vector a)+finalizeMaxMinusMin nGroups mx mn = do+    out <- VUM.new nGroups+    let go !k+            | k >= nGroups = pure ()+            | otherwise = do+                a <- VUM.unsafeRead mx k+                b <- VUM.unsafeRead mn k+                VUM.unsafeWrite out k (a - b)+                go (k + 1)+    go 0+    VU.unsafeFreeze out++-------------------------------------------------------------------------------+-- Top-2 sum (order-independent multiset selection; one float add at finalize)+-------------------------------------------------------------------------------++{- | Sum of the two largest values per group. Each accumulator holds the+(largest, second-largest) pair seen so far; merging two pairs keeps the top two+of the four candidates. No float ADDITION happens until the single @m1 + m2@ at+finalize, so the result is byte-identical to the gather kernel regardless of+chunking. Mirrors the gather kernel exactly, including the @realToFrac@ per+element and the @-inf -> 0@ guards at finalize.+-}+top2Dense ::+    (VU.Unbox a, Real a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double+{-# SPECIALIZE top2Dense ::+    VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Double+    #-}+{-# SPECIALIZE top2Dense ::+    VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double+    #-}+top2Dense g nGroups v+    | not (shouldPar n) = unsafePerformIO $ do+        (m1, m2) <- top2Chunk g v nGroups 0 n+        finalizeTop2 nGroups m1 m2+    | otherwise = unsafePerformIO $ do+        parts <- runPartialsOver n capabilities (top2Chunk g v nGroups)+        (m1, m2) <- mergeTop2 nGroups parts+        finalizeTop2 nGroups m1 m2+  where+    !n = VU.length v+{- INLINEABLE (not NOINLINE) so the SPECIALIZE pragmas above take effect; pure+function of its arguments, so unsafePerformIO sharing is not a concern. -}+{-# INLINEABLE top2Dense #-}++{- | Second-largest value per group: the exact same per-worker+(largest, second-largest) accumulator and merge as 'top2Dense'+('top2Chunk'/'mergeTop2'), finalized to the second max alone. A group of+size < 2 finalizes its @-inf@ seed to NaN — documented behaviour (the+db-benchmark Q8 data has no size-1 @id6@ groups). Order-independent multiset+selection, so byte-identical at any @-N@.+-}+top2SndDense ::+    (VU.Unbox a, Real a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double+{-# SPECIALIZE top2SndDense ::+    VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Double+    #-}+{-# SPECIALIZE top2SndDense ::+    VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double+    #-}+top2SndDense g nGroups v+    | not (shouldPar n) = unsafePerformIO $ do+        (m1, m2) <- top2Chunk g v nGroups 0 n+        finalizeTop2Snd nGroups m1 m2+    | otherwise = unsafePerformIO $ do+        parts <- runPartialsOver n capabilities (top2Chunk g v nGroups)+        (m1, m2) <- mergeTop2 nGroups parts+        finalizeTop2Snd nGroups m1 m2+  where+    !n = VU.length v+{- INLINEABLE (not NOINLINE) so the SPECIALIZE pragmas above take effect; pure+function of its arguments, so unsafePerformIO sharing is not a concern. -}+{-# INLINEABLE top2SndDense #-}++top2Chunk ::+    (VU.Unbox a, Real a) =>+    VU.Vector Int ->+    VU.Vector a ->+    Int ->+    Int ->+    Int ->+    IO (VUM.IOVector Double, VUM.IOVector Double)+top2Chunk g v nGroups lo hi = do+    let ninf = negate (1 / 0) :: Double+    m1 <- VUM.replicate nGroups ninf+    m2 <- VUM.replicate nGroups ninf+    let go !i+            | i >= hi = 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 lo+    pure (m1, m2)++{- | Top two of the four candidates @{a1, a2, b1, b2}@ per group (each pair+already ordered @m1 >= m2@, @-inf@ seeds included).+-}+mergeTop2 ::+    Int ->+    [(VUM.IOVector Double, VUM.IOVector Double)] ->+    IO (VUM.IOVector Double, VUM.IOVector Double)+mergeTop2 nGroups parts = case parts of+    [] -> error "mergeTop2: no partials"+    ((m10, m20) : rest) -> do+        let add (m1, m2) = do+                let go !k+                        | k >= nGroups = pure ()+                        | otherwise = do+                            a1 <- VUM.unsafeRead m10 k+                            a2 <- VUM.unsafeRead m20 k+                            b1 <- VUM.unsafeRead m1 k+                            b2 <- VUM.unsafeRead m2 k+                            if b1 > a1+                                then do+                                    VUM.unsafeWrite m10 k b1+                                    VUM.unsafeWrite m20 k (max a1 b2)+                                else VUM.unsafeWrite m20 k (max a2 b1)+                            go (k + 1)+                go 0+        mapM_ add rest+        pure (m10, m20)++finalizeTop2 ::+    Int -> VUM.IOVector Double -> VUM.IOVector Double -> IO (VU.Vector Double)+finalizeTop2 nGroups m1 m2 = do+    out <- VUM.new nGroups+    let go !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+                go (k + 1)+    go 0+    VU.unsafeFreeze out++finalizeTop2Snd ::+    Int -> VUM.IOVector Double -> VUM.IOVector Double -> IO (VU.Vector Double)+finalizeTop2Snd nGroups _m1 m2 = do+    out <- VUM.new nGroups+    let go !k+            | k >= nGroups = pure ()+            | otherwise = do+                a2 <- VUM.unsafeRead m2 k+                VUM.unsafeWrite out k (if isInfinite a2 then 0 / 0 else a2)+                go (k + 1)+    go 0+    VU.unsafeFreeze out++-------------------------------------------------------------------------------+-- 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@.+-}+meanIntDense :: VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Double+meanIntDense g nGroups v+    | not (shouldPar n) = unsafePerformIO $ do+        (s, c) <- meanIntChunk g v nGroups 0 n+        finalizeMeanInt nGroups s c+    | otherwise = unsafePerformIO $ do+        parts <- runPartialsOver n capabilities (meanIntChunk g v nGroups)+        (s, c) <- mergePair nGroups parts+        finalizeMeanInt nGroups s c+  where+    !n = VU.length v+{-# NOINLINE meanIntDense #-}++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++{- | Sum the per-worker Double partials into the first worker's accumulator IN+WORKER ORDER: deterministic at a fixed @-N@, but the float summation order is+chunk-major rather than the sequential row order.+-}+mergeDblSum :: Int -> [VUM.IOVector Double] -> IO (VU.Vector Double)+mergeDblSum 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)++{- | As 'mergePair' but for the Double (sum, count) partials of the Double mean;+worker-order float sums (see 'mergeDblSum').+-}+mergeMeanDbl ::+    Int ->+    [(VUM.IOVector Double, VUM.IOVector Int)] ->+    IO (VUM.IOVector Double, VUM.IOVector Int)+mergeMeanDbl 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++mergeExtremaDbl :: Bool -> Int -> [VUM.IOVector Double] -> IO (VU.Vector Double)+mergeExtremaDbl isMin nGroups parts = case parts of+    [] ->+        VU.unsafeFreeze+            =<< VUM.replicate nGroups (if isMin then 1 / 0 else negate (1 / 0))+    (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
+ src-internal/DataFrame/Internal/Aggregation/Kernel/Fused.hs view
@@ -0,0 +1,757 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | Fused multi-reduction aggregation kernels: several reductions over the+SAME grouping, evaluated in one pass over the rows.++Two shapes, picked by group-domain width:++* 'mkFusedAgg' \/ 'runFusedAggs' — the STREAMING pass, driven by @rowToGroup@.+  No @valueIndices@ and no placement pass; each worker keeps a private+  accumulator per reduction and the partials merge over group slices. Capped at+  'streamGroupCap', above which the per-worker arrays overflow cache.++* 'mkGatherAgg' \/ 'runGatherAggs' — the GATHER pass, driven by the grouped+  @(valueIndices, offsets)@ layout, for group domains too wide to stream.++Both amortize the memory traffic of the grouping across every reduction in the+batch instead of re-reading it once per reduction.+-}+module DataFrame.Internal.Aggregation.Kernel.Fused (+    FusedAgg,+    mkFusedAgg,+    runFusedAggs,+    GatherAgg,+    mkGatherAgg,+    runGatherAggs,+) where++import Control.Exception (evaluate)+import Control.Monad (replicateM, 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.Aggregation.Kernel.Scatter (+    groupRangeBounds,+    groupSlices,+    overGroupsAcc,+    streamGroupCap,+ )+import DataFrame.Internal.Aggregation.Reduction (Reduction (..))+import DataFrame.Internal.Column (Column (..), fromUnboxedVector)+import DataFrame.Internal.Control.Concurrent (+    capabilities,+    forkJoin,+    parThreshold,+    parallelBounds_,+    shouldParallelize,+ )++-- | Whether to fan out at this row count.+shouldPar :: Int -> Bool+shouldPar = shouldParallelize parThreshold++-------------------------------------------------------------------------------+-- Fused multi-reduction streaming pass+-------------------------------------------------------------------------------++{- | One fused reduction of a multi-reduction streaming pass: how to allocate a+per-worker accumulator, fold a row range into it, merge another worker's+accumulator into it over a group range (callers merge in worker order), and+finalize the fully merged accumulator into the output column.+-}+data FusedAgg+    = forall s.+        FusedAgg+        (IO s)+        -- \^ allocate one worker's accumulator+        (s -> Int -> Int -> IO ())+        -- \^ accumulate rows [lo, hi) in ascending order+        (s -> s -> Int -> Int -> IO ())+        -- \^ merge the second accumulator into the first over groups [lo, hi)+        (s -> IO Column)++-- \^ finalize the merged accumulator++{- | Below this many groups a Double sum/mean does NOT join the fused streaming+pass: the per-expression kernel it would replace+('DataFrame.Internal.AggKernelDirect.sumDblDirect' under its @seqFloatGroups@+policy, value mirrored here) runs those as ONE sequential row-order pass that+is byte-identical to the gather kernels and the interpreter, and the vectorized+parity gate asserts exactly that. Above it the per-expression kernel already+merges chunk partials in worker order, so fusing changes nothing semantically+new. Int reductions and count/min/max are exact under any merge and always+fuse.+-}+fusedSeqFloatGroups :: Int+fusedSeqFloatGroups = 65536++{- | Build the fused-pass reduction for one @(reduction, column)@ pair, or+'Nothing' when the pair cannot stream (nullable/boxed columns; the+order-sensitive var/std/top2 reductions, which keep their per-expression+kernels; or a Double sum/mean below 'fusedSeqFloatGroups', which keeps its+byte-identical sequential pass). Sum/min/max/count/mean over Int are exact+under any chunk merge; the admitted Double sum and mean merge their per-worker+partials in worker order (deterministic at a fixed @-N@, chunk-major float+summation order).+-}+mkFusedAgg :: Int -> VU.Vector Int -> Reduction -> Column -> Maybe FusedAgg+mkFusedAgg nGroups rtg red col+    | nGroups <= 0 || nGroups > streamGroupCap = Nothing+    | otherwise = case col of+        UnboxedColumn Nothing (v :: VU.Vector a)+            | Just Refl <- testEquality (typeRep @a) (typeRep @Int) ->+                case red of+                    RCount -> Just (countFusedAgg nGroups rtg)+                    RSum -> Just (sumIntFusedAgg nGroups rtg v)+                    RMean -> Just (meanIntFusedAgg nGroups rtg v)+                    RMin -> Just (extremaIntFusedAgg True nGroups rtg v)+                    RMax -> Just (extremaIntFusedAgg False nGroups rtg v)+                    _ -> Nothing+            | Just Refl <- testEquality (typeRep @a) (typeRep @Double) ->+                case red of+                    RCount -> Just (countFusedAgg nGroups rtg)+                    RSum+                        | nGroups > fusedSeqFloatGroups ->+                            Just (sumDblFusedAgg nGroups rtg v)+                    RMean+                        | nGroups > fusedSeqFloatGroups ->+                            Just (meanDblFusedAgg nGroups rtg v)+                    RMin -> Just (extremaDblFusedAgg True nGroups rtg v)+                    RMax -> Just (extremaDblFusedAgg False nGroups rtg v)+                    -- Top-2 selection is an exact multiset selection (no float+                    -- arithmetic before finalize), so its per-worker merge is+                    -- byte-identical to the per-expression kernels at any -N.+                    RTop2Snd -> Just (top2SndDblFusedAgg nGroups rtg v)+                    _ -> Nothing+        _ -> Nothing++{- | Execute all fused reductions in ONE pass over @rowToGroup@ and the value+columns: the rows are split into one contiguous chunk per capability, each+worker walks its chunk in 4096-row blocks running every reduction's+accumulate-step on the block (the block's @rowToGroup@ slice stays in L1+across the k steps), and each reduction then merges its per-worker partials in+fixed worker order (parallel over group slices) and finalizes. Per-group update+order within each worker is ascending original row order, and chunk boundaries+are a fixed function of the row and capability counts, so the result is+deterministic at a fixed @-N@; on a single capability it is bit-identical to+running each unfused sequential scatter kernel separately.++Pure w.r.t. its immutable inputs (deterministic fan-out and merge order), so+the 'unsafePerformIO' is safe.+-}+runFusedAggs :: Int -> Int -> [FusedAgg] -> [Column]+runFusedAggs n nGroups aggs = unsafePerformIO $ do+    let !caps' = if shouldPar n then capabilities else 1+        !per = (max 1 n + caps' - 1) `div` caps'+    opened <- mapM (openFusedAgg caps' nGroups) aggs++    let stepsFor w = map (\(steps, _) -> steps !! w) opened+    _ <-+        forkJoin+            [ blockRun (stepsFor w) lo hi+            | w <- [0 .. caps' - 1]+            , let lo = min n (w * per)+            , let hi = min n (lo + per)+            ]+    mapM snd opened+{-# NOINLINE runFusedAggs #-}++{- | Open one fused reduction for @caps'@ workers: its per-worker step+functions (worker order) and the merge+finalize action.+-}+openFusedAgg :: Int -> Int -> FusedAgg -> IO ([Int -> Int -> IO ()], IO Column)+openFusedAgg caps' nGroups (FusedAgg new step mergeR fin) = do+    ss <- replicateM caps' new+    let finish = case ss of+            [] -> error "runFusedAggs: no workers"+            (s0 : rest) -> do+                _ <-+                    forkJoin+                        [ mapM_ (\s -> mergeR s0 s lo hi) rest+                        | (lo, hi) <- groupSlices nGroups+                        ]+                fin s0+    pure (map step ss, finish)++-- | Rows per fused block: the block's rowToGroup slice (32KB) stays in L1.+fusedBlock :: Int+fusedBlock = 4096++blockRun :: [Int -> Int -> IO ()] -> Int -> Int -> IO ()+blockRun steps lo0 hi = go lo0+  where+    go !lo+        | lo >= hi = pure ()+        | otherwise = do+            let !e = min hi (lo + fusedBlock)+            mapM_ (\s -> s lo e) steps+            go e++-- Individual fused reductions. Each step loop is monomorphic.++countFusedAgg :: Int -> VU.Vector Int -> FusedAgg+countFusedAgg nGroups rtg =+    FusedAgg+        (VUM.replicate nGroups (0 :: Int))+        (countStepK rtg)+        addIntRange+        (fmap fromUnboxedVector . VU.unsafeFreeze)++countStepK :: VU.Vector Int -> VUM.IOVector Int -> Int -> Int -> IO ()+countStepK rtg acc lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !k = VU.unsafeIndex rtg i+            c <- VUM.unsafeRead acc k+            VUM.unsafeWrite acc k (c + 1)+            go (i + 1)++sumIntFusedAgg :: Int -> VU.Vector Int -> VU.Vector Int -> FusedAgg+sumIntFusedAgg nGroups rtg v =+    FusedAgg+        (VUM.replicate nGroups (0 :: Int))+        (sumStepInt rtg v)+        addIntRange+        (fmap fromUnboxedVector . VU.unsafeFreeze)++sumStepInt ::+    VU.Vector Int -> VU.Vector Int -> VUM.IOVector Int -> Int -> Int -> IO ()+sumStepInt rtg v acc lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !k = VU.unsafeIndex rtg i+            c <- VUM.unsafeRead acc k+            VUM.unsafeWrite acc k (c + VU.unsafeIndex v i)+            go (i + 1)++sumDblFusedAgg :: Int -> VU.Vector Int -> VU.Vector Double -> FusedAgg+sumDblFusedAgg nGroups rtg v =+    FusedAgg+        (VUM.replicate nGroups (0 :: Double))+        (sumStepDbl rtg v)+        addDblRange+        (fmap fromUnboxedVector . VU.unsafeFreeze)++sumStepDbl ::+    VU.Vector Int -> VU.Vector Double -> VUM.IOVector Double -> Int -> Int -> IO ()+sumStepDbl rtg v acc lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !k = VU.unsafeIndex rtg i+            c <- VUM.unsafeRead acc k+            VUM.unsafeWrite acc k (c + VU.unsafeIndex v i)+            go (i + 1)++{- | The mean aggs hold their (sum, count) pair INTERLEAVED in one array —+slots @2g@/@2g+1@ share a cache line (pairs are 16-byte aligned, so they never+straddle one), halving the accumulator misses of the random per-row update+against two separate arrays (measured ~13% off a fused sum+mean pass at 1e6+groups / 1e8 rows on -N16). The count is exact in both layouts (an integer, or+integer-valued Double additions well below 2^53), so sums, merges and the+finalize divide are bit-identical to the two-array layout.+-}+meanIntFusedAgg :: Int -> VU.Vector Int -> VU.Vector Int -> FusedAgg+meanIntFusedAgg nGroups rtg v =+    FusedAgg+        (VUM.replicate (2 * nGroups) (0 :: Int))+        (meanStepInt rtg v)+        (\a b lo hi -> addIntRange a b (2 * lo) (2 * hi))+        ( \s -> do+            sv <- VU.unsafeFreeze s+            pure+                ( fromUnboxedVector+                    ( VU.generate+                        nGroups+                        ( \g ->+                            let !sx = VU.unsafeIndex sv (2 * g)+                                !cx = VU.unsafeIndex sv (2 * g + 1)+                             in if cx == 0+                                    then 0 / 0+                                    else fromIntegral sx / fromIntegral cx :: Double+                        )+                    )+                )+        )++meanStepInt ::+    VU.Vector Int ->+    VU.Vector Int ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    IO ()+meanStepInt rtg v s lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !k2 = 2 * VU.unsafeIndex rtg i+            sv <- VUM.unsafeRead s k2+            VUM.unsafeWrite s k2 (sv + VU.unsafeIndex v i)+            cv <- VUM.unsafeRead s (k2 + 1)+            VUM.unsafeWrite s (k2 + 1) (cv + 1)+            go (i + 1)++-- | See 'meanIntFusedAgg' for the interleaved accumulator layout.+meanDblFusedAgg :: Int -> VU.Vector Int -> VU.Vector Double -> FusedAgg+meanDblFusedAgg nGroups rtg v =+    FusedAgg+        (VUM.replicate (2 * nGroups) (0 :: Double))+        (meanStepDbl rtg v)+        (\a b lo hi -> addDblRange a b (2 * lo) (2 * hi))+        ( \s -> do+            sv <- VU.unsafeFreeze s+            pure+                ( fromUnboxedVector+                    ( VU.generate+                        nGroups+                        ( \g ->+                            let !sx = VU.unsafeIndex sv (2 * g)+                                !cx = VU.unsafeIndex sv (2 * g + 1)+                             in if cx == 0 then 0 / 0 else sx / cx+                        )+                    )+                )+        )++meanStepDbl ::+    VU.Vector Int ->+    VU.Vector Double ->+    VUM.IOVector Double ->+    Int ->+    Int ->+    IO ()+meanStepDbl rtg v s lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !k2 = 2 * VU.unsafeIndex rtg i+            sv <- VUM.unsafeRead s k2+            VUM.unsafeWrite s k2 (sv + VU.unsafeIndex v i)+            cv <- VUM.unsafeRead s (k2 + 1)+            VUM.unsafeWrite s (k2 + 1) (cv + 1)+            go (i + 1)++extremaIntFusedAgg :: Bool -> Int -> VU.Vector Int -> VU.Vector Int -> FusedAgg+extremaIntFusedAgg isMin nGroups rtg v =+    FusedAgg+        (VUM.replicate nGroups (if isMin then maxBound else minBound :: Int))+        (extremaStepInt isMin rtg v)+        (combineIntRange isMin)+        (fmap fromUnboxedVector . VU.unsafeFreeze)++extremaStepInt ::+    Bool ->+    VU.Vector Int ->+    VU.Vector Int ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    IO ()+extremaStepInt isMin rtg v acc lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !k = VU.unsafeIndex rtg i+                !x = VU.unsafeIndex v i+            c <- VUM.unsafeRead acc k+            VUM.unsafeWrite acc k (if isMin then min c x else max c x)+            go (i + 1)++extremaDblFusedAgg ::+    Bool -> Int -> VU.Vector Int -> VU.Vector Double -> FusedAgg+extremaDblFusedAgg isMin nGroups rtg v =+    FusedAgg+        (VUM.replicate nGroups (if isMin then 1 / 0 else negate (1 / 0) :: Double))+        (extremaStepDbl isMin rtg v)+        (combineDblRange isMin)+        (fmap fromUnboxedVector . VU.unsafeFreeze)++extremaStepDbl ::+    Bool ->+    VU.Vector Int ->+    VU.Vector Double ->+    VUM.IOVector Double ->+    Int ->+    Int ->+    IO ()+extremaStepDbl isMin rtg v acc lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !k = VU.unsafeIndex rtg i+                !x = VU.unsafeIndex v i+            c <- VUM.unsafeRead acc k+            VUM.unsafeWrite acc k (if isMin then min c x else max c x)+            go (i + 1)++{- | Fused second-largest over a Double column. The per-group+(largest, second-largest) pair is INTERLEAVED at slots @2g@/@2g+1@ (one cache+line per group, as 'meanDblFusedAgg'); the update is the same top-2 selection+as every other top2 kernel, the merge keeps the top two of the four candidates+per group (exact — no float arithmetic), and the finalize returns the second+max, NaN for a group of size < 2 (the @-inf@ seed; see+'DataFrame.Internal.AggKernel.top2SndScatter').+-}+top2SndDblFusedAgg :: Int -> VU.Vector Int -> VU.Vector Double -> FusedAgg+top2SndDblFusedAgg nGroups rtg v =+    FusedAgg+        (VUM.replicate (2 * nGroups) (negate (1 / 0) :: Double))+        (top2SndStepDbl rtg v)+        mergeTop2Range+        ( \s -> do+            sv <- VU.unsafeFreeze s+            pure+                ( fromUnboxedVector+                    ( VU.generate+                        nGroups+                        ( \g ->+                            let !a2 = VU.unsafeIndex sv (2 * g + 1)+                             in if isInfinite a2 then 0 / 0 else a2+                        )+                    )+                )+        )++top2SndStepDbl ::+    VU.Vector Int ->+    VU.Vector Double ->+    VUM.IOVector Double ->+    Int ->+    Int ->+    IO ()+top2SndStepDbl rtg v s lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !k2 = 2 * VU.unsafeIndex rtg i+                !x = VU.unsafeIndex v i+            a1 <- VUM.unsafeRead s k2+            if x > a1+                then do+                    VUM.unsafeWrite s k2 x+                    VUM.unsafeWrite s (k2 + 1) a1+                else do+                    a2 <- VUM.unsafeRead s (k2 + 1)+                    when (x > a2) $ VUM.unsafeWrite s (k2 + 1) x+            go (i + 1)++-- | Top two of the four candidates per group (pairs already ordered m1 >= m2).+mergeTop2Range ::+    VUM.IOVector Double -> VUM.IOVector Double -> Int -> Int -> IO ()+mergeTop2Range a b lo hi = go lo+  where+    go !g+        | g >= hi = pure ()+        | otherwise = do+            let !g2 = 2 * g+            a1 <- VUM.unsafeRead a g2+            a2 <- VUM.unsafeRead a (g2 + 1)+            b1 <- VUM.unsafeRead b g2+            b2 <- VUM.unsafeRead b (g2 + 1)+            if b1 > a1+                then do+                    VUM.unsafeWrite a g2 b1+                    VUM.unsafeWrite a (g2 + 1) (max a1 b2)+                else VUM.unsafeWrite a (g2 + 1) (max a2 b1)+            go (g + 1)++addIntRange :: VUM.IOVector Int -> VUM.IOVector Int -> Int -> Int -> IO ()+addIntRange a b lo hi = go lo+  where+    go !g+        | g >= hi = pure ()+        | otherwise = do+            x <- VUM.unsafeRead a g+            y <- VUM.unsafeRead b g+            VUM.unsafeWrite a g (x + y)+            go (g + 1)++addDblRange :: VUM.IOVector Double -> VUM.IOVector Double -> Int -> Int -> IO ()+addDblRange a b lo hi = go lo+  where+    go !g+        | g >= hi = pure ()+        | otherwise = do+            x <- VUM.unsafeRead a g+            y <- VUM.unsafeRead b g+            VUM.unsafeWrite a g (x + y)+            go (g + 1)++combineIntRange ::+    Bool -> VUM.IOVector Int -> VUM.IOVector Int -> Int -> Int -> IO ()+combineIntRange isMin a b lo hi = go lo+  where+    go !g+        | g >= hi = pure ()+        | otherwise = do+            x <- VUM.unsafeRead a g+            y <- VUM.unsafeRead b g+            VUM.unsafeWrite a g (if isMin then min x y else max x y)+            go (g + 1)++combineDblRange ::+    Bool -> VUM.IOVector Double -> VUM.IOVector Double -> Int -> Int -> IO ()+combineDblRange isMin a b lo hi = go lo+  where+    go !g+        | g >= hi = pure ()+        | otherwise = do+            x <- VUM.unsafeRead a g+            y <- VUM.unsafeRead b g+            VUM.unsafeWrite a g (if isMin then min x y else max x y)+            go (g + 1)++-------------------------------------------------------------------------------+-- Fused multi-reduction gather pass (nGroups above 'streamGroupCap')+-------------------------------------------------------------------------------++{- | One fused reduction of a multi-reduction GATHER pass: allocate the output+array, fold a contiguous group range (each group's rows via the shared+@valueIndices@ slice, accumulator in registers, one write per group), finalize.+Group ranges are disjoint across workers, so there is no merge and every+reduction reproduces the exact per-group fold order and formula of its unfused+gather kernel ('reduceParTyped') — results are bit-identical to running the+kernels separately, at any @-N@.+-}+data GatherAgg+    = forall s.+        GatherAgg+        (IO s)+        (s -> Int -> Int -> IO ())+        (s -> IO Column)++{- | Build the fused gather reduction for one @(reduction, column)@ pair, or+'Nothing' when the pair cannot fuse (nullable/boxed columns, or var/std/top2,+which keep their per-expression gather kernels). @vis@/@offs@ are captured+lazily: nothing is forced until the pass actually runs.+-}+mkGatherAgg ::+    Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    Reduction ->+    Column ->+    Maybe GatherAgg+mkGatherAgg nGroups vis offs red col = case col of+    UnboxedColumn Nothing (v :: VU.Vector a)+        | Just Refl <- testEquality (typeRep @a) (typeRep @Int) ->+            case red of+                RCount -> Just countGather+                RSum -> Just (outGather (gatherSumInt vis offs v))+                RMean -> Just (outGatherD (gatherMeanInt vis offs v))+                RMin -> Just (outGather (gatherExtremaInt True vis offs v))+                RMax -> Just (outGather (gatherExtremaInt False vis offs v))+                _ -> Nothing+        | Just Refl <- testEquality (typeRep @a) (typeRep @Double) ->+            case red of+                RCount -> Just countGather+                RSum -> Just (outGatherD (gatherSumDbl vis offs v))+                RMean -> Just (outGatherD (gatherMeanDbl vis offs v))+                RMin -> Just (outGatherD (gatherExtremaDbl True vis offs v))+                RMax -> Just (outGatherD (gatherExtremaDbl False vis offs v))+                _ -> Nothing+    _ -> Nothing+  where+    countGather =+        GatherAgg+            (VUM.new nGroups)+            (gatherCount offs)+            (fmap fromUnboxedVector . VU.unsafeFreeze)+    outGather step =+        GatherAgg+            (VUM.new nGroups :: IO (VUM.IOVector Int))+            step+            (fmap fromUnboxedVector . VU.unsafeFreeze)+    outGatherD step =+        GatherAgg+            (VUM.new nGroups :: IO (VUM.IOVector Double))+            step+            (fmap fromUnboxedVector . VU.unsafeFreeze)++{- | Number of groups each fused-gather block hands to every reduction before+moving on: the block's @valueIndices@ slice stays hot in cache across the k+per-reduction loops.+-}+gatherBlock :: Int+gatherBlock = 32++{- | Execute all fused gather reductions in one traversal: workers own disjoint+contiguous group ranges (row-balanced, same policy as every gather kernel), and+walk them in 'gatherBlock'-group blocks running each reduction's fold on the+block. Deterministic and bit-identical to the unfused kernels (see+'GatherAgg'). Forces @valueIndices@ once, before the fan-out.++Pure w.r.t. its immutable inputs, so the 'unsafePerformIO' is safe.+-}+runGatherAggs ::+    VU.Vector Int -> VU.Vector Int -> Int -> [GatherAgg] -> [Column]+runGatherAggs vis offs nGroups aggs = unsafePerformIO $ do+    _ <- evaluate (VU.length vis)+    let !caps = capabilities+        !bounds = groupRangeBounds offs nGroups caps+    opened <-+        mapM (\(GatherAgg new step fin) -> do s <- new; pure (step s, fin s)) aggs+    parallelBounds_ caps bounds $ \gs ge ->+        let go !g+                | g >= ge = pure ()+                | otherwise = do+                    let !e = min ge (g + gatherBlock)+                    mapM_ (\(st, _) -> st g e) opened+                    go e+         in go gs+    mapM snd opened+{-# NOINLINE runGatherAggs #-}++-- Monomorphic per-reduction gather folds; each replicates the exact per-group+-- recurrence of its unfused kernel above ('countPar'/'sumPar'/'extremaPar'/+-- 'meanPar'), so fused results are bit-identical.++gatherCount :: VU.Vector Int -> VUM.IOVector Int -> Int -> Int -> IO ()+gatherCount offs out gs ge = go gs+  where+    go !g+        | g >= ge = pure ()+        | otherwise = do+            let !c = VU.unsafeIndex offs (g + 1) - VU.unsafeIndex offs g+            VUM.unsafeWrite out g c+            go (g + 1)++gatherSumInt ::+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    IO ()+gatherSumInt vis offs v out gs ge =+    overGroupsAcc vis offs gs ge 0 (\acc row -> acc + VU.unsafeIndex v row) $+        VUM.unsafeWrite out++gatherSumDbl ::+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Double ->+    VUM.IOVector Double ->+    Int ->+    Int ->+    IO ()+gatherSumDbl vis offs v out gs ge =+    overGroupsAcc vis offs gs ge 0 (\acc row -> acc + VU.unsafeIndex v row) $+        VUM.unsafeWrite out++gatherExtremaInt ::+    Bool ->+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    IO ()+gatherExtremaInt isMin vis offs v out gs ge =+    overGroupsAcc+        vis+        offs+        gs+        ge+        (if isMin then maxBound else minBound)+        ( \acc row ->+            let !x = VU.unsafeIndex v row+             in if isMin then min acc x else max acc x+        )+        (VUM.unsafeWrite out)++gatherExtremaDbl ::+    Bool ->+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Double ->+    VUM.IOVector Double ->+    Int ->+    Int ->+    IO ()+gatherExtremaDbl isMin vis offs v out gs ge =+    overGroupsAcc+        vis+        offs+        gs+        ge+        (if isMin then 1 / 0 else negate (1 / 0))+        ( \acc row ->+            let !x = VU.unsafeIndex v row+             in if isMin then min acc x else max acc x+        )+        (VUM.unsafeWrite out)++-- | Exact replica of 'meanPar''s per-group loop (Int element type).+gatherMeanInt ::+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    VUM.IOVector Double ->+    Int ->+    Int ->+    IO ()+gatherMeanInt vis offs v out gs ge = grp gs+  where+    grp !g+        | g >= ge = pure ()+        | otherwise = do+            let !e = VU.unsafeIndex offs (g + 1)+                inner !pos !acc+                    | pos >= e = acc+                    | otherwise =+                        inner+                            (pos + 1)+                            (acc + fromIntegral (VU.unsafeIndex v (VU.unsafeIndex vis pos)))+                !s0 = VU.unsafeIndex offs g+                !total = inner s0 0+                !c = e - s0+            VUM.unsafeWrite out g (if c == 0 then 0 / 0 else total / fromIntegral c)+            grp (g + 1)++-- | Exact replica of 'meanPar''s per-group loop (Double element type).+gatherMeanDbl ::+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Double ->+    VUM.IOVector Double ->+    Int ->+    Int ->+    IO ()+gatherMeanDbl vis offs v out gs ge = grp gs+  where+    grp !g+        | g >= ge = pure ()+        | otherwise = do+            let !e = VU.unsafeIndex offs (g + 1)+                inner !pos !acc+                    | pos >= e = acc+                    | otherwise =+                        inner+                            (pos + 1)+                            (acc + VU.unsafeIndex v (VU.unsafeIndex vis pos))+                !s0 = VU.unsafeIndex offs g+                !total = inner s0 0+                !c = e - s0+            VUM.unsafeWrite out g (if c == 0 then 0 / 0 else total / fromIntegral c)+            grp (g + 1)
+ src-internal/DataFrame/Internal/Aggregation/Kernel/Moments.hs view
@@ -0,0 +1,389 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | The fused two-column moment kernel: one pass over @x@ and @y@ producing the+six sufficient statistics @(n, Sx, Sy, Sxx, Syy, Sxy)@ per group, from which the+whole correlation\/regression family (mean, variance, covariance, correlation,+OLS slope) is algebra requiring no further look at the rows.++The sequential and parallel passes live together because they must agree on+floating-point accumulation order. Moments are additive, so a row-range split+with a merge would be correct in exact arithmetic — but float addition is not+associative, so 'momentScatterPar' partitions by GROUP range instead: every+group accumulates start-to-finish inside a single worker, in the same order as+'momentScatter'. That is what makes the two byte-identical at any @-N@.+-}+module DataFrame.Internal.Aggregation.Kernel.Moments (+    Moments (..),+    momentScatter,+    momentScatterPar,+    momentStreamPar,+) where++import Control.Monad.ST (runST)+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.Aggregation.Kernel.Scatter (+    groupRangeBounds,+    groupSlices,+    rtgFromVis,+    streamGroupCap,+ )+import DataFrame.Internal.Aggregation.Reduction (cleanDoubleVector)+import DataFrame.Internal.Column (Column (..), fromUnboxedVector)+import DataFrame.Internal.Control.Concurrent (+    capabilities,+    forkJoin,+    parThreshold,+    parallelBounds_,+    shouldParallelize,+ )++-- | Whether to fan out at this row count.+shouldPar :: Int -> Bool+shouldPar = shouldParallelize parThreshold++{- | 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 <- cleanDoubleVector colX+    ys <- cleanDoubleVector 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++{- | 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 <- cleanDoubleVector colX+        ys <- cleanDoubleVector 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)+    parallelBounds_ caps bounds $ \gs ge ->+        -- The six running sums carried in registers per group, written once.+        let grp !g+                | g >= ge = pure ()+                | otherwise = do+                    let !e = VU.unsafeIndex offs (g + 1)+                        inner !pos !ax !ay !axx !ayy !axy+                            | pos >= e = do+                                VUM.unsafeWrite sx g ax+                                VUM.unsafeWrite sy g ay+                                VUM.unsafeWrite sxx g axx+                                VUM.unsafeWrite syy g ayy+                                VUM.unsafeWrite sxy g axy+                            | otherwise =+                                let !row = VU.unsafeIndex vis pos+                                    !x = VU.unsafeIndex xs row+                                    !y = VU.unsafeIndex ys row+                                 in inner+                                        (pos + 1)+                                        (ax + x)+                                        (ay + y)+                                        (axx + x * x)+                                        (ayy + y * y)+                                        (axy + x * y)+                        !s0 = VU.unsafeIndex offs g+                    VUM.unsafeWrite cnt g (e - s0)+                    inner s0 0 0 0 0 0+                    grp (g + 1)+         in grp gs+    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)++-------------------------------------------------------------------------------+-- Streaming fused two-column moments (Q9)+-------------------------------------------------------------------------------++{- | Streaming counterpart of 'momentScatterPar': one fused pass over+@rowToGroup@ and the two TYPED value columns (Int values convert to Double+in-register — bit-identical to the @VU.map fromIntegral@ materialization it+replaces, with no 800MB intermediate column and no sequential conversion pass).+Each worker accumulates the six per-group sums over its contiguous row chunk in+original row order; partials merge in fixed worker order (counts exactly, the+five Double sums in chunk-major float order — deterministic at a fixed @-N@,+but a different summation order than the per-group gather kernel).+'Nothing' above 'streamGroupCap' or unless both columns are clean unboxed+Int/Double; the caller then keeps the gather path.+-}+momentStreamPar :: VU.Vector Int -> Int -> Column -> Column -> Maybe Moments+momentStreamPar rtg nGroups colX colY+    | nGroups <= 0 || nGroups > streamGroupCap = Nothing+    | otherwise = case (colX, colY) of+        ( UnboxedColumn Nothing (vx :: VU.Vector x)+            , UnboxedColumn Nothing (vy :: VU.Vector y)+            )+                | Just Refl <- testEquality (typeRep @x) (typeRep @Int)+                , Just Refl <- testEquality (typeRep @y) (typeRep @Int) ->+                    Just (momentStreamII rtg nGroups vx vy)+                | Just Refl <- testEquality (typeRep @x) (typeRep @Int)+                , Just Refl <- testEquality (typeRep @y) (typeRep @Double) ->+                    Just (momentStreamID rtg nGroups vx vy)+                | Just Refl <- testEquality (typeRep @x) (typeRep @Double)+                , Just Refl <- testEquality (typeRep @y) (typeRep @Int) ->+                    Just (momentStreamDI rtg nGroups vx vy)+                | Just Refl <- testEquality (typeRep @x) (typeRep @Double)+                , Just Refl <- testEquality (typeRep @y) (typeRep @Double) ->+                    Just (momentStreamDD rtg nGroups vx vy)+        _ -> Nothing+{-# NOINLINE momentStreamPar #-}++{- | Monomorphic entry points (see 'reduceParInt' for why the 'testEquality'+dispatch needs them).+-}+momentStreamII ::+    VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int -> Moments+momentStreamII = momentStreamTyped+{-# NOINLINE momentStreamII #-}++momentStreamID ::+    VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Double -> Moments+momentStreamID = momentStreamTyped+{-# NOINLINE momentStreamID #-}++momentStreamDI ::+    VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Int -> Moments+momentStreamDI = momentStreamTyped+{-# NOINLINE momentStreamDI #-}++momentStreamDD ::+    VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double -> Moments+momentStreamDD = momentStreamTyped+{-# NOINLINE momentStreamDD #-}++-- | The six per-group running sums of one worker chunk.+data MomentAcc = MomentAcc+    { maCnt :: !(VUM.IOVector Int)+    , maSx :: !(VUM.IOVector Double)+    , maSy :: !(VUM.IOVector Double)+    , maSxx :: !(VUM.IOVector Double)+    , maSyy :: !(VUM.IOVector Double)+    , maSxy :: !(VUM.IOVector Double)+    }++newMomentAcc :: Int -> IO MomentAcc+newMomentAcc nGroups =+    MomentAcc+        <$> VUM.replicate nGroups 0+        <*> VUM.replicate nGroups 0+        <*> VUM.replicate nGroups 0+        <*> VUM.replicate nGroups 0+        <*> VUM.replicate nGroups 0+        <*> VUM.replicate nGroups 0++momentStreamTyped ::+    forall a b.+    (VU.Unbox a, VU.Unbox b, Real a, Real b) =>+    VU.Vector Int ->+    Int ->+    VU.Vector a ->+    VU.Vector b ->+    Moments+{- The SPECIALIZE pragmas matter for the same reason as 'reduceParTyped': the+per-element @realToFrac@ must rewrite to @int2Double@/@id@ at a concrete type+or it goes through 'Rational' at runtime. -}+{-# SPECIALIZE momentStreamTyped ::+    VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int -> Moments+    #-}+{-# SPECIALIZE momentStreamTyped ::+    VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Double -> Moments+    #-}+{-# SPECIALIZE momentStreamTyped ::+    VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Int -> Moments+    #-}+{-# SPECIALIZE momentStreamTyped ::+    VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double -> Moments+    #-}+momentStreamTyped rtg nGroups vx vy = unsafePerformIO $ do+    let !n = VU.length rtg+        !caps' = if shouldPar n then capabilities else 1+        !per = (max 1 n + caps' - 1) `div` caps'+    parts <-+        forkJoin+            [ momentStreamChunk rtg nGroups vx vy lo hi+            | w <- [0 .. caps' - 1]+            , let lo = min n (w * per)+            , let hi = min n (lo + per)+            ]+    case parts of+        [] -> error "momentStreamTyped: no partials"+        (p0 : rest) -> do+            _ <-+                forkJoin+                    [ mapM_ (\p -> mergeMomentRange p0 p lo hi) rest+                    | (lo, hi) <- groupSlices nGroups+                    ]+            freezeMoments p0+{-# INLINEABLE momentStreamTyped #-}++momentStreamChunk ::+    (VU.Unbox a, VU.Unbox b, Real a, Real b) =>+    VU.Vector Int ->+    Int ->+    VU.Vector a ->+    VU.Vector b ->+    Int ->+    Int ->+    IO MomentAcc+momentStreamChunk rtg nGroups vx vy lo hi = do+    acc@(MomentAcc cnt sx sy sxx syy sxy) <- newMomentAcc nGroups+    let go !i+            | i >= hi = pure ()+            | otherwise = do+                let !k = VU.unsafeIndex rtg i+                    !x = realToFrac (VU.unsafeIndex vx i) :: Double+                    !y = realToFrac (VU.unsafeIndex vy i) :: Double+                c <- VUM.unsafeRead cnt k+                VUM.unsafeWrite cnt k (c + 1)+                ax <- VUM.unsafeRead sx k+                VUM.unsafeWrite sx k (ax + x)+                ay <- VUM.unsafeRead sy k+                VUM.unsafeWrite sy k (ay + y)+                axx <- VUM.unsafeRead sxx k+                VUM.unsafeWrite sxx k (axx + x * x)+                ayy <- VUM.unsafeRead syy k+                VUM.unsafeWrite syy k (ayy + y * y)+                axy <- VUM.unsafeRead sxy k+                VUM.unsafeWrite sxy k (axy + x * y)+                go (i + 1)+    go lo+    pure acc+{-# INLINE momentStreamChunk #-}++mergeMomentRange :: MomentAcc -> MomentAcc -> Int -> Int -> IO ()+mergeMomentRange a b lo hi = go lo+  where+    go !g+        | g >= hi = pure ()+        | otherwise = do+            addI (maCnt a) (maCnt b) g+            addD (maSx a) (maSx b) g+            addD (maSy a) (maSy b) g+            addD (maSxx a) (maSxx b) g+            addD (maSyy a) (maSyy b) g+            addD (maSxy a) (maSxy b) g+            go (g + 1)+    addI p q g = do+        x <- VUM.unsafeRead p g+        y <- VUM.unsafeRead q g+        VUM.unsafeWrite p g (x + y)+    addD p q g = do+        x <- VUM.unsafeRead p g+        y <- VUM.unsafeRead q g+        VUM.unsafeWrite p g (x + y)++freezeMoments :: MomentAcc -> IO Moments+freezeMoments (MomentAcc cnt sx sy sxx syy sxy) =+    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)
+ src-internal/DataFrame/Internal/Aggregation/Kernel/Scatter.hs view
@@ -0,0 +1,818 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | The group-range scatter-accumulate reduction kernel: reduces a value column+over the grouped layout @(valueIndices, offsets)@.++Sequential and parallel are the same algorithm at different row counts, so they+live together. 'scatterReducePar' cuts the GROUP axis into @caps@ ranges of+roughly equal row count and lets workers write disjoint slots of one shared+output — no per-worker accumulator, no merge — which keeps each group's+accumulation order identical to 'scatterReduce' and the results byte-identical+at any @-N@. Below 'parThreshold' it delegates to 'scatterReduce' directly.++Contrast "DataFrame.Internal.Aggregation.Kernel.Dense", which scatters off+@rowToGroup@ with no gather but needs a small dense group domain.+-}+module DataFrame.Internal.Aggregation.Kernel.Scatter (+    scatterReduce,+    scatterReducePar,+    maxMinusMinScatterPar,+    top2SndScatter,++    -- * Group-range helpers+    -- $shared+    groupRangeBounds,+    rtgFromVis,+    overGroupsAcc,+    groupSlices,+    streamGroupCap,+) where++import Control.Monad (when)+import Control.Monad.ST (ST, runST)+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.Aggregation.Reduction (Reduction (..))+import DataFrame.Internal.Column (+    Column (..),+    Columnable,+    fromUnboxedVector,+    materializePacked,+ )+import DataFrame.Internal.Control.Concurrent (+    capabilities,+    chunksFor,+    parThreshold,+    parallelBounds_,+    shouldParallelize,+ )++{- $shared+Also used by "DataFrame.Internal.Aggregation.Kernel.Moments", which partitions+the group axis the same way.+-}++-- | Whether to fan out at this row count.+shouldPar :: Int -> Bool+shouldPar = shouldParallelize parThreshold++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)+    RTop2Snd -> fromUnboxedVector (top2SndScatter 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 #-}++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++scatterReducePar ::+    Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> Maybe Column+scatterReducePar red vis offs nGroups col+    | not (shouldParallelize parThreshold (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 (reduceParInt red vis offs nGroups v)+                Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+                    Just Refl -> Just (reduceParDouble red vis offs nGroups v)+                    Nothing -> Nothing+        p@(PackedText _ _) -> scatterReducePar red vis offs nGroups (materializePacked p)+        _ -> Nothing+{-# NOINLINE scatterReducePar #-}++{- | Monomorphic entry points: the 'testEquality' dispatch above only yields an+unsafe coercion, so a direct call to the polymorphic 'reduceParTyped' there+would stay at the abstract element type and never meet its SPECIALIZE rules;+calling through these fixed-type wrappers (the coercion lands on the argument)+does.+-}+reduceParInt ::+    Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Int -> Column+reduceParInt red vis offs nGroups v = reduceParTyped red vis offs nGroups v intIdent++reduceParDouble ::+    Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Double -> Column+reduceParDouble red vis offs nGroups v = reduceParTyped red vis offs nGroups v dblIdent++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++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+{- The SPECIALIZE pragmas matter: without them the @realToFrac@ in the+mean/var/top2 kernels survives to runtime as a dictionary call through+'Rational' (the Double->Double/Int->Double rewrite rules only fire once the+type is concrete), costing ~4x on the whole pass. -}+{-# SPECIALIZE reduceParTyped ::+    Reduction ->+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector Int ->+    Idents Int ->+    Column+    #-}+{-# SPECIALIZE reduceParTyped ::+    Reduction ->+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector Double ->+    Idents Double ->+    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))+            RTop2Snd ->+                fromUnboxedVector (unsafePerformIO (top2SndPar vis offs nGroups v caps bounds))+{-# INLINEABLE reduceParTyped #-}++{- | For each group in @[gs, ge)@, fold the group's rows (in @valueIndices@+order, i.e. ascending original-row order) into an accumulator held in+registers, then hand the final accumulator to @done@ exactly once. Keeping the+running state out of memory leaves one write per group instead of a+read-modify-write per row; the per-group fold order is unchanged, so results+stay byte-identical to the row-wise variant.+-}+overGroupsAcc ::+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    Int ->+    acc ->+    (acc -> Int -> acc) ->+    (Int -> acc -> IO ()) ->+    IO ()+overGroupsAcc vis offs gs ge seed step done = grp gs+  where+    grp !g+        | g >= ge = pure ()+        | otherwise = do+            let !e = VU.unsafeIndex offs (g + 1)+                inner !pos !acc+                    | pos >= e = pure acc+                    | otherwise = inner (pos + 1) (step acc (VU.unsafeIndex vis pos))+            acc <- inner (VU.unsafeIndex offs g) seed+            done g acc+            grp (g + 1)+{-# INLINE overGroupsAcc #-}++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)+    parallelBounds_ caps bounds $ \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+    parallelBounds_ caps bounds $ \gs ge ->+        overGroupsAcc vis offs gs ge 0 (\acc row -> acc + VU.unsafeIndex v row) $+            VUM.unsafeWrite out+    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+    parallelBounds_ caps bounds $ \gs ge ->+        overGroupsAcc+            vis+            offs+            gs+            ge+            seed+            (\acc row -> combine acc (VU.unsafeIndex v row))+            $ VUM.unsafeWrite out+    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+    out <- VUM.replicate nGroups (0 :: Double)+    parallelBounds_ caps bounds $ \gs ge ->+        let grp !g+                | g >= ge = pure ()+                | otherwise = do+                    let !e = VU.unsafeIndex offs (g + 1)+                        inner !pos !acc+                            | pos >= e = acc+                            | otherwise =+                                inner+                                    (pos + 1)+                                    (acc + realToFrac (VU.unsafeIndex v (VU.unsafeIndex vis pos)))+                        !s0 = VU.unsafeIndex offs g+                        !total = inner s0 0+                        !c = e - s0+                    VUM.unsafeWrite out g (if c == 0 then 0 / 0 else total / fromIntegral c)+                    grp (g + 1)+         in grp gs+    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+    out <- VUM.replicate nGroups (0 :: Double)+    parallelBounds_ caps bounds $ \gs ge ->+        -- Per-group Welford state (count, mean, M2) carried in registers; the+        -- update order per group is the same ascending row order as before.+        let grp !g+                | g >= ge = pure ()+                | otherwise = do+                    let !e = VU.unsafeIndex offs (g + 1)+                        inner !pos !c !mu !mm+                            | pos >= e =+                                let var = if c < 2 then 0 else mm / fromIntegral (c - 1)+                                 in if takeSqrt then sqrt var else var+                            | otherwise =+                                let !x = realToFrac (VU.unsafeIndex v (VU.unsafeIndex vis pos))+                                    !c' = c + 1+                                    !delta = x - mu+                                    !mu' = mu + delta / fromIntegral c'+                                    !mm' = mm + delta * (x - mu')+                                 in inner (pos + 1) c' mu' mm'+                        !res = inner (VU.unsafeIndex offs g) (0 :: Int) 0 0+                    VUM.unsafeWrite out g res+                    grp (g + 1)+         in grp gs+    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+    out <- VUM.replicate nGroups (0 :: Double)+    parallelBounds_ caps bounds $ \gs ge ->+        -- The (largest, second-largest) pair carried in registers per group.+        let grp !g+                | g >= ge = pure ()+                | otherwise = do+                    let !e = VU.unsafeIndex offs (g + 1)+                        inner !pos !a1 !a2+                            | pos >= e =+                                (if isInfinite a1 then 0 else a1)+                                    + (if isInfinite a2 then 0 else a2)+                            | otherwise =+                                let !x = realToFrac (VU.unsafeIndex v (VU.unsafeIndex vis pos))+                                 in if x > a1+                                        then inner (pos + 1) x a1+                                        else inner (pos + 1) a1 (max a2 x)+                        !res = inner (VU.unsafeIndex offs g) ninf ninf+                    VUM.unsafeWrite out g res+                    grp (g + 1)+         in grp gs+    VU.unsafeFreeze out+{-# INLINE top2Par #-}++{- | Second-largest value per group: the same (largest, second-largest)+register pair as 'top2Par', finalized to the second max alone. Size-1 groups+finalize the @-inf@ seed to NaN (documented; see+'DataFrame.Internal.AggKernel.top2SndScatter').+-}+top2SndPar ::+    (VU.Unbox a, Real a) =>+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector a ->+    Int ->+    VU.Vector Int ->+    IO (VU.Vector Double)+top2SndPar vis offs nGroups v caps bounds = do+    let ninf = negate (1 / 0) :: Double+    out <- VUM.replicate nGroups (0 :: Double)+    parallelBounds_ caps bounds $ \gs ge ->+        -- The (largest, second-largest) pair carried in registers per group.+        let grp !g+                | g >= ge = pure ()+                | otherwise = do+                    let !e = VU.unsafeIndex offs (g + 1)+                        inner !pos !a1 !a2+                            | pos >= e = if isInfinite a2 then 0 / 0 else a2+                            | otherwise =+                                let !x = realToFrac (VU.unsafeIndex v (VU.unsafeIndex vis pos))+                                 in if x > a1+                                        then inner (pos + 1) x a1+                                        else inner (pos + 1) a1 (max a2 x)+                        !res = inner (VU.unsafeIndex offs g) ninf ninf+                    VUM.unsafeWrite out g res+                    grp (g + 1)+         in grp gs+    VU.unsafeFreeze out+{-# INLINE top2SndPar #-}++-------------------------------------------------------------------------------+-- Parallel fused max(a) - min(b) (Q7 at wide group domains)+-------------------------------------------------------------------------------++{- | Fused @max a - min b@ over the group-range layout: ONE traversal of+@valueIndices@ accumulating both extrema, parallel by disjoint group range with+no cross-worker merge. min/max are order-independent, so the result is+byte-identical to running the two gather extrema passes separately; the fusion+halves the index traffic. 'Nothing' below the parallel threshold or unless both+columns are clean unboxed and same-typed (Int/Int keeps the Int result of the+interpreter; Double/Double the Double one) — the caller then keeps its two-pass+fallback.+-}+maxMinusMinScatterPar ::+    VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -> Maybe Column+maxMinusMinScatterPar vis offs nGroups ca cb+    | not (shouldPar (VU.length vis)) || nGroups <= 1 = Nothing+    | otherwise = case (ca, cb) of+        ( UnboxedColumn Nothing (va :: VU.Vector x)+            , UnboxedColumn Nothing (vb :: VU.Vector y)+            )+                | Just Refl <- testEquality (typeRep @x) (typeRep @Int)+                , Just Refl <- testEquality (typeRep @y) (typeRep @Int) ->+                    Just (maxMinusMinParInt vis offs nGroups va vb caps bounds)+                | Just Refl <- testEquality (typeRep @x) (typeRep @Double)+                , Just Refl <- testEquality (typeRep @y) (typeRep @Double) ->+                    Just (maxMinusMinParDbl vis offs nGroups va vb caps bounds)+        _ -> Nothing+  where+    !caps = capabilities+    !bounds = groupRangeBounds offs nGroups caps+{-# NOINLINE maxMinusMinScatterPar #-}++{- | Monomorphic entry points (see 'reduceParInt' for why the 'testEquality'+dispatch needs them).+-}+maxMinusMinParInt ::+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector Int ->+    Column+maxMinusMinParInt vis offs nGroups va vb caps bounds =+    fromUnboxedVector+        ( unsafePerformIO+            (maxMinusMinPar minBound maxBound vis offs nGroups va vb caps bounds)+        )+{-# NOINLINE maxMinusMinParInt #-}++maxMinusMinParDbl ::+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector Double ->+    VU.Vector Double ->+    Int ->+    VU.Vector Int ->+    Column+maxMinusMinParDbl vis offs nGroups va vb caps bounds =+    fromUnboxedVector+        ( unsafePerformIO+            (maxMinusMinPar (negate (1 / 0)) (1 / 0) vis offs nGroups va vb caps bounds)+        )+{-# NOINLINE maxMinusMinParDbl #-}++maxMinusMinPar ::+    (VU.Unbox a, Num a, Ord a) =>+    a ->+    a ->+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector a ->+    VU.Vector a ->+    Int ->+    VU.Vector Int ->+    IO (VU.Vector a)+{-# SPECIALIZE maxMinusMinPar ::+    Int ->+    Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector Int ->+    IO (VU.Vector Int)+    #-}+{-# SPECIALIZE maxMinusMinPar ::+    Double ->+    Double ->+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector Double ->+    VU.Vector Double ->+    Int ->+    VU.Vector Int ->+    IO (VU.Vector Double)+    #-}+maxMinusMinPar maxSeed minSeed vis offs nGroups va vb caps bounds = do+    out <- VUM.new nGroups+    parallelBounds_ caps bounds $ \gs ge ->+        -- Both extrema carried in registers per group; one traversal of the+        -- shared index slice reads both value columns.+        let grp !g+                | g >= ge = pure ()+                | otherwise = do+                    let !e = VU.unsafeIndex offs (g + 1)+                        inner !pos !mx !mn+                            | pos >= e = mx - mn+                            | otherwise =+                                let !row = VU.unsafeIndex vis pos+                                 in inner+                                        (pos + 1)+                                        (max mx (VU.unsafeIndex va row))+                                        (min mn (VU.unsafeIndex vb row))+                        !res = inner (VU.unsafeIndex offs g) maxSeed minSeed+                    VUM.unsafeWrite out g res+                    grp (g + 1)+         in grp gs+    VU.unsafeFreeze out++-------------------------------------------------------------------------------+-- Streaming (rowToGroup-scatter) kernels: no valueIndices, no placement pass+-------------------------------------------------------------------------------++{- | Group-count cap for the FUSED streaming rtg-scatter kernels+('DataFrame.Internal.Aggregation.Kernel.Moments.momentStreamPar',+'DataFrame.Internal.Aggregation.Kernel.Fused.runFusedAggs'). Above+'DataFrame.Internal.Grouping.Direct.directThreshold' the per-worker accumulator+arrays overflow cache, so a SINGLE streaming reduction loses to a gather pass —+the per-expression dispatch keeps that threshold. A fused multi-reduction pass+amortizes those misses across all its reductions AND avoids the deferred+@valueIndices@ placement entirely, which flips the comparison (measured at 1e6+groups / 1e8 rows on -N16: rowToGroup 0.4s + fused 3-sum stream 1.4s, against+placement 1.1s + fused gather 1.05s), so the fused cap extends to+'directGroupThreshold' — every direct-grouped frame can stream. Wider+groupings are necessarily hash-path (eager @valueIndices@) and use the fused+GATHER kernel ('DataFrame.Internal.Aggregation.Kernel.Fused.runGatherAggs')+instead. Memory: @capabilities * nGroups@ words per accumulator array, at most+~128MB transient at -N16.+-}+streamGroupCap :: Int+streamGroupCap = 1048576++{- | Near-equal contiguous slices of the group domain for parallel merges.+Below 4096 groups (or single-capability) the merge stays on one thread.+-}+groupSlices :: Int -> [(Int, Int)]+groupSlices = chunksFor 4096++{- | Second-largest value per group: the same (largest, second-largest)+accumulator pair as 'top2Scatter', but the finalize returns the second max+alone. A group of size 1 (or 0) leaves the @-inf@ seed in the second slot, so+its output is NaN — documented behaviour (the db-benchmark Q8 data has no+size-1 @id6@ groups). Like 'top2Scatter''s @-inf -> 0@ guard, an actual+infinite data value in the second slot is indistinguishable from the seed.+-}+top2SndScatter ::+    (VU.Unbox a, Real a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double+top2SndScatter 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+                a2 <- VUM.unsafeRead m2 k+                VUM.unsafeWrite out k (if isInfinite a2 then 0 / 0 else a2)+                fin (k + 1)+    fin 0+    VU.unsafeFreeze out+{-# INLINE top2SndScatter #-}
+ src-internal/DataFrame/Internal/Aggregation/Plan.hs view
@@ -0,0 +1,222 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | The aggregation fast-path planner. 'planAgg' recognises a supported+aggregate shape over a clean unboxed Int/Double column and returns an 'AggPlan';+'planMoments' recognises the six-fold regression shape and returns a+'MomentPlan'. Planning only — the kernels live under "Kernel".+-}+module DataFrame.Internal.Aggregation.Plan (+    AggPlan (..),+    planAgg,+    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 DataFrame.Internal.Aggregation.Reduction (Reduction (..))+import DataFrame.Internal.Column (Column (..))+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)+        "top2Snd" -> require name (PlanScatter RTop2Snd 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])
+ src-internal/DataFrame/Internal/Aggregation/Reduction.hs view
@@ -0,0 +1,44 @@+{-# LANGUAGE ScopedTypeVariables #-}++{- | The vocabulary shared by the aggregation planner and the aggregation+kernels: the set of recognised reductions, and the admission gate that decides+whether a column may enter a Double-valued fast path.++This module holds no kernel and no policy, so the planner can depend on it+without depending on an implementation.+-}+module DataFrame.Internal.Aggregation.Reduction (+    Reduction (..),+    cleanDoubleVector,+) where++import qualified Data.Vector.Unboxed as VU++import DataFrame.Internal.Column (+    Column (..),+    materializePacked,+ )+import DataFrame.Internal.Column.Conversion (toDoubleVector)++{- | 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+    | RTop2Snd+    deriving (Eq, Show)++cleanDoubleVector :: Column -> Maybe (VU.Vector Double)+cleanDoubleVector col = case col of+    UnboxedColumn Nothing _ -> either (const Nothing) Just (toDoubleVector col)+    BoxedColumn Nothing _ -> either (const Nothing) Just (toDoubleVector col)+    p@(PackedText _ _) -> cleanDoubleVector (materializePacked p)+    _ -> Nothing
+ src-internal/DataFrame/Internal/Algorithms/Hash.hs view
@@ -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.Algorithms.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 #-}
+ src-internal/DataFrame/Internal/Algorithms/Rank/Radix.hs view
@@ -0,0 +1,101 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-}++{- | Stable rank of a set of group representatives by ascending unsigned hash+order. Shared by the sequential and parallel group-by canonical-ordering steps+so they stay bit-for-bit identical. @O(ng)@ stable LSD radix sort.+-}+module DataFrame.Internal.Algorithms.Rank.Radix (+    rankByHash,+    sortKey,+) where++import Control.Monad (when)+import Control.Monad.Primitive (PrimMonad)+import Data.Bits (unsafeShiftR, (.&.))+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.Word (Word64)++{- | Unsigned sort key of a hash: ascending 'Word64' order of @sortKey h@ equals+ascending signed-'Int' order of @h@. Reinterpreted to 'Int' for the byte-wise+radix passes (the byte mask makes the sign extension irrelevant).+-}+sortKey :: Int -> Int+sortKey h = fromIntegral (fromIntegral h + 0x8000000000000000 :: Word64)+{-# INLINE sortKey #-}++-- | See the module header. @readHash@ supplies the hash of local group @gid@.+rankByHash ::+    forall m. (PrimMonad m) => (Int -> m Int) -> Int -> m (VU.Vector Int)+rankByHash readHash ng = do+    rankM <- VUM.new (max 1 ng)+    if ng <= 1+        then when (ng == 1) (VUM.unsafeWrite rankM 0 0)+        else do+            keysA <- VUM.new ng+            orderA <- VUM.new ng+            let seed !i+                    | i >= ng = pure ()+                    | otherwise = do+                        h <- readHash i+                        VUM.unsafeWrite keysA i (sortKey h)+                        VUM.unsafeWrite orderA i i+                        seed (i + 1)+            seed 0+            keysB <- VUM.new ng+            orderB <- VUM.new ng+            counts <- VUM.new 256+            let pass ::+                    Int ->+                    VUM.MVector (VUM.PrimState m) Int ->+                    VUM.MVector (VUM.PrimState m) Int ->+                    VUM.MVector (VUM.PrimState m) Int ->+                    VUM.MVector (VUM.PrimState m) Int ->+                    m ()+                pass !shiftBits !srcK !srcO !dstK !dstO = do+                    VUM.set counts 0+                    let count !i+                            | i >= ng = pure ()+                            | otherwise = do+                                k <- VUM.unsafeRead srcK i+                                let !b = (k `unsafeShiftR` shiftBits) .&. 0xff+                                VUM.unsafeRead counts b >>= VUM.unsafeWrite counts b . (+ 1)+                                count (i + 1)+                    count 0+                    let scan !b !acc+                            | b >= 256 = pure ()+                            | otherwise = do+                                c <- VUM.unsafeRead counts b+                                VUM.unsafeWrite counts b acc+                                scan (b + 1) (acc + c)+                    scan 0 0+                    let place !i+                            | i >= ng = pure ()+                            | otherwise = do+                                k <- VUM.unsafeRead srcK i+                                o <- VUM.unsafeRead srcO i+                                let !b = (k `unsafeShiftR` shiftBits) .&. 0xff+                                pos <- VUM.unsafeRead counts b+                                VUM.unsafeWrite counts b (pos + 1)+                                VUM.unsafeWrite dstK pos k+                                VUM.unsafeWrite dstO pos o+                                place (i + 1)+                    place 0+            pass 0 keysA orderA keysB orderB+            pass 8 keysB orderB keysA orderA+            pass 16 keysA orderA keysB orderB+            pass 24 keysB orderB keysA orderA+            pass 32 keysA orderA keysB orderB+            pass 40 keysB orderB keysA orderA+            pass 48 keysA orderA keysB orderB+            pass 56 keysB orderB keysA orderA+            let inv !r+                    | r >= ng = pure ()+                    | otherwise = do+                        g <- VUM.unsafeRead orderA r+                        VUM.unsafeWrite rankM g r+                        inv (r + 1)+            inv 0+    VU.unsafeFreeze rankM+{-# INLINEABLE rankByHash #-}
+ src-internal/DataFrame/Internal/Algorithms/Sort/Radix/Parallel.hs view
@@ -0,0 +1,332 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-}++{- | Parallel stable sort of row indices by ascending unsigned order of a per-row+'Int' hash, used by the join build side. A counting sort buckets rows into+key-ordered partitions that workers LSD-radix-sort in parallel, with no merge step.+-}+module DataFrame.Internal.Algorithms.Sort.Radix.Parallel (+    parSortByHash,+    parSortThreshold,+) where++import Control.Concurrent (getNumCapabilities)+import Control.Monad (forM_, when)+import Data.Bits (countLeadingZeros, unsafeShiftR, (.&.))+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.Word (Word64)+import DataFrame.Internal.Algorithms.Rank.Radix (sortKey)+import DataFrame.Internal.Control.Concurrent (+    capabilities,+    forkJoin,+    forkJoin_,+    pooledIndices,+ )+import System.IO.Unsafe (unsafePerformIO)++{- | Below this many rows the partition/fork overhead is not worth it; the+caller's sequential LSD radix path is used instead.+-}+parSortThreshold :: Int+parSortThreshold = 500000++{- | Top-bits partition index of a hash: the high @64 - shift@ bits of its+unsigned 'sortKey'. Ascending partition order equals ascending key order.+-}+partIx :: Int -> Int -> Int+partIx shift h = fromIntegral ((fromIntegral (sortKey h) :: Word64) `unsafeShiftR` shift)+{-# INLINE partIx #-}++-- | Number of partitions: a power of two, at least @4 * caps@, 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 stable sort of @[0, n)@ by ascending unsigned hash order. See the+module header for the ordering contract.+-}+parSortByHash :: Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)+parSortByHash n hashes+    | n <= 1 =+        (hashes, VU.enumFromN 0 n)+    | n < parSortThreshold || capabilities <= 1 =+        seqSortByHash n hashes+    | otherwise = unsafePerformIO (parSortByHashIO n hashes)+{-# NOINLINE parSortByHash #-}++-------------------------------------------------------------------------------+-- Sequential LSD radix sort (also the per-partition worker kernel)+-------------------------------------------------------------------------------++{- | Stable LSD radix sort of @[0, n)@ by ascending 'sortKey' of their hash, 8+bits per pass over the full 64-bit key. Returns @(sortedHashes, sortedIndices)@.+-}+seqSortByHash :: Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)+seqSortByHash n hashes = unsafePerformIO $ do+    keysA <- VUM.new n+    orderA <- VUM.new n+    let seed !i+            | i >= n = pure ()+            | otherwise = do+                VUM.unsafeWrite keysA i (sortKey (VU.unsafeIndex hashes i))+                VUM.unsafeWrite orderA i i+                seed (i + 1)+    seed 0+    keysB <- VUM.new n+    orderB <- VUM.new n+    radixPasses n keysA orderA keysB orderB+    order <- VU.unsafeFreeze orderA+    pure (VU.unsafeBackpermute hashes order, order)++{- | Run all eight stable 8-bit LSD passes, ping-ponging between the two+key/order buffer pairs so the sorted order lands back in @(keysA, orderA)@.+@keysA[i]@ must already hold @sortKey (hash of orderA[i])@ on entry.+-}+radixPasses ::+    Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    IO ()+radixPasses = radixPassesN 8++{- | Run the first @np@ stable 8-bit LSD passes (bits @0 .. 8*np-1@),+ping-ponging between the buffer pairs. For odd @np@ the sorted order lands in+@(keysB, orderB)@, for even in @(keysA, orderA)@. Callers whose rows share+their top bytes (per-partition sorts partitioned on the top byte) can pass+@np = 7@: the eighth pass would be a stable identity copy.+-}+radixPassesN ::+    Int ->+    Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    IO ()+radixPassesN np n keysA orderA keysB orderB = do+    counts <- VUM.new 256+    let pass ::+            Int ->+            VUM.IOVector Int ->+            VUM.IOVector Int ->+            VUM.IOVector Int ->+            VUM.IOVector Int ->+            IO ()+        pass !shiftBits !srcK !srcO !dstK !dstO = do+            VUM.set counts 0+            let count !i+                    | i >= n = pure ()+                    | otherwise = do+                        k <- VUM.unsafeRead srcK i+                        let !b = (k `unsafeShiftR` shiftBits) .&. 0xff+                        VUM.unsafeRead counts b >>= VUM.unsafeWrite counts b . (+ 1)+                        count (i + 1)+            count 0+            let scan !b !acc+                    | b >= 256 = pure ()+                    | otherwise = do+                        c <- VUM.unsafeRead counts b+                        VUM.unsafeWrite counts b acc+                        scan (b + 1) (acc + c)+            scan 0 0+            let place !i+                    | i >= n = pure ()+                    | otherwise = do+                        k <- VUM.unsafeRead srcK i+                        o <- VUM.unsafeRead srcO i+                        let !b = (k `unsafeShiftR` shiftBits) .&. 0xff+                        pos <- VUM.unsafeRead counts b+                        VUM.unsafeWrite counts b (pos + 1)+                        VUM.unsafeWrite dstK pos k+                        VUM.unsafeWrite dstO pos o+                        place (i + 1)+            place 0+        run !k+            | k >= np = pure ()+            | even k = pass (8 * k) keysA orderA keysB orderB >> run (k + 1)+            | otherwise = pass (8 * k) keysB orderB keysA orderA >> run (k + 1)+    run 0++-------------------------------------------------------------------------------+-- Parallel path: counting-sort partition, then per-partition sort in parallel+-------------------------------------------------------------------------------++parSortByHashIO :: Int -> VU.Vector Int -> IO (VU.Vector Int, VU.Vector Int)+parSortByHashIO n hashes = do+    caps <- getNumCapabilities+    let !p = numPartitionsFor caps+        !shift = 64 - intLog2 p+    (partStart, partRows, partHashes) <- partitionRows n hashes p shift+    outOrder <- VUM.new n+    outKeys <- VUM.new n+    sortPartitions caps p partStart partRows partHashes outOrder outKeys+    order <- VU.unsafeFreeze outOrder+    sortedHashes <- VU.unsafeFreeze outKeys+    pure (sortedHashes, order)++{- | Bucket every row index into its top-bits partition by a counting sort.+Returns the exclusive prefix sum @partStart@ (length @p+1@, @partStart[p] == n@),+the row indices laid out partition-by-partition in ascending key order, and+each sorted position's hash in the same layout (so downstream passes read+hashes sequentially instead of a random @hashes[row]@ gather per row).++Runs chunked across capabilities: per-chunk partition histograms are prefix+summed (in chunk order) into disjoint per-chunk write cursors, so the scatter+threads never contend and each partition keeps its rows in ascending original+row order — bit-for-bit the sequential counting sort's layout.+-}+partitionRows ::+    Int ->+    VU.Vector Int ->+    Int ->+    Int ->+    IO (VU.Vector Int, VU.Vector Int, VU.Vector Int)+partitionRows n hashes p shift = do+    caps <- getNumCapabilities+    let chunks = rowChunks caps n+    cursors <- forkJoin [histChunk hashes p shift lo hi | (lo, hi) <- chunks]+    -- Exclusive prefix over partitions (outer) and chunks (inner): partStart+    -- from the totals, and each chunk's histogram rewritten into its cursor.+    partStartM <- VUM.new (p + 1)+    let seed !pp !acc+            | pp >= p = VUM.unsafeWrite partStartM p acc+            | otherwise = do+                VUM.unsafeWrite partStartM pp acc+                let inner [] !a = pure a+                    inner (cur : rest) !a = do+                        t <- VUM.unsafeRead cur pp+                        VUM.unsafeWrite cur pp a+                        inner rest (a + t)+                acc' <- inner cursors acc+                seed (pp + 1) acc'+    seed 0 0+    rowsM <- VUM.new (max 1 n)+    rowHashM <- VUM.new (max 1 n)+    forkJoin_+        [ scatterChunk hashes shift cur rowsM rowHashM lo hi+        | ((lo, hi), cur) <- zip chunks cursors+        ]+    partStart <- VU.unsafeFreeze partStartM+    partRows <- VU.unsafeFreeze rowsM+    partHashes <- VU.unsafeFreeze rowHashM+    pure (partStart, partRows, partHashes)++-- | Contiguous near-equal row chunks, one per capability; empties dropped.+rowChunks :: Int -> Int -> [(Int, Int)]+rowChunks caps n =+    [ (lo, hi)+    | w <- [0 .. caps - 1]+    , let lo = min n (w * per)+    , let hi = min n (lo + per)+    , lo < hi+    ]+  where+    !per = (n + max 1 caps - 1) `div` max 1 caps++-- | Per-partition counts of one row chunk.+histChunk :: VU.Vector Int -> Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)+histChunk hashes p shift lo hi = do+    acc <- VUM.replicate p (0 :: Int)+    let go !i+            | i >= hi = pure acc+            | otherwise = do+                let !pp = partIx shift (VU.unsafeIndex hashes i)+                c <- VUM.unsafeRead acc pp+                VUM.unsafeWrite acc pp (c + 1)+                go (i + 1)+    go lo++{- | Scatter one row chunk into the partitioned layout using the chunk's+pre-summed cursor (disjoint write regions per chunk, no contention).+-}+scatterChunk ::+    VU.Vector Int ->+    Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    IO ()+scatterChunk hashes shift cursor rowsM rowHashM lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !h = VU.unsafeIndex hashes i+                !pp = partIx shift h+            pos <- VUM.unsafeRead cursor pp+            VUM.unsafeWrite rowsM pos i+            VUM.unsafeWrite rowHashM pos h+            VUM.unsafeWrite cursor pp (pos + 1)+            go (i + 1)++{- | Stable-sort each partition by full key, writing sorted original indices+into @outOrder@ and their hashes into @outKeys@ at the partition's slot range.+Forks @caps@ workers that pull partition indices off a shared atomic counter.+Within a partition the counting sort already left rows in ascending original+order, so the LSD radix sort's stability reproduces the global @(key, row)@+order. Partitions below two elements are already sorted (counting sort kept+original order) and are copied directly.++@partHashes@ is the partition-layout hash vector from 'partitionRows', so+seeding reads hashes sequentially; only 7 LSD passes run (the top byte is the+partition byte, constant within a partition), and the sorted hash is recovered+from the sort key ('sortKey' is self-inverse) instead of a random gather.+-}+sortPartitions ::+    Int ->+    Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    IO ()+sortPartitions caps p partStart partRows partHashes outOrder outKeys =+    pooledIndices caps p sortOne+  where+    sortOne !pp = do+        let !s = VU.unsafeIndex partStart pp+            !e = VU.unsafeIndex partStart (pp + 1)+            !sz = e - s+        when (sz > 0) $+            if sz == 1+                then do+                    VUM.unsafeWrite outOrder s (VU.unsafeIndex partRows s)+                    VUM.unsafeWrite outKeys s (VU.unsafeIndex partHashes s)+                else do+                    keysA <- VUM.new sz+                    orderA <- VUM.new sz+                    let seed !i+                            | i >= sz = pure ()+                            | otherwise = do+                                VUM.unsafeWrite keysA i (sortKey (VU.unsafeIndex partHashes (s + i)))+                                VUM.unsafeWrite orderA i (VU.unsafeIndex partRows (s + i))+                                seed (i + 1)+                    seed 0+                    keysB <- VUM.new sz+                    orderB <- VUM.new sz+                    radixPassesN 7 sz keysA orderA keysB orderB+                    let emit !i+                            | i >= sz = pure ()+                            | otherwise = do+                                o <- VUM.unsafeRead orderB i+                                k <- VUM.unsafeRead keysB i+                                VUM.unsafeWrite outOrder (s + i) o+                                -- sortKey is an involution: recover the hash.+                                VUM.unsafeWrite outKeys (s + i) (sortKey k)+                                emit (i + 1)+                    emit 0
+ src-internal/DataFrame/Internal/Column.hs view
@@ -0,0 +1,29 @@+{- |+Umbrella re-export of the column implementation. The module is split by+concern:++* "DataFrame.Internal.Column.Types" — type-level machinery ('Rep', 'SBool', ...)+* "DataFrame.Internal.Column.Base" — the 'Column' GADT and core definitions+* "DataFrame.Internal.Column.Bitmap" — the validity bitmap and its helpers+* "DataFrame.Internal.Column.Properties" — predicates and introspection+* "DataFrame.Internal.Column.Conversion" — vector\/list conversions+* "DataFrame.Internal.Column.Operations" — bulk transformations++Import this module to get the whole surface; import a submodule directly when+you only need one layer.+-}+module DataFrame.Internal.Column (+    module DataFrame.Internal.Column.Base,+    module DataFrame.Internal.Column.Bitmap,+    module DataFrame.Internal.Column.Conversion,+    module DataFrame.Internal.Column.Operations,+    module DataFrame.Internal.Column.Properties,+    module DataFrame.Internal.Column.Types,+) where++import DataFrame.Internal.Column.Base+import DataFrame.Internal.Column.Bitmap+import DataFrame.Internal.Column.Conversion+import DataFrame.Internal.Column.Operations+import DataFrame.Internal.Column.Properties+import DataFrame.Internal.Column.Types
+ src-internal/DataFrame/Internal/Column/Base.hs view
@@ -0,0 +1,333 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}++{- |+Core column definitions: the type-erased 'Column' GADT and its mutable/typed+companions, the 'Columnable' constraint, and the representation-level+primitives ('materializePacked', 'materializeMerged') that the non-orphan+'Show'/'Eq' instances depend on.++Predicates live in "DataFrame.Internal.Column.Properties", vector/list+conversions in "DataFrame.Internal.Column.Conversion", and bulk transformations+in "DataFrame.Internal.Column.Operations".+-}+module DataFrame.Internal.Column.Base where++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.Monad (forM_)+import Control.Monad.ST (runST)+import Data.Maybe (fromMaybe, isNothing)+import Data.Type.Equality (TestEquality (..))+import DataFrame.Internal.Column.Bitmap+import DataFrame.Internal.Column.Types+import DataFrame.Internal.Data.PackedText (+    PackedTextData (..),+    packedIndexText,+    packedLength,+    packedSlice,+    packedTake,+    sliceEqBytes,+ )+import Type.Reflection (typeRep, type (:~:) (Refl))++-- | 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)+    )++{- | 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+    -- TODO: mchavinda - investigate splitting this into a separate intermediate+    --                   representation.+    -- Efficient intermediate formats.z+    -- 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 join's same-named non-key column pair ('mkMergedColumns'): both sides+    -- keep their native (packed/dict/unboxed) representation; per-row 'These'+    -- values only materialize on element access ('materializeMerged').+    MergedColumn :: !Column -> !Column -> Column++instance Show Column where+    show :: Column -> String+    show c@(MergedColumn _ _) = show (materializeMerged c)+    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+                        )+    (==) lhs@(MergedColumn _ _) rhs = materializeMerged lhs == rhs+    (==) lhs rhs@(MergedColumn _ _) = lhs == materializeMerged rhs+    (==) (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 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++{- | 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++instance (Show a) => Show (TypedColumn a) where+    show :: (Show a) => TypedColumn a -> String+    show (TColumn col) = show col++-- | Unwrap a 'TypedColumn' back to its type-erased 'Column'.+unwrapTypedColumn :: TypedColumn a -> Column+unwrapTypedColumn (TColumn value) = value++{- | 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) =+    BoxedColumn bm (VB.generate (packedLength p) (packedIndexText p))+materializePacked c = c+{-# INLINE materializePacked #-}++-- | Return the 'Maybe Bitmap' from a column.+columnBitmap :: Column -> Maybe Bitmap+columnBitmap (BoxedColumn bm _) = bm+columnBitmap (UnboxedColumn bm _) = bm+columnBitmap (PackedText bm _) = bm+columnBitmap (MergedColumn _ _) = Nothing++{- | 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++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 v =+        let+            n = VB.length v+            nullIdxs = VU.filter (isNothing . VB.unsafeIndex v) (VU.enumFromN 0 n)+            bm =+                if VU.null nullIdxs then allValidBitmap n else buildBitmapFromNulls' n nullIdxs+         in+            case sUnbox @a of+                STrue -> UnboxedColumn (Just bm) $ runST $ do+                    mv <- VUM.new n+                    VG.iforM_ v $ \i mx -> forM_ mx (VUM.unsafeWrite mv i)+                    VU.unsafeFreeze mv+                SFalse ->+                    BoxedColumn+                        (Just bm)+                        (VB.map (fromMaybe (errorWithoutStackTrace "toColumnRep: Nothing slot")) v)++-- | O(1) Gets the number of elements in the column.+columnLength :: Column -> Int+columnLength (MergedColumn a b) = min (columnLength a) (columnLength b)+columnLength (BoxedColumn _ xs) = VB.length xs+columnLength (UnboxedColumn _ xs) = VU.length xs+columnLength (PackedText _ p) = packedLength p+{-# INLINE columnLength #-}++-- | O(n) Takes the first n values of a column.+takeColumn :: Int -> Column -> Column+takeColumn n (MergedColumn a b) = MergedColumn (takeColumn n a) (takeColumn n b)+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 #-}++{- | Merge two columns using `These`. O(1): the sides are kept in their+native representation and 'These' values materialize on element access.+-}+mkMergedColumns :: Column -> Column -> Column+mkMergedColumns = MergedColumn+{-# INLINE mkMergedColumns #-}++-- | Decode a 'MergedColumn' into the eager @BoxedColumn (These a b)@ form.+materializeMerged :: Column -> Column+materializeMerged (MergedColumn colA colB) =+    mergeEager (materializeMerged colA) (materializeMerged colB)+materializeMerged c = c++mergedHead :: Column -> Column+mergedHead (MergedColumn a b) =+    materializeMerged (MergedColumn (takeColumn 1 a) (takeColumn 1 b))+mergedHead c = c++{- | The eager element-wise merge ('These' per row, boxed). Bitmaps are+honored for every representation pair: a null side yields 'This'/'That',+both-null is an error (the join kernels never produce such a row).+-}+mergeEager :: Column -> Column -> Column+mergeEager colA colB = case (colA, colB) of+    (MergedColumn a b, _) -> mergeEager (mergeEager a b) colB+    (_, MergedColumn a b) -> mergeEager colA (mergeEager a b)+    (PackedText _ _, _) -> mergeEager (materializePacked colA) colB+    (_, PackedText _ _) -> mergeEager colA (materializePacked colB)+    (BoxedColumn bmA c1, BoxedColumn bmB c2) ->+        merged bmA bmB (VG.length c1) (VG.length c2) (c1 VG.!) (c2 VG.!)+    (BoxedColumn bmA c1, UnboxedColumn bmB c2) ->+        merged bmA bmB (VG.length c1) (VG.length c2) (c1 VG.!) (c2 VG.!)+    (UnboxedColumn bmA c1, BoxedColumn bmB c2) ->+        merged bmA bmB (VG.length c1) (VG.length c2) (c1 VG.!) (c2 VG.!)+    (UnboxedColumn bmA c1, UnboxedColumn bmB c2) ->+        merged bmA bmB (VG.length c1) (VG.length c2) (c1 VG.!) (c2 VG.!)+  where+    merged ::+        (Columnable a, Columnable b) =>+        Maybe Bitmap ->+        Maybe Bitmap ->+        Int ->+        Int ->+        (Int -> a) ->+        (Int -> b) ->+        Column+    merged bmA bmB lenA lenB atA atB =+        BoxedColumn Nothing $ VB.generate (min lenA lenB) $ \i ->+            case (validAt bmA i, validAt bmB i) of+                (True, True) -> These (atA i) (atB i)+                (True, False) -> This (atA i)+                (False, True) -> That (atB i)+                (False, False) -> error "mkMergedColumns: both null"+    validAt mbm i = maybe True (`bitmapTestBit` i) mbm+    {-# INLINE validAt #-}
+ src-internal/DataFrame/Internal/Column/Bitmap.hs view
@@ -0,0 +1,188 @@+{-# LANGUAGE BangPatterns #-}++module DataFrame.Internal.Column.Bitmap where++import Control.Monad (foldM_, forM_, when)+import Control.Monad.ST (ST, runST)+import Data.Bits (+    complement,+    popCount,+    setBit,+    shiftL,+    shiftR,+    testBit,+    (.&.),+    (.|.),+ )+import Data.List (foldl')+import Data.Maybe (fromMaybe, isNothing)+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.Word (Word8)++-- | A bit-packed validity bitmap. Bit @i@ = 1 means row @i@ is valid (not null).+type Bitmap = VU.Vector Word8++-- | A bitmap attached to its row counts so we can splice it.+data Validity = Validity !(Maybe Bitmap) {-# UNPACK #-} !Int++vBitmap :: Validity -> Maybe Bitmap+vBitmap (Validity bm _) = bm++vRowCount :: Validity -> Int+vRowCount (Validity _ n) = n++{- | Test whether row @i@ is valid (not null) in a bitmap.++The bit-level arithmetric is dense but read this as:+`shiftR` 3 is equivalent to `div` 8, and `.&. 7` is equivalent to `mod` 8.+-}+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 = runST (allValidBitmap' n >>= VU.unsafeFreeze)+{-# INLINE allValidBitmap #-}++allValidBitmap' :: Int -> ST s (VUM.MVector s Word8)+allValidBitmap' n =+    let+        bytes = (n + 7) `shiftR` 3+        lastBits = n .&. 7+        lastByte = if lastBits == 0 then 0xFF else (1 `shiftL` lastBits) - 1+     in+        if bytes == 0+            then VUM.new 0+            else do+                mv <- VUM.replicate bytes (0xFF :: Word8) :: ST s (VUM.MVector s Word8)+                when (lastBits /= 0) $ VUM.unsafeWrite mv (bytes - 1) lastByte+                pure mv+{-# 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 idxs = buildBitmapFromNulls' n (VU.fromList idxs)++buildBitmapFromNulls' :: Int -> VU.Vector Int -> VU.Vector Word8+buildBitmapFromNulls' n nullIdxs = runST $ do+    bm' <- allValidBitmap' n+    VU.forM_ nullIdxs $ \i -> do+        let byteIdx = i `shiftR` 3+            bitIdx = i .&. 7+        v <- VUM.unsafeRead bm' byteIdx+        VUM.unsafeWrite bm' byteIdx (clearBit8 v bitIdx)+    VU.unsafeFreeze bm'+  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++{- | Count the set bits among the first @n@ bits of a bitmap. A bitmap does+not know the length of the column it describes, and 'bitmapSlice' keeps whole+bytes on its aligned path, so the bits past @n@ may still describe rows+outside the slice.+-}+popCountUpTo :: Int -> Bitmap -> Int+popCountUpTo n bm = whole + partial+  where+    !fullBytes = min (n `shiftR` 3) (VU.length bm)+    !rest = n .&. 7+    whole = VU.foldl' (\acc b -> acc + popCount b) 0 (VU.take fullBytes bm)+    partial+        | rest == 0 || fullBytes >= VU.length bm = 0+        | otherwise =+            popCount (VU.unsafeIndex bm fullBytes .&. ((1 `shiftL` rest) - 1))+{-# INLINE popCountUpTo #-}++-- | 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).+andBitmaps :: Bitmap -> Bitmap -> Bitmap+andBitmaps = VU.zipWith (.&.)++{- | 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.+-}+concatValidity :: [Validity] -> Maybe Bitmap+concatValidity parts+    | all (isNothing . vBitmap) parts = Nothing+    | otherwise = Just $ VU.create $ do+        let total = sum (map vRowCount 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 (Validity !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
+ src-internal/DataFrame/Internal/Column/Builder.hs view
@@ -0,0 +1,306 @@+{-# 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.Column.Builder (+    ColumnBuilder (..),+    NumBuilder,+    IntBuilder,+    DoubleBuilder,+    TextBuilder,+    TextChunk (..),+    newIntBuilder,+    newDoubleBuilder,+    newNumBuilder,+    newTextBuilder,+    appendInt,+    appendDouble,+    appendNum,+    appendText,+    appendTextSlice,+    appendTextSliceFromPtr,+    freezeTextChunk,+    concatColumns,+    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 (+    Column (UnboxedColumn),+    Columnable,+ )+import DataFrame.Internal.Column.Bitmap (packValidity)+import DataFrame.Internal.Column.Merge (+    TextChunk (..),+    concatColumns,+    mergeTextChunks,+ )+import Foreign.Ptr (Ptr)++{- | Operations shared by all column builders.++NB: Do not use after freezing+-}+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'.+    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.+-}+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.+{-# SPECIALIZE newNumBuilder :: Int -> Int -> ST s (NumBuilder Int s) #-}+{-# SPECIALIZE newNumBuilder :: Double -> Int -> ST s (NumBuilder Double s) #-}+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.++Representation is packed. I.e all field bytes go into one exponentially+grown byte array with rows recorded as offsets.+-}+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 later merging+('mergeTextChunks').+-}+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]
+ src-internal/DataFrame/Internal/Column/Conversion.hs view
@@ -0,0 +1,441 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}++{- |+Conversions between 'Column' and ordinary vectors\/lists, plus the typed+extraction functions ('toVector', 'toDoubleVector', ...) that recover a+column's element type.+-}+module DataFrame.Internal.Column.Conversion where++import qualified Data.Text as T+import qualified Data.Vector as VB+import qualified Data.Vector.Generic as VG+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM++import Control.Exception (throw)+import Control.Monad.ST (ST)+import Data.Kind (Type)+import Data.Type.Equality (TestEquality (..))+import Data.Word (Word8)+import DataFrame.Errors (+    DataFrameException (ExpectedNonNullableException, TypeMismatchException),+    TypeErrorContext (+        MkTypeErrorContext,+        callingFunctionName,+        errorColumnName,+        expectedType,+        userType+    ),+ )+import DataFrame.Internal.Column.Base+import DataFrame.Internal.Column.Bitmap+import DataFrame.Internal.Column.Properties+import DataFrame.Internal.Column.Types+import DataFrame.Internal.Data.PackedText (packedIndexText, packedLength)+import System.Random (RandomGen, UniformRange, uniformR)+import Type.Reflection (TypeRep, Typeable, typeRep, type (:~:) (Refl))++{- | 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)++{- | 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)+    MergedColumn _ _ -> toVector (materializeMerged 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+                                    }+                                )+{-# INLINEABLE toVector #-}++-- 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)+        MergedColumn _ _ -> toDoubleVector (materializeMerged column)+        UnboxedColumn (Just _) _ -> Left ExpectedNonNullableException+        UnboxedColumn Nothing (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Double) of+            Just Refl -> Right f+            Nothing -> case sFloating @a of+                STrue -> Right (VU.map realToFrac f)+                SFalse -> case sIntegral @a of+                    STrue -> Right (VU.map fromIntegral f)+                    SFalse ->+                        Left $+                            TypeMismatchException+                                ( MkTypeErrorContext+                                    { userType = Right (typeRep @Double)+                                    , expectedType = Right (typeRep @a)+                                    , callingFunctionName = Just "toDoubleVector"+                                    , errorColumnName = Nothing+                                    }+                                )+        BoxedColumn (Just _) (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of+            Just Refl -> Left ExpectedNonNullableException+            Nothing ->+                Left $+                    TypeMismatchException+                        ( MkTypeErrorContext+                            { userType = Right (typeRep @Double)+                            , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())+                            , callingFunctionName = Just "toDoubleVector"+                            , errorColumnName = Nothing+                            }+                        )+        BoxedColumn Nothing (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 @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)+        MergedColumn _ _ -> toFloatVector (materializeMerged 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)+        MergedColumn _ _ -> toIntVector (materializeMerged 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))++-- | Convert any Column to a vector of Text labels (one per row).+columnToTextVec :: Column -> VB.Vector T.Text+columnToTextVec c@(MergedColumn _ _) = columnToTextVec (materializeMerged c)+columnToTextVec (BoxedColumn bm (col' :: VB.Vector a)) =+    case bm of+        Nothing -> case testEquality (typeRep @a) (typeRep @T.Text) of+            Just Refl -> col'+            Nothing -> VB.map (T.pack . show) col'+        Just bitmap ->+            VB.imap+                (\i x -> if bitmapTestBit bitmap i then T.pack (show x) else "null")+                col'+columnToTextVec (UnboxedColumn bm col') =+    case bm of+        Nothing -> VB.map (T.pack . show) (VB.convert col')+        Just bitmap ->+            VB.generate (VU.length col') $ \i ->+                if bitmapTestBit bitmap i then T.pack (show (col' VU.! i)) else "null"+columnToTextVec (PackedText bm p) =+    VB.generate (packedLength p) $ \i -> case bm of+        Just bitmap | not (bitmapTestBit bitmap i) -> "null"+        _ -> packedIndexText p i++-- 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+                }
+ src-internal/DataFrame/Internal/Column/Encode.hs view
@@ -0,0 +1,185 @@+{-# 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.++TODO: mchavinda - revise if this module is still necessary.+-}+module DataFrame.Internal.Column.Encode (+    dictEncodeColumn,+    dictEncodeColumnUpTo,+    dictCompactColumn,+    dictMaxCardinality,+) where++import Control.Monad (when)+import Control.Monad.ST (runST)+import qualified Data.Text as T+import qualified Data.Text.Array as A+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.Algorithms.Hash (+    fnvOffset,+    mixBytes,+    mixText,+    nullSalt,+ )+import DataFrame.Internal.Column (Column (..))+import DataFrame.Internal.Column.Bitmap (Bitmap, bitmapTestBit)+import DataFrame.Internal.Data.HashTable (htInsert, newHashTable)+import DataFrame.Internal.Data.PackedText (+    PackedTextData (..),+    mkOffsets,+    mkSel,+    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))++dictCompactColumn :: Column -> Column+dictCompactColumn col@(PackedText bm p) =+    case encodePacked dictMaxCardinality bm p of+        Just (codes, card)+            | 2 * card <= packedLength p ->+                PackedText bm (dictPacked p codes card)+        _ -> col+dictCompactColumn col = col++dictPacked :: PackedTextData -> VU.Vector Int -> Int -> PackedTextData+dictPacked p codes card = runST $ do+    let n = VU.length codes+    reps <- VUM.replicate card (-1)+    let findReps !i !remaining+            | remaining <= 0 || i >= n = pure ()+            | otherwise = do+                let c = VU.unsafeIndex codes i+                cur <- VUM.unsafeRead reps c+                if cur < 0+                    then VUM.unsafeWrite reps c i >> findReps (i + 1) (remaining - 1)+                    else findReps (i + 1) remaining+    findReps 0 card+    repsV <- VU.unsafeFreeze reps+    let lens = VU.map (\r -> let (_, _, l) = packedSlice p r in l) repsV+        offs = VU.scanl' (+) 0 lens+        total = VU.last offs+    marr <- A.new (max 1 total)+    let copyRep !c =+            when (c < card) $ do+                let r = VU.unsafeIndex repsV c+                    (arr, o, l) = packedSlice p r+                A.copyI l marr (VU.unsafeIndex offs c) arr o+                copyRep (c + 1)+    copyRep 0+    arr <- A.unsafeFreeze marr+    pure+        ( PackedTextData+            { ptBytes = arr+            , ptOffsets = mkOffsets offs+            , ptSel = Just (mkSel card codes)+            , ptCanonicalSel = True+            }+        )
+ src-internal/DataFrame/Internal/Column/Merge.hs view
@@ -0,0 +1,138 @@+{-# 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.Column.Merge (+    TextChunk (..),+    concatColumns,+    mergeTextChunks,+    packedFromTextChunk,+    concatValidity,+    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_)+import Control.Monad.ST (runST)+import Data.Type.Equality (testEquality, (:~:) (Refl))+import DataFrame.Internal.Column (+    Column (..),+    Columnable,+    isMergedColumn,+    isPackedText,+    materializeMerged,+    materializePacked,+ )+import DataFrame.Internal.Column.Bitmap (+    Bitmap,+    Validity (Validity),+    concatValidity,+ )+import DataFrame.Internal.Data.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 = concatValidity [Validity (tcBitmap c) (tcRows c) | c <- cs]+    pure (PackedText bm (mkPackedContiguous farr foffs))++{- | Merge per-chunk columns into one column.++TODO: mchavinda - this is very similar to mappendColumns can could possibly+be defined in terms of it but I'll have to ivnestigate further.+-}+concatColumns :: [Column] -> Column+concatColumns [] = error "DataFrame.Internal.Column.Builder.concatColumns: empty list"+concatColumns [c] = c+-- Normalize on the whole list, not the head: a packed or merged chunk in any+-- position must demote every chunk to the common boxed form.+concatColumns cols@(c0 : _)+    | any isMergedColumn cols = concatColumns (map materializeMerged cols)+    | any isPackedText cols = concatColumns (map materializePacked cols)+concatColumns cols@(c0 : _) = case c0 of+    PackedText _ _ -> concatColumns (map materializePacked cols)+    MergedColumn _ _ -> concatColumns (map materializeMerged cols)+    UnboxedColumn _ (_ :: VU.Vector a) ->+        let parts = map (unboxedPart @a) cols+            !merged = VU.concat (map snd parts)+            !bm = concatValidity [Validity 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 = concatValidity [Validity 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.Column.Builder.concatColumns: chunk column types differ"
+ src-internal/DataFrame/Internal/Column/Operations.hs view
@@ -0,0 +1,1369 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}++{- |+Bulk operations over columns: mapping, folding, slicing, gathering, zipping,+appending, and the mutable-column IO helpers.+-}+module DataFrame.Internal.Column.Operations 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.Monad (when)+import Control.Monad.ST (runST)+import Data.Bits (setBit, shiftL, shiftR)+import Data.Int (Int32)+import Data.Kind (Type)+import Data.Maybe (catMaybes, fromMaybe, isNothing)+import Data.Type.Equality (TestEquality (..))+import Data.Word (Word8)+import DataFrame.Errors (+    DataFrameException (EmptyDataSetException, TypeMismatchException),+    TypeErrorContext (+        MkTypeErrorContext,+        callingFunctionName,+        errorColumnName,+        expectedType,+        userType+    ),+ )+import DataFrame.Internal.Column.Base+import DataFrame.Internal.Column.Bitmap+import DataFrame.Internal.Column.Conversion+import DataFrame.Internal.Column.Properties+import DataFrame.Internal.Column.Types+import DataFrame.Internal.Control.Concurrent (+    parThreshold,+    parallelChunks_,+    shouldParallelize,+ )+import DataFrame.Internal.Data.PackedText (+    PackedOffsets,+    PackedSel (..),+    PackedTextData (..),+    offCount,+    packedGather,+    packedLength,+ )+import System.IO.Unsafe (unsafePerformIO)+import Type.Reflection (+    TypeRep,+    Typeable,+    typeOf,+    typeRep,+    withTypeable,+    type (:~:) (Refl),+ )++{- | 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` ()+forceColumn (MergedColumn a b) =+    forceColumn a `seq` forceColumn b `seq` checkMergedNoBothNull a b++{- | 'MergedColumn' defers element construction, so forcing must still surface+the one deferred error — a row null on both sides — inside strict IO/executor+boundaries. O(rows) bitmap walk, no allocation; both-null needs a bitmap on+each side, so anything else passes immediately.+-}+checkMergedNoBothNull :: Column -> Column -> ()+checkMergedNoBothNull a b = case (columnBitmap a, columnBitmap b) of+    (Just ba, Just bb) ->+        let !n = min (columnLength a) (columnLength b)+            go !i+                | i >= n = ()+                | bitmapTestBit ba i || bitmapTestBit bb i = go (i + 1)+                | otherwise = error "mkMergedColumns: both null"+         in go 0+    _ -> ()++-- | 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)+newMutableColumn n c@(MergedColumn _ _) = newMutableColumn n (materializeMerged c)++-- | Copy a column chunk into a mutable column starting at offset @off@.+copyIntoMutableColumn :: MutableColumn -> Int -> Column -> IO ()+copyIntoMutableColumn mv off c@(MergedColumn _ _) =+    copyIntoMutableColumn mv off (materializeMerged c)+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++-- | 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)+    c@(MergedColumn _ _) -> mapColumn f (materializeMerged 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 $+                        parGenerateUnboxed 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+                            (parGenerateUnboxed (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 $+                        parGenerateUnboxed 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+                        (parGenerateUnboxed (VU.length col) (f . VU.unsafeIndex 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)+    c@(MergedColumn _ _) -> imapColumn f (materializeMerged 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(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 (MergedColumn a b) =+    MergedColumn (sliceColumn start n a) (sliceColumn start n b)+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 (PackedText bm p)+    -- packedGather decodes an out-of-range index as the empty string, where+    -- the boxed and unboxed arms reject the slice, so check first.+    | start < 0 || n < 0 || start + n > packedLength p =+        errorWithoutStackTrace+            ( "sliceColumn: invalid slice ("+                ++ show start+                ++ ","+                ++ show n+                ++ ","+                ++ show (packedLength p)+                ++ ")"+            )+    | otherwise =+        PackedText+            (fmap (bitmapSlice start n) bm)+            (packedGather (VU.enumFromN start n) p)+{-# INLINE sliceColumn #-}++-- | O(n) Selects the elements at a given set of indices. Does not change the order.++-------------------------------------------------------------------------------+-- Parallel element-wise kernels+-------------------------------------------------------------------------------++{- | Parallel unboxed 'VU.generate': splits the index space into one contiguous+chunk per capability, evaluates each chunk into its disjoint slice of a single+pre-allocated mutable vector, then freezes. Element @i@ depends only on @f i@,+so the result is bit-identical to the sequential 'VU.generate' regardless of+capability count. Falls back to 'VU.generate' below 'parThreshold'.+-}+{-# SPECIALIZE parGenerateUnboxed ::+    Int -> (Int -> Double) -> VU.Vector Double+    #-}+{-# SPECIALIZE parGenerateUnboxed ::+    Int -> (Int -> Float) -> VU.Vector Float+    #-}+{-# SPECIALIZE parGenerateUnboxed :: Int -> (Int -> Int) -> VU.Vector Int #-}+{-# SPECIALIZE parGenerateUnboxed :: Int -> (Int -> Bool) -> VU.Vector Bool #-}+parGenerateUnboxed :: (VU.Unbox c) => Int -> (Int -> c) -> VU.Vector c+parGenerateUnboxed n f+    | not (shouldParallelize parThreshold n) = VU.generate n f+    | otherwise = unsafePerformIO $ do+        mv <- VUM.unsafeNew n+        parallelChunks_ parThreshold n (fillGenerate mv f)+        VU.unsafeFreeze mv+{-# NOINLINE parGenerateUnboxed #-}++{- | The chunk body shared by 'parGenerateUnboxed' and+'parGenerateUnboxedInline'. INLINE so the latter's monomorphic wrappers each+get their own copy with @f@ inlined.+-}+fillGenerate ::+    (VU.Unbox c) => VUM.IOVector c -> (Int -> c) -> Int -> Int -> IO ()+fillGenerate mv f !lo !hi =+    let go !i+            | i >= hi = pure ()+            | otherwise = VUM.unsafeWrite mv i (f i) >> go (i + 1)+     in go lo+{-# INLINE fillGenerate #-}++{- | Parallel unboxed gather: element @i@ of the result is+@v ! (ix ! i)@ (unsafe indexing — callers pass in-bounds index vectors, e.g.+grouping-produced representative rows). Same chunking as 'parGenerateUnboxed'+(one contiguous chunk per capability into disjoint slices of one buffer), so+the result is bit-identical to the sequential backpermute at any capability+count; falls back to a sequential loop below 'parThreshold'.+-}+parBackpermuteUnboxed ::+    (VU.Unbox a) => VU.Vector a -> VU.Vector Int -> VU.Vector a+parBackpermuteUnboxed v ix =+    parGenerateUnboxed (VU.length ix) (VU.unsafeIndex v . VU.unsafeIndex ix)+{-# INLINE parBackpermuteUnboxed #-}++{- | 'parGenerateUnboxed' with an INLINE body: each monomorphic NOINLINE+wrapper below gets its own copy of the fill loop with @f@ inlined, so the+per-element unknown closure call (and the boxed result it returns) disappears+— on a 1e8-row gather that call+alloc dominated the whole pass. Same chunking,+bit-identical results; wrappers must stay NOINLINE so the 'unsafePerformIO'+runs once per call.+-}+parGenerateUnboxedInline :: (VU.Unbox c) => Int -> (Int -> c) -> VU.Vector c+parGenerateUnboxedInline n f+    | not (shouldParallelize parThreshold n) = VU.generate n f+    | otherwise = unsafePerformIO $ do+        mv <- VUM.unsafeNew n+        parallelChunks_ parThreshold n (fillGenerate mv f)+        VU.unsafeFreeze mv+{-# INLINE parGenerateUnboxedInline #-}++-- | Closure-free parallel 'Int' gather: @out!i = v ! (ix!i)@.+parBackpermuteInt :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int+parBackpermuteInt v ix =+    parGenerateUnboxedInline+        (VU.length ix)+        (VU.unsafeIndex v . VU.unsafeIndex ix)+{-# NOINLINE parBackpermuteInt #-}++-- | Closure-free parallel 'Double' gather: @out!i = v ! (ix!i)@.+parBackpermuteDouble :: VU.Vector Double -> VU.Vector Int -> VU.Vector Double+parBackpermuteDouble v ix =+    parGenerateUnboxedInline+        (VU.length ix)+        (VU.unsafeIndex v . VU.unsafeIndex ix)+{-# NOINLINE parBackpermuteDouble #-}++{- | Closure-free double-indirection gather: @out!g = vis ! (offs!g)@ over+@length offs - 1@ groups (the representative-row build of the 'Grouped'+pattern).+-}+parBackpermute2Int :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int+parBackpermute2Int vis offs =+    parGenerateUnboxedInline+        (max 0 (VU.length offs - 1))+        (VU.unsafeIndex vis . VU.unsafeIndex offs)+{-# NOINLINE parBackpermute2Int #-}++{- | Parallel boxed gather. The read side uses 'VB.unsafeIndexM' so the array+slot is fetched eagerly (element pointers are shared, elements themselves stay+un-forced, exactly as the sequential 'VB.generate' gather). Bit-identical+element values; sequential below 'parThreshold'.+-}+parBackpermuteBoxed :: VB.Vector a -> VU.Vector Int -> VB.Vector a+parBackpermuteBoxed v ix+    | not (shouldParallelize parThreshold n) =+        VB.generate n ((v `VB.unsafeIndex`) . (ix `VU.unsafeIndex`))+    | otherwise = unsafePerformIO $ do+        mv <- VBM.unsafeNew n+        parallelChunks_ parThreshold n $ \ !lo !hi ->+            let go !i+                    | i >= hi = pure ()+                    | otherwise = do+                        x <- VB.unsafeIndexM v (VU.unsafeIndex ix i)+                        VBM.unsafeWrite mv i x+                        go (i + 1)+             in go lo+        VB.unsafeFreeze mv+  where+    !n = VU.length ix+{-# NOINLINE parBackpermuteBoxed #-}++{- | Clamp sentinel (negative) indices to 0 in one closure-free parallel pass.+The clamped rows read row 0's value; callers mask them via the sentinel bitmap.+-}+parClampNonNeg :: VU.Vector Int -> VU.Vector Int+parClampNonNeg ix =+    parGenerateUnboxedInline+        (VU.length ix)+        (\i -> let !x = VU.unsafeIndex ix i in max x 0)+{-# NOINLINE parClampNonNeg #-}++{- | Validity bitmap from sentinel indices (bit @i@ valid iff @ix!i >= 0@),+built one byte (8 rows) per element in parallel.+-}+parBitmapNonNeg :: VU.Vector Int -> Bitmap+parBitmapNonNeg ix =+    let !n = VU.length ix+        !nBytes = (n + 7) `shiftR` 3+     in parGenerateUnboxedInline nBytes $ \b ->+            let !base = b `shiftL` 3+                go !acc !bit+                    | bit >= 8 = acc+                    | otherwise =+                        let !idx = base + bit+                            !acc' =+                                if idx < n && VU.unsafeIndex ix idx >= 0+                                    then setBit acc bit+                                    else acc+                         in go acc' (bit + 1)+             in go (0 :: Word8) 0+{-# NOINLINE parBitmapNonNeg #-}++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+        )+        (parBackpermuteBoxed column indexes)+atIndicesStable indexes (UnboxedColumn bm (column :: VU.Vector a)) =+    UnboxedColumn+        ( fmap+            ( \bm0 ->+                buildBitmapFromValid $+                    VU.map (\i -> if bitmapTestBit bm0 i then 1 else 0) indexes+            )+            bm+        )+        -- Int/Double hit the closure-free monomorphic kernels; anything else+        -- keeps the generic (per-element closure call) path.+        ( case testEquality (typeRep @a) (typeRep @Int) of+            Just Refl -> parBackpermuteInt column indexes+            Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+                Just Refl -> parBackpermuteDouble column indexes+                Nothing -> parBackpermuteUnboxed column indexes+        )+atIndicesStable indexes (MergedColumn a b) =+    MergedColumn (atIndicesStable indexes a) (atIndicesStable indexes b)+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 c@(MergedColumn _ _) =+    gatherWithSentinel indices (materializeMerged c)+gatherWithSentinel indices col =+    let !n = VU.length indices+        !newBm = parBitmapNonNeg indices+        withBm srcBm = case srcBm of+            Nothing -> Just newBm+            Just sb -> Just (andBitmaps newBm (gatherSrcBm sb))+        -- Sequential fallback: gather an existing source bitmap through the+        -- raw indices (negative-guarded). Only runs when the source side is+        -- itself nullable, which join build sides normally are not.+        gatherSrcBm sb =+            buildBitmapFromValid $ VU.generate n $ \i ->+                let idx = VU.unsafeIndex indices i+                 in if idx >= 0 && bitmapTestBit sb idx then 1 else 0+     in case col of+            -- packedGather composes -1 sentinels into the selector natively+            -- (and shares the byte buffers), so text takes the raw indices.+            PackedText srcBm p -> PackedText (withBm srcBm) (packedGather indices p)+            BoxedColumn srcBm v+                -- An empty source means every index is a sentinel; clamping+                -- would read row 0 of an empty vector.+                | VB.null v -> BoxedColumn (withBm srcBm) (allNullBoxed n v)+                | otherwise ->+                    BoxedColumn+                        (withBm srcBm)+                        (parBackpermuteBoxed v (parClampNonNeg indices))+            UnboxedColumn srcBm v+                | VU.null v -> UnboxedColumn (withBm srcBm) (allNullUnboxed n v)+                | otherwise ->+                    -- Reuse atIndicesStable's Int/Double monomorphic kernel+                    -- dispatch for the payload; the bitmap is replaced below.+                    case atIndicesStable (parClampNonNeg indices) (UnboxedColumn Nothing v) of+                        UnboxedColumn _ dat -> UnboxedColumn (withBm srcBm) dat+                        other -> other+{-# INLINE gatherWithSentinel #-}++{- | An @n@-row payload for a gather whose source is empty: every index is a+sentinel, so the sentinel bitmap masks every row and no element is ever read.+The source vector is passed only to fix the element type.+-}+allNullBoxed :: Int -> VB.Vector a -> VB.Vector a+allNullBoxed n _ = VB.replicate n (error "gatherWithSentinel: null row forced")++-- | 'allNullBoxed' for unboxed payloads; the buffer is left uninitialised.+allNullUnboxed :: (VU.Unbox a) => Int -> VU.Vector a -> VU.Vector a+allNullUnboxed n _ = runST (VUM.new n >>= VU.unsafeFreeze)++-- | 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)+    c@(MergedColumn _ _) -> findIndices predicate (materializeMerged 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)+    c@(MergedColumn _ _) -> ifoldrColumn f acc (materializeMerged 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)+    c@(MergedColumn _ _) -> foldlColumn f acc (materializeMerged 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)+    c@(MergedColumn _ _) -> foldl1Column f (materializeMerged 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+        MergedColumn _ _ -> foldl1DirectGroups f (materializeMerged 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+        MergedColumn _ _ ->+            foldLinearGroups f seed (materializeMerged 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)+    c@(MergedColumn _ _) -> headColumn (mergedHead 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@(MergedColumn _ _) r = zipColumns (materializeMerged l) r+zipColumns l r@(MergedColumn _ _) = zipColumns l (materializeMerged r)+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 #-}++-- | 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 ->+                        let !n = min (VU.length column) (VU.length other)+                         in UnboxedColumn Nothing $+                                parGenerateUnboxed n $ \i ->+                                    f (VU.unsafeIndex column i) (VU.unsafeIndex other i)+                    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@(MergedColumn _ _) = ensureOptional (materializeMerged c)+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@(MergedColumn a b)+    | n <= min (columnLength a) (columnLength b) = c+    | otherwise = expandColumn n (materializeMerged c)+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 (VU.enumFromN (VG.length col) extra))+                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 (VU.enumFromN (VG.length col) extra))+                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@(MergedColumn a b)+    | n <= min (columnLength a) (columnLength b) = c+    | otherwise = leftExpandColumn n (materializeMerged c)+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 (VU.enumFromN 0 extra))+                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 (VU.enumFromN 0 extra))+                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.+-}+mappendColumns :: Column -> Column -> Either DataFrameException Column+mappendColumns left right = case (left, right) of+    (MergedColumn _ _, _) -> mappendColumns (materializeMerged left) right+    (_, MergedColumn _ _) -> mappendColumns left (materializeMerged right)+    (PackedText _ _, _) -> mappendColumns (materializePacked left) right+    (_, PackedText _ _) -> mappendColumns 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 "mappendColumns"+                        , errorColumnName = Nothing+                        }+                    )++{- | Like 'mappendColumns' 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 isMergedColumn all' =+        concatManyColumns (map materializeMerged 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))+    MergedColumn _ _ -> concatManyColumns (map materializeMerged (c0 : cs))++mappendColumnsEither :: Column -> Column -> Column+mappendColumnsEither l@(MergedColumn _ _) r =+    mappendColumnsEither (materializeMerged l) r+mappendColumnsEither l r@(MergedColumn _ _) =+    mappendColumnsEither l (materializeMerged r)+mappendColumnsEither l@(PackedText _ _) r = mappendColumnsEither (materializePacked l) r+mappendColumnsEither l r@(PackedText _ _) = mappendColumnsEither l (materializePacked r)+mappendColumnsEither (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+mappendColumnsEither (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+mappendColumnsEither (BoxedColumn _ left) (UnboxedColumn _ right) =+    BoxedColumn Nothing $ fmap Left left <> fmap Right (VG.convert right)+mappendColumnsEither (UnboxedColumn _ left) (BoxedColumn _ right) =+    BoxedColumn Nothing $ fmap Left (VG.convert left) <> fmap Right right++-------------------------------------------------------------------------------+-- Fused multi-column gather+-------------------------------------------------------------------------------++{- | Gather ONE in-bounds index vector through several columns in a single+parallel pass. Result columns are identical to @map (atIndicesStable ixs)@;+columns whose shape the fused kernel does not cover (bitmapped, boxed, merged,+unusual element types) fall back to per-column 'atIndicesStable'. All fused+outputs are backed by one shared deferred computation: forcing ANY of them+runs the single pass that fills ALL of them. The pass reads the index vector+once per block instead of once per column, and every iteration of a block+issues each column's independent random load back-to-back, so their cache/TLB+misses overlap instead of forming one latency chain per column (the aggregate+key materialization of a 1e8-group result was 6 sequential latency-bound+passes without this).+-}+atIndicesStableMulti :: VU.Vector Int -> [Column] -> [Column]+atIndicesStableMulti ixs cols =+    let specs = map mgClassify cols+        nFused = length [() | Just _ <- specs]+     in if nFused < 2+            then map (atIndicesStable ixs) cols+            else+                let fused = multiGatherRun ixs (catMaybes specs)+                    go [] _ = []+                    go (Nothing : ss) !k = atIndicesStable ixs (cols !! k) : go ss (k + 1)+                    go (Just _ : ss) !k =+                        let !r = mgRank k+                         in (fused VB.! r) : go ss (k + 1)+                    -- fused-output rank of column position k.+                    mgRank k = length [() | Just _ <- take k specs]+                 in go specs 0++-- | One fusable source column shape (all bitmap-free).+data MGSpec+    = -- | Clean unboxed 'Int' payload.+      MGInt !(VU.Vector Int)+    | -- | Clean unboxed 'Double' payload.+      MGDouble !(VU.Vector Double)+    | {- | Packed text with an 'Int32' selector (base row count necessarily+      fits 'Int32'); carries the payload for rebuilding and the canon flag.+      -}+      MGSel32 !PackedTextData !(VU.Vector Int32)+    | -- | Packed text with an 'Int' selector and an 'Int32'-sized base.+      MGSel64To32 !PackedTextData !(VU.Vector Int)+    | -- | Packed text with an 'Int' selector and a wide base.+      MGSel64To64 !PackedTextData !(VU.Vector Int)++mgClassify :: Column -> Maybe MGSpec+mgClassify (UnboxedColumn Nothing (v :: VU.Vector a)) =+    case testEquality (typeRep @a) (typeRep @Int) of+        Just Refl -> Just (MGInt v)+        Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+            Just Refl -> Just (MGDouble v)+            Nothing -> Nothing+mgClassify (PackedText Nothing p) = case ptSel p of+    Just (Sel32 s) -> Just (MGSel32 p s)+    Just (Sel64 s)+        | offCount (ptOffsets p) - 1 <= mgInt32Max -> Just (MGSel64To32 p s)+        | otherwise -> Just (MGSel64To64 p s)+    Nothing -> Nothing+mgClassify _ = Nothing++mgInt32Max :: Int+mgInt32Max = fromIntegral (maxBound :: Int32)++{- | Run the fused pass over the fusable specs; element @r@ of the result is+spec @r@'s gathered column. Deferred as one shared thunk (see+'atIndicesStableMulti'). Pure w.r.t. its immutable inputs, so the+'unsafePerformIO' is safe.+-}+multiGatherRun :: VU.Vector Int -> [MGSpec] -> VB.Vector Column+multiGatherRun ixs specs = unsafePerformIO $ do+    let !n = VU.length ixs+    opened <- mapM (mgOpen n ixs) specs+    let fills = map fst opened+        !block = 4096+        worker !lo !hi+            | lo >= hi = pure ()+            | otherwise = do+                let !e = min hi (lo + block)+                mapM_ (\fill -> fill lo e) fills+                worker e hi+    parallelChunks_ parThreshold n worker+    VB.fromList <$> mapM snd opened+{-# NOINLINE multiGatherRun #-}++{- | Allocate a spec's destination; return its block-fill action and its+finalizer. Selector gathers reproduce 'packedGather' exactly (composition+clamps against the source selector length and base row count); unboxed gathers+reproduce the unclamped 'parBackpermuteUnboxed'.+-}+mgOpen :: Int -> VU.Vector Int -> MGSpec -> IO (Int -> Int -> IO (), IO Column)+mgOpen n ixs spec = case spec of+    MGInt v -> do+        mv <- VUM.unsafeNew n+        pure+            ( mgFillInt ixs v mv+            , UnboxedColumn Nothing <$> VU.unsafeFreeze mv+            )+    MGDouble v -> do+        mv <- VUM.unsafeNew n+        pure+            ( mgFillDouble ixs v mv+            , UnboxedColumn Nothing <$> VU.unsafeFreeze mv+            )+    MGSel32 p s -> do+        mv <- VUM.unsafeNew n+        pure+            ( mgFillSel32 ixs s (offCount (ptOffsets p) - 1) mv+            , (\out -> PackedText Nothing p{ptSel = Just (Sel32 out)}) <$> VU.unsafeFreeze mv+            )+    MGSel64To32 p s -> do+        mv <- VUM.unsafeNew n+        pure+            ( mgFillSel64To32 ixs s (offCount (ptOffsets p) - 1) mv+            , (\out -> PackedText Nothing p{ptSel = Just (Sel32 out)}) <$> VU.unsafeFreeze mv+            )+    MGSel64To64 p s -> do+        mv <- VUM.unsafeNew n+        pure+            ( mgFillSel64To64 ixs s (offCount (ptOffsets p) - 1) mv+            , (\out -> PackedText Nothing p{ptSel = Just (Sel64 out)}) <$> VU.unsafeFreeze mv+            )++mgFillInt ::+    VU.Vector Int -> VU.Vector Int -> VUM.IOVector Int -> Int -> Int -> IO ()+mgFillInt ixs v dst lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            VUM.unsafeWrite dst i (VU.unsafeIndex v (VU.unsafeIndex ixs i))+            go (i + 1)+{-# NOINLINE mgFillInt #-}++mgFillDouble ::+    VU.Vector Int -> VU.Vector Double -> VUM.IOVector Double -> Int -> Int -> IO ()+mgFillDouble ixs v dst lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            VUM.unsafeWrite dst i (VU.unsafeIndex v (VU.unsafeIndex ixs i))+            go (i + 1)+{-# NOINLINE mgFillDouble #-}++mgFillSel32 ::+    VU.Vector Int ->+    VU.Vector Int32 ->+    Int ->+    VUM.IOVector Int32 ->+    Int ->+    Int ->+    IO ()+mgFillSel32 ixs s !base dst lo hi = go lo+  where+    !sn = VU.length s+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !j = VU.unsafeIndex ixs i+                !out =+                    if j >= 0 && j < sn+                        then+                            let !r = fromIntegral (VU.unsafeIndex s j) :: Int+                             in if r >= 0 && r < base then fromIntegral r else -1+                        else -1+            VUM.unsafeWrite dst i out+            go (i + 1)+{-# NOINLINE mgFillSel32 #-}++mgFillSel64To32 ::+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VUM.IOVector Int32 ->+    Int ->+    Int ->+    IO ()+mgFillSel64To32 ixs s !base dst lo hi = go lo+  where+    !sn = VU.length s+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !j = VU.unsafeIndex ixs i+                !out =+                    if j >= 0 && j < sn+                        then+                            let !r = VU.unsafeIndex s j+                             in if r >= 0 && r < base then fromIntegral r else -1+                        else -1+            VUM.unsafeWrite dst i out+            go (i + 1)+{-# NOINLINE mgFillSel64To32 #-}++mgFillSel64To64 ::+    VU.Vector Int -> VU.Vector Int -> Int -> VUM.IOVector Int -> Int -> Int -> IO ()+mgFillSel64To64 ixs s !base dst lo hi = go lo+  where+    !sn = VU.length s+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !j = VU.unsafeIndex ixs i+                !out =+                    if j >= 0 && j < sn+                        then+                            let !r = VU.unsafeIndex s j+                             in if r >= 0 && r < base then r else -1+                        else -1+            VUM.unsafeWrite dst i out+            go (i + 1)+{-# NOINLINE mgFillSel64To64 #-}
+ src-internal/DataFrame/Internal/Column/Properties.hs view
@@ -0,0 +1,153 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}++{- |+Predicates and introspection over a 'Column': representation tests, null/type+queries, and the human-readable type descriptions used in error messages.+-}+module DataFrame.Internal.Column.Properties where++import qualified Data.Text as T+import qualified Data.Vector as VB+import qualified Data.Vector.Unboxed as VU++import Data.Kind (Type)+import Data.Maybe (isJust)+import Data.Type.Equality (TestEquality (..))+import DataFrame.Internal.Column.Base+import DataFrame.Internal.Column.Bitmap+import DataFrame.Internal.Column.Types+import DataFrame.Internal.Data.PackedText (packedLength)+import Type.Reflection (+    TypeRep,+    Typeable,+    eqTypeRep,+    typeRep,+    pattern App,+    type (:~:) (Refl),+    type (:~~:) (HRefl),+ )++-- | Whether a column is a 'PackedText'.+isPackedText :: Column -> Bool+isPackedText (PackedText _ _) = True+isPackedText _ = False+{-# INLINE isPackedText #-}++-- | Whether a column is a 'MergedColumn'.+isMergedColumn :: Column -> Bool+isMergedColumn (MergedColumn _ _) = True+isMergedColumn _ = False+{-# INLINE isMergedColumn #-}++-- | 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) =+    not (VB.null col) && popCountUpTo (VB.length col) bm == 0+allMissing (UnboxedColumn (Just bm) col) =+    not (VU.null col) && popCountUpTo (VU.length col) bm == 0+allMissing (PackedText (Just bm) p) =+    packedLength p > 0 && popCountUpTo (packedLength p) bm == 0+allMissing _ = False++-- | Checks if a column contains numeric values.+isNumeric :: Column -> Bool+isNumeric c@(MergedColumn _ _) = isNumeric (mergedHead c)+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)+    c@(MergedColumn _ _) -> hasElemType @a (mergedHead c)+  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"+    MergedColumn _ _ -> columnVersionString (mergedHead column)++{- | 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+    MergedColumn _ _ -> columnTypeString (mergedHead column)+  where+    showMaybeType :: forall a. (Typeable a) => String+    showMaybeType =+        let s = show (typeRep @a)+         in "Maybe " ++ if ' ' `elem` s then "(" ++ s ++ ")" else s++-- | 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++-- | O(n) Gets the number of non-null elements in the column.+numElements :: Column -> Int+numElements (MergedColumn a b) = min (columnLength a) (columnLength b)+numElements (BoxedColumn Nothing xs) = VB.length xs+numElements (BoxedColumn (Just bm) xs) = popCountUpTo (VB.length xs) bm+numElements (UnboxedColumn Nothing xs) = VU.length xs+numElements (UnboxedColumn (Just bm) xs) = popCountUpTo (VU.length xs) bm+numElements (PackedText Nothing p) = packedLength p+numElements (PackedText (Just bm) p) = popCountUpTo (packedLength p) bm+{-# INLINE numElements #-}
+ src-internal/DataFrame/Internal/Column/Types.hs view
@@ -0,0 +1,161 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE DeriveTraversable #-}+{-# LANGUAGE ExistentialQuantification #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}++module DataFrame.Internal.Column.Types where++import Data.Int (Int16, Int32, Int64, Int8)+import Data.Kind (Constraint, Type)+import Data.Typeable (Typeable)+import qualified Data.Vector.Unboxed as VU+import Data.Word (Word16, Word32, Word64, Word8)++type Columnable' a = (Typeable a, Show a, Eq a)++{- | Inline replacement for @Data.These.These@ to keep @dataframe-core@ free+of the @these@ package dependency. Only the three constructors and the+derived classes are used internally.+-}+data These a b = This a | That b | These a b+    deriving (Eq, Ord, Show, Read, Functor, Foldable, Traversable)++{- | A type with column representations used to select the+"right" representation when specializing the `toColumn` function.+-}+data Rep+    = RBoxed+    | RUnboxed+    | RNullableBoxed++-- | Type-level if statement.+type family If (cond :: Bool) (yes :: k) (no :: k) :: k where+    If 'True yes _ = yes+    If 'False _ no = no++-- | All unboxable types (according to the `vector` package).+type family Unboxable (a :: Type) :: Bool where+    Unboxable Int = 'True+    Unboxable Int8 = 'True+    Unboxable Int16 = 'True+    Unboxable Int32 = 'True+    Unboxable Int64 = 'True+    Unboxable Word = 'True+    Unboxable Word8 = 'True+    Unboxable Word16 = 'True+    Unboxable Word32 = 'True+    Unboxable Word64 = 'True+    Unboxable Char = 'True+    Unboxable Bool = 'True+    Unboxable Double = 'True+    Unboxable Float = 'True+    Unboxable _ = 'False++type family Numeric (a :: Type) :: Bool where+    Numeric Integer = 'True+    Numeric Int = 'True+    Numeric Int8 = 'True+    Numeric Int16 = 'True+    Numeric Int32 = 'True+    Numeric Int64 = 'True+    Numeric Word = 'True+    Numeric Word8 = 'True+    Numeric Word16 = 'True+    Numeric Word32 = 'True+    Numeric Word64 = 'True+    Numeric Double = 'True+    Numeric Float = 'True+    Numeric _ = 'False++-- | Compute the column representation tag for any 'a'.+type family KindOf a :: Rep where+    KindOf (Maybe a) = 'RNullableBoxed+    KindOf a = If (Unboxable a) 'RUnboxed 'RBoxed++-- | Type-level boolean for constraint/type comparison.+data SBool (b :: Bool) where+    STrue :: SBool 'True+    SFalse :: SBool 'False++-- | The runtime witness for our type-level branching.+class SBoolI (b :: Bool) where+    sbool :: SBool b++instance SBoolI 'True where sbool = STrue+instance SBoolI 'False where sbool = SFalse++-- | Runtime witness for whether @a@ is unboxable.+sUnbox :: forall a. (SBoolI (Unboxable a)) => SBool (Unboxable a)+sUnbox = sbool @(Unboxable a)++sNumeric :: forall a. (SBoolI (Numeric a)) => SBool (Numeric a)+sNumeric = sbool @(Numeric a)++type family When (flag :: Bool) (c :: Constraint) :: Constraint where+    When 'True c = c+    When 'False c = ()++type UnboxIf a = When (Unboxable a) (VU.Unbox a)++type family IntegralTypes (a :: Type) :: Bool where+    IntegralTypes Integer = 'True+    IntegralTypes Int = 'True+    IntegralTypes Int8 = 'True+    IntegralTypes Int16 = 'True+    IntegralTypes Int32 = 'True+    IntegralTypes Int64 = 'True+    IntegralTypes Word = 'True+    IntegralTypes Word8 = 'True+    IntegralTypes Word16 = 'True+    IntegralTypes Word32 = 'True+    IntegralTypes Word64 = 'True+    IntegralTypes _ = 'False++sIntegral :: forall a. (SBoolI (IntegralTypes a)) => SBool (IntegralTypes a)+sIntegral = sbool @(IntegralTypes a)++type IntegralIf a = When (IntegralTypes a) (Integral a)++type family FloatingTypes (a :: Type) :: Bool where+    FloatingTypes Float = 'True+    FloatingTypes Double = 'True+    FloatingTypes _ = 'False++sFloating :: forall a. (SBoolI (FloatingTypes a)) => SBool (FloatingTypes a)+sFloating = sbool @(FloatingTypes a)++type FloatingIf a = When (FloatingTypes a) (Real a, Fractional a)++{- | Numeric type promotion: resolves the common type for mixed arithmetic.+Double dominates over Float/Int; Float dominates over Int; same types stay unchanged.+-}+type family Promote (a :: Type) (b :: Type) :: Type where+    Promote a a = a+    Promote Double _ = Double+    Promote _ Double = Double+    Promote Float _ = Float+    Promote _ Float = Float+    Promote Int64 _ = Int64+    Promote _ Int64 = Int64+    Promote Int32 _ = Int32+    Promote _ Int32 = Int32+    Promote a _ = a++{- | Like 'Promote', but integral × integral → Double for use with './' .+Double\/Float still dominate; any two integral types (same or mixed) become Double.+-}+type family PromoteDiv (a :: Type) (b :: Type) :: Type where+    PromoteDiv Double _ = Double+    PromoteDiv _ Double = Double+    PromoteDiv Float _ = Float+    PromoteDiv _ Float = Float+    PromoteDiv _ _ = Double
+ src-internal/DataFrame/Internal/Control/Concurrent.hs view
@@ -0,0 +1,177 @@+{-# LANGUAGE BangPatterns #-}++-- | Shared concurrency primitives for the dataframe packages.+module DataFrame.Internal.Control.Concurrent (+    -- * Capabilities+    capabilities,+    capabilitiesIO,+    shouldParallelize,+    parThreshold,++    -- * Chunk planning (pure)+    splitChunkRange,+    chunksFor,+    boundsChunks,++    -- * Thread fan-out+    forkJoin,+    forkJoin_,++    -- * Chunked fan-out (per-chunk callbacks only)+    parallelChunks,+    parallelChunks_,+    parallelBounds_,++    -- * Work-stealing pools+    pooledIndices,+    pooledRun,+) where++import Control.Concurrent (forkFinally, getNumCapabilities)+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)+import Control.Exception (ErrorCall (..), SomeException, throwIO)+import Control.Monad (when)+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 System.IO.Unsafe (unsafePerformIO)++capabilities :: Int+capabilities = unsafePerformIO getNumCapabilities+{-# NOINLINE capabilities #-}++capabilitiesIO :: IO Int+capabilitiesIO = getNumCapabilities+{-# INLINE capabilitiesIO #-}++shouldParallelize :: Int -> Int -> Bool+shouldParallelize threshold n = n >= threshold && capabilities > 1+{-# INLINE shouldParallelize #-}++{- | Row count below which a fan-out does not pay for itself, for the fixed-cost+per-row loops (grouping, the aggregation kernels). Pair it with+'shouldParallelize'. Kernels with a materially different per-row cost —+the join probe, the radix sort — set their own thresholds.+-}+parThreshold :: Int+parThreshold = 200000++splitChunkRange :: Int -> Int -> [(Int, Int)]+splitChunkRange k n+    | n <= 0 = []+    | otherwise =+        [ (lo, lo + len)+        | w <- [0 .. k' - 1]+        , let lo = w * q + min w r+        , let len = q + fromEnum (w < r)+        , len > 0+        ]+  where+    -- A non-positive width is a caller bug; clamp to one chunk rather than+    -- return [] and silently skip the rows.+    !k' = max 1 k+    (!q, !r) = n `quotRem` k'++chunksFor :: Int -> Int -> [(Int, Int)]+chunksFor !threshold !n+    | not (shouldParallelize threshold n) = [(0, n)]+    | otherwise = splitChunkRange capabilities n+{-# INLINE chunksFor #-}++{- | Adjacent pairs of a precomputed bounds vector of length @caps + 1@. Empty+ranges are NOT dropped: the bounds are the caller's and the slot count is+often meaningful. Endpoints are forced here, on the spawning thread, so a+worker never starts by evaluating an index thunk that retains @bs@.+-}+boundsChunks :: Int -> VU.Vector Int -> [(Int, Int)]+boundsChunks caps bs =+    [ (lo, hi)+    | w <- [0 .. caps - 1]+    , let !lo = VU.unsafeIndex bs w+    , let !hi = VU.unsafeIndex bs (w + 1)+    ]+{-# INLINE boundsChunks #-}++rethrow :: Either SomeException a -> IO a+rethrow = either throwIO pure+{-# INLINE rethrow #-}++forkJoin :: [IO a] -> IO [a]+forkJoin [] = pure []+forkJoin [act] = fmap (: []) act+forkJoin actions = do+    vars <- mapM spawn actions+    results <- mapM takeMVar vars+    mapM rethrow results+  where+    spawn act = do+        var <- newEmptyMVar+        _ <- forkFinally act (putMVar var)+        pure var+{-# INLINEABLE forkJoin #-}++forkJoin_ :: [IO ()] -> IO ()+forkJoin_ [] = pure ()+forkJoin_ [act] = act+forkJoin_ actions = do+    vars <- mapM spawn actions+    results <- mapM takeMVar vars+    mapM_ rethrow results+  where+    spawn act = do+        var <- newEmptyMVar+        _ <- forkFinally act (putMVar var)+        pure var+{-# INLINEABLE forkJoin_ #-}++parallelChunks :: Int -> Int -> (Int -> Int -> IO a) -> IO [a]+parallelChunks threshold n body =+    forkJoin [body lo hi | (!lo, !hi) <- chunksFor threshold n]+{-# NOINLINE parallelChunks #-} -- INLINE worsens performance here.++-- | 'parallelChunks' for chunk bodies run only for their effects.+parallelChunks_ :: Int -> Int -> (Int -> Int -> IO ()) -> IO ()+parallelChunks_ threshold n body =+    forkJoin_ [body lo hi | (!lo, !hi) <- chunksFor threshold n]+{-# INLINE parallelChunks_ #-}++parallelBounds_ :: Int -> VU.Vector Int -> (Int -> Int -> IO ()) -> IO ()+parallelBounds_ caps bs body =+    forkJoin_ [body lo hi | (!lo, !hi) <- boundsChunks caps bs]+{-# NOINLINE parallelBounds_ #-}++pooledIndices :: Int -> Int -> (Int -> IO ()) -> IO ()+pooledIndices width count body+    | count <= 0 = pure ()+    | width <= 1 = mapM_ body [0 .. count - 1]+    | otherwise = do+        next <- newIORef 0+        let worker = do+                i <- atomicModifyIORef' next (\j -> (j + 1, j))+                when (i < count) (body i >> worker)+        forkJoin_ (replicate (min width count) worker)+{-# NOINLINE pooledIndices #-}++pooledRun :: Int -> [IO a] -> IO [a]+pooledRun width actions+    | width >= n = forkJoin actions+    | otherwise = do+        next <- newIORef 0+        out <- VM.unsafeNew n+        acts <- VM.unsafeNew n+        sequence_ [VM.unsafeWrite acts i a | (i, a) <- zip [0 ..] actions]+        let worker = do+                i <- atomicModifyIORef' next (\j -> (j + 1, j))+                when (i < n) $ do+                    act <- VM.unsafeRead acts i+                    VM.unsafeWrite acts i consumed+                    r <- act+                    VM.write out i r+                    worker+        forkJoin_ (replicate width worker)+        V.toList <$> V.freeze out+  where+    n = length actions+    consumed = throwIO (ErrorCall "pooledRun: slot already consumed")+{-# INLINEABLE pooledRun #-}
+ src-internal/DataFrame/Internal/Data/HashTable.hs view
@@ -0,0 +1,108 @@+{-# 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.Data.HashTable (+    HashTable (..),+    newHashTable,+    htInsert,+    nextPow2Above,+) where++import Control.Monad.Primitive (PrimMonad, PrimState)+import Data.Bits (popCount, unsafeShiftR, (.&.))+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.Word (Word64)++{- | 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 (homeSlot mask hash)+  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 #-}++{- | Home slot of a hash: the top @log2 cap@ bits of a Fibonacci multiply. The+row hash's final FxHash step is a multiply, which leaves its LOW bits poorly+diffused — raw-text keys cluster into contiguous linear-probe pileups when+slotted by @hash .&. mask@, so the home slot must come from the top bits.+-}+homeSlot :: Int -> Int -> Int+homeSlot !mask !hash =+    fromIntegral+        ( (fromIntegral hash * (0x9E3779B97F4A7C15 :: Word64))+            `unsafeShiftR` (64 - popCount mask)+        )+{-# INLINE homeSlot #-}
+ src-internal/DataFrame/Internal/Data/PackedText.hs view
@@ -0,0 +1,361 @@+{-# LANGUAGE BangPatterns #-}++{- | Packed-text payload + byte-slice primitives. A 'PackedTextData' shares one+UTF-8 byte buffer across all rows of a string column, with @n+1@ row offsets, so+no per-row 'Data.Text.Text' header is materialized until decode is demanded.+Offsets and selection vectors are stored 'Int32' whenever their values fit+(Arrow-style), halving the per-row footprint of large string columns.+-}+module DataFrame.Internal.Data.PackedText (+    PackedTextData (..),+    PackedOffsets (..),+    PackedSel (..),+    offAt,+    offCount,+    selAt,+    selLength,+    mkPackedContiguous,+    mkPackedContiguous32,+    mkOffsets,+    mkSel,+    packedGather,+    packedTake,+    packedRowOffsets,+    packedLength,+    packedSlice,+    packedIndexText,+    sliceEqBytes,+    sliceCmpBytes,+) 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 Data.Int (Int32)+import Data.Ord (comparing)+import Data.Text.Internal (Text (Text))+import DataFrame.Internal.Control.Concurrent (+    parThreshold,+    parallelChunks_,+    shouldParallelize,+ )+import DataFrame.Internal.Data.PackedText.Utf8 (+    isValidUtf8Slice,+    lenientDecodeSlice,+ )+import System.IO.Unsafe (unsafePerformIO)++{- | Row byte-offsets, physically 'Int32' when every value fits (total buffer+bytes < 2^31) and 'Int' otherwise. Values are non-negative byte positions.+-}+data PackedOffsets+    = Offs32 {-# UNPACK #-} !(VU.Vector Int32)+    | Offs64 {-# UNPACK #-} !(VU.Vector Int)++-- | Offset at index @i@, widened to 'Int'.+offAt :: PackedOffsets -> Int -> Int+offAt (Offs32 v) i = fromIntegral (VU.unsafeIndex v i)+offAt (Offs64 v) i = VU.unsafeIndex v i+{-# INLINE offAt #-}++-- | Number of offset entries (row count + 1).+offCount :: PackedOffsets -> Int+offCount (Offs32 v) = VU.length v+offCount (Offs64 v) = VU.length v+{-# INLINE offCount #-}++{- | A selection layer mapping logical rows to base rows; @-1@ marks an+invalid/null row. 'Int32' when the base row count fits.+-}+data PackedSel+    = Sel32 {-# UNPACK #-} !(VU.Vector Int32)+    | Sel64 {-# UNPACK #-} !(VU.Vector Int)++-- | Base row for logical row @i@ (may be @-1@).+selAt :: PackedSel -> Int -> Int+selAt (Sel32 v) i = fromIntegral (VU.unsafeIndex v i)+selAt (Sel64 v) i = VU.unsafeIndex v i+{-# INLINE selAt #-}++selLength :: PackedSel -> Int+selLength (Sel32 v) = VU.length v+selLength (Sel64 v) = VU.length v+{-# INLINE selLength #-}++{- | A shared UTF-8 byte buffer plus @n+1@ row offsets (base row @r@ spans bytes+@[offsets!r, offsets!(r+1))@); validity lives in the column's bitmap. @ptSel@ is+an optional selection layer letting a gather/join/sort result share the buffer.++@ptCanonicalSel@ marks a selection that is a canonical dictionary encoding:+equal byte slices always map to the same base row (codes). Set by dictionary+compaction; preserved by gather/take over an already-canonical selection (a+row keeps its code); 'False' for a gather over an unselected base, where two+logical rows can select different but equal-byted base rows. Grouping keys on+codes directly when it holds.+-}+data PackedTextData = PackedTextData+    { ptBytes :: {-# UNPACK #-} !A.Array+    , ptOffsets :: !PackedOffsets+    , ptSel :: !(Maybe PackedSel)+    , ptCanonicalSel :: !Bool+    }++int32Max :: Int+int32Max = fromIntegral (maxBound :: Int32)++-- | Narrow an 'Int' offset vector when the final offset (total bytes) fits.+mkOffsets :: VU.Vector Int -> PackedOffsets+mkOffsets offs+    | not (VU.null offs) && VU.last offs <= int32Max =+        Offs32 (VU.map fromIntegral offs)+    | otherwise = Offs64 offs+{-# INLINE mkOffsets #-}++{- | Narrow an 'Int' base-row vector (@-1@ sentinels allowed) when the base+row count fits in 'Int32'.+-}+mkSel :: Int -> VU.Vector Int -> PackedSel+mkSel base rows+    | base <= int32Max = Sel32 (VU.map fromIntegral rows)+    | otherwise = Sel64 rows+{-# INLINE mkSel #-}++-- | Build a contiguous packed payload (no selection): the freeze-path shape.+mkPackedContiguous :: A.Array -> VU.Vector Int -> PackedTextData+mkPackedContiguous arr offs = PackedTextData arr (mkOffsets offs) Nothing False+{-# INLINE mkPackedContiguous #-}++-- | 'mkPackedContiguous' from offsets already produced at 'Int32' width.+mkPackedContiguous32 :: A.Array -> VU.Vector Int32 -> PackedTextData+mkPackedContiguous32 arr offs = PackedTextData arr (Offs32 offs) Nothing False+{-# INLINE mkPackedContiguous32 #-}++{- | Reindex a packed payload by a selection vector, sharing the byte buffer;+logical row @i@ becomes base row @indices!i@. A negative or out-of-range index+decodes to the empty slice. Composes with an existing selection; canonicality+survives composition (a kept row keeps its code) but not a first selection+over the unselected base.+-}+packedGather :: VU.Vector Int -> PackedTextData -> PackedTextData+packedGather indices (PackedTextData arr offs msel canon) =+    let !base = offCount offs - 1+        canon' = case msel of+            Nothing -> False+            Just _ -> canon+        {- Base row for logical output row i composes any existing selection;+        only the SELECTOR is rebuilt (the byte buffer and offsets are shared+        untouched — a canonical dict column keeps its bytes). Each (source+        width x output width) shape runs a monomorphic closure-free kernel,+        generated directly at its final 'mkSel' width (one parallel pass).+        Same values and widths as the historical closure-driven build. -}+        sel' = case msel of+            Nothing+                | base <= int32Max -> Sel32 (gatherBaseTo32 indices base)+                | otherwise -> Sel64 (gatherBaseTo64 indices base)+            Just (Sel32 s)+                | base <= int32Max -> Sel32 (gatherSel32To32 indices s base)+                | otherwise -> Sel64 (gatherSel32To64 indices s base)+            Just (Sel64 s)+                | base <= int32Max -> Sel32 (gatherSel64To32 indices s base)+                | otherwise -> Sel64 (gatherSel64To64 indices s base)+     in PackedTextData arr offs (Just sel') canon'++-- Monomorphic 'packedGather' selector kernels. All reproduce exactly+-- @clamp r = if r >= 0 && r < base then r else -1@ over the composed pick.++gatherBaseTo32 :: VU.Vector Int -> Int -> VU.Vector Int32+gatherBaseTo32 indices !base =+    parGenSelInline (VU.length indices) $ \i ->+        let !r = VU.unsafeIndex indices i+         in if r >= 0 && r < base then fromIntegral r else -1+{-# NOINLINE gatherBaseTo32 #-}++gatherBaseTo64 :: VU.Vector Int -> Int -> VU.Vector Int+gatherBaseTo64 indices !base =+    parGenSelInline (VU.length indices) $ \i ->+        let !r = VU.unsafeIndex indices i+         in if r >= 0 && r < base then r else -1+{-# NOINLINE gatherBaseTo64 #-}++gatherSel32To32 :: VU.Vector Int -> VU.Vector Int32 -> Int -> VU.Vector Int32+gatherSel32To32 indices s !base =+    let !sn = VU.length s+     in parGenSelInline (VU.length indices) $ \i ->+            let !j = VU.unsafeIndex indices i+             in if j >= 0 && j < sn+                    then+                        let !r = fromIntegral (VU.unsafeIndex s j) :: Int+                         in if r >= 0 && r < base then fromIntegral r else -1+                    else -1+{-# NOINLINE gatherSel32To32 #-}++gatherSel32To64 :: VU.Vector Int -> VU.Vector Int32 -> Int -> VU.Vector Int+gatherSel32To64 indices s !base =+    let !sn = VU.length s+     in parGenSelInline (VU.length indices) $ \i ->+            let !j = VU.unsafeIndex indices i+             in if j >= 0 && j < sn+                    then+                        let !r = fromIntegral (VU.unsafeIndex s j) :: Int+                         in if r >= 0 && r < base then r else -1+                    else -1+{-# NOINLINE gatherSel32To64 #-}++gatherSel64To32 :: VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Int32+gatherSel64To32 indices s !base =+    let !sn = VU.length s+     in parGenSelInline (VU.length indices) $ \i ->+            let !j = VU.unsafeIndex indices i+             in if j >= 0 && j < sn+                    then+                        let !r = VU.unsafeIndex s j+                         in if r >= 0 && r < base then fromIntegral r else -1+                    else -1+{-# NOINLINE gatherSel64To32 #-}++gatherSel64To64 :: VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Int+gatherSel64To64 indices s !base =+    let !sn = VU.length s+     in parGenSelInline (VU.length indices) $ \i ->+            let !j = VU.unsafeIndex indices i+             in if j >= 0 && j < sn+                    then+                        let !r = VU.unsafeIndex s j+                         in if r >= 0 && r < base then r else -1+                    else -1+{-# NOINLINE gatherSel64To64 #-}++{- | Selector generate over 'parallelChunks_': one contiguous index chunk per+capability written into disjoint slices of one buffer — element @i@ depends+only on @f i@, so the result is bit-identical to 'VU.generate' at any @-N@.+Same policy as 'DataFrame.Internal.Column.Operations.parGenerateUnboxed', which+sits above this module.+-}++{- | 'parGenSel' with an INLINE body: each monomorphic NOINLINE kernel above+gets its own copy of the fill loop with the pick function inlined — no unknown+closure call (or boxed result allocation) per element. Bit-identical results;+callers must be NOINLINE so the 'unsafePerformIO' runs once per call.+-}+parGenSelInline :: (VU.Unbox c) => Int -> (Int -> c) -> VU.Vector c+parGenSelInline n f+    | not (shouldParallelize parThreshold n) = VU.generate n f+    | otherwise = unsafePerformIO $ do+        mv <- VUM.unsafeNew n+        parallelChunks_ parThreshold n $ \ !lo !hi ->+            let fill !i+                    | i >= hi = pure ()+                    | otherwise = VUM.unsafeWrite mv i (f i) >> fill (i + 1)+             in fill lo+        VU.unsafeFreeze mv+{-# INLINE parGenSelInline #-}++parGenSel :: (VU.Unbox c) => Int -> (Int -> c) -> VU.Vector c+{-# SPECIALIZE parGenSel :: Int -> (Int -> Int32) -> VU.Vector Int32 #-}+{-# SPECIALIZE parGenSel :: Int -> (Int -> Int) -> VU.Vector Int #-}+parGenSel n f+    | not (shouldParallelize parThreshold n) = VU.generate n f+    | otherwise = unsafePerformIO $ do+        mv <- VUM.unsafeNew n+        parallelChunks_ parThreshold n $ \ !lo !hi ->+            let fill !i+                    | i >= hi = pure ()+                    | otherwise = VUM.unsafeWrite mv i (f i) >> fill (i + 1)+             in fill lo+        VU.unsafeFreeze mv+{-# NOINLINE parGenSel #-}++{- | Take the first @k@ logical rows, sharing the byte buffer via a capped+selection layer. O(k), no byte copy or decode — cheap @take@/display on a+large packed column.+-}+packedTake :: Int -> PackedTextData -> PackedTextData+packedTake k (PackedTextData arr offs msel canon) =+    let !base = offCount offs - 1+        !k' = max 0 k+        (sel', canon') = case msel of+            Just (Sel32 s) -> (Sel32 (VU.take k' s), canon)+            Just (Sel64 s) -> (Sel64 (VU.take k' s), canon)+            Nothing -> (mkSel base (VU.enumFromN 0 (min k' base)), False)+     in PackedTextData arr offs (Just sel') canon'+{-# INLINE packedTake #-}++-- | Map a logical row index to its base row, honoring any selection layer.+baseRow :: PackedTextData -> Int -> Int+baseRow (PackedTextData _ _ Nothing _) i = i+baseRow (PackedTextData _ _ (Just sel) _) i = selAt sel i+{-# INLINE baseRow #-}++-- | Row count: @length sel@ when selected, else @length offsets - 1@.+packedLength :: PackedTextData -> Int+packedLength (PackedTextData _ offs Nothing _) = offCount offs - 1+packedLength (PackedTextData _ _ (Just sel) _) = selLength sel+{-# INLINE packedLength #-}++-- | Raw byte slice for logical row @i@: @(buffer, offset, length)@. The hot accessor.+packedSlice :: PackedTextData -> Int -> (A.Array, Int, Int)+packedSlice p@(PackedTextData arr offs _ _) i =+    let !r = baseRow p i+     in if r < 0+            then (arr, 0, 0)+            else+                let o = offAt offs r in (arr, o, offAt offs (r + 1) - o)+{-# INLINE packedSlice #-}++{- | The shared buffer + contiguous @n+1@ offsets when the payload is the+unselected base; a selected (gathered) payload returns 'Nothing' (its rows are+non-contiguous). Lets contiguous consumers skip the selection indirection.+-}+packedRowOffsets :: PackedTextData -> Maybe (A.Array, PackedOffsets)+packedRowOffsets (PackedTextData arr offs Nothing _) = Just (arr, offs)+packedRowOffsets _ = Nothing+{-# INLINE packedRowOffsets #-}++{- | On-demand single 'Data.Text.Text' for row @i@, using the same+validate-or-lenient decode as the freeze path so output is bit-identical.+-}+packedIndexText :: PackedTextData -> Int -> T.Text+packedIndexText p i =+    let (arr, o, l) = packedSlice p i+     in decodeField arr o l+{-# INLINE packedIndexText #-}++-- Decode one field exactly as the boxed freeze path does per row.+decodeField :: A.Array -> Int -> Int -> T.Text+decodeField arr o l+    | l == 0 = T.empty+    | isValidUtf8Slice arr o l = Text arr o l+    | otherwise = lenientDecodeSlice arr o l+{-# INLINE decodeField #-}++{- | Byte-wise equality of two slices. UTF-8 is injective on valid scalar+sequences and lenient decode is deterministic, so this agrees with+@Text@'s '==' on the decoded values.+-}+sliceEqBytes :: A.Array -> Int -> Int -> A.Array -> Int -> Int -> Bool+sliceEqBytes a ao al b bo bl+    | al /= bl = False+    | otherwise = go 0+  where+    go !k+        | k >= al = True+        | A.unsafeIndex a (ao + k) == A.unsafeIndex b (bo + k) = go (k + 1)+        | otherwise = False+{-# INLINE sliceEqBytes #-}++{- | Unsigned byte-lexicographic comparison (memcmp semantics). For+well-formed UTF-8 this matches 'Data.Text.compare' exactly, since UTF-8+byte order equals codepoint order for all valid scalars.+-}+sliceCmpBytes :: A.Array -> Int -> Int -> A.Array -> Int -> Int -> Ordering+sliceCmpBytes a ao al b bo bl = go 0+  where+    !m = min al bl+    go !k+        | k >= m = compare al bl+        | otherwise = case comparing id (A.unsafeIndex a (ao + k)) (A.unsafeIndex b (bo + k)) of+            EQ -> go (k + 1)+            r -> r+{-# INLINE sliceCmpBytes #-}
+ src-internal/DataFrame/Internal/Data/PackedText/Utf8.hs view
@@ -0,0 +1,95 @@+{-# LANGUAGE BangPatterns #-}++{- | UTF-8 validation and @decodeUtf8Lenient@-parity slice decoding used by+'DataFrame.Internal.Column.Builder' to turn shared byte buffers into 'Text'.+-}+module DataFrame.Internal.Data.PackedText.Utf8 (+    isValidUtf8Slice,+    isUtf8Boundary,+    lenientDecodeSlice,+    sliceTextVector,+) where++import qualified Data.Text as T+import qualified Data.Text.Array as A+import qualified Data.Vector as VB+import qualified Data.Vector.Mutable as VBM+import qualified Data.Vector.Unboxed as VU++import Data.Text.Internal (Text (..))+import Data.Text.Internal.Encoding.Utf8 (+    DecoderResult (..),+    utf8DecodeContinue,+    utf8DecodeStart,+ )+import Data.Text.Internal.Validate (isValidUtf8ByteArray)+import Data.Word (Word8)++-- | Whether @len@ bytes starting at @off@ are well-formed UTF-8.+isValidUtf8Slice :: A.Array -> Int -> Int -> Bool+isValidUtf8Slice = isValidUtf8ByteArray+{-# INLINE isValidUtf8Slice #-}++{- | Whether a byte may start a code point (i.e. is not a continuation+byte). Field slices of a valid buffer are themselves valid iff every+field starts on a boundary.+-}+isUtf8Boundary :: Word8 -> Bool+isUtf8Boundary w = w < 0x80 || w >= 0xC0+{-# INLINE isUtf8Boundary #-}++{- | Decode a byte slice exactly like @decodeUtf8Lenient@: greedy decode at+each position; any byte that cannot begin a complete, valid sequence within+the slice becomes one U+FFFD and decoding resumes at the next byte.+-}+lenientDecodeSlice :: A.Array -> Int -> Int -> T.Text+lenientDecodeSlice arr off len = T.pack (go off)+  where+    !end = off + len+    go !i+        | i >= end = []+        | otherwise = case tryDecode i of+            Just (c, i') -> c : go i'+            Nothing -> '\xFFFD' : go (i + 1)+    tryDecode !i = loop (utf8DecodeStart (A.unsafeIndex arr i)) (i + 1)+      where+        loop (Accept c) !j = Just (c, j)+        loop Reject _ = Nothing+        loop (Incomplete st cp) !j+            | j >= end = Nothing+            | otherwise = loop (utf8DecodeContinue (A.unsafeIndex arr j) st cp) (j + 1)++{- | Slice forced 'Text' values off a shared array; row @i@ spans bytes+@[offs!i, offs!(i+1))@. Fast path validates the whole span once when every field+starts on a code-point boundary; else per-field validation with lenient decode.+-}+sliceTextVector :: A.Array -> VU.Vector Int -> VB.Vector T.Text+sliceTextVector arr offs = VB.create $ do+    mv <- VBM.unsafeNew n+    let fill dec = go 0+          where+            go !i+                | i >= n = pure ()+                | otherwise = do+                    let o = VU.unsafeIndex offs i+                        !t = dec o (VU.unsafeIndex offs (i + 1) - o)+                    VBM.unsafeWrite mv i t+                    go (i + 1)+    if fast then fill mkSlice else fill decodeField+    pure mv+  where+    n = VU.length offs - 1+    base = VU.unsafeIndex offs 0+    used = VU.unsafeIndex offs n+    boundariesOk !i+        | i >= n = True+        | otherwise =+            let o = VU.unsafeIndex offs i+             in (o >= used || isUtf8Boundary (A.unsafeIndex arr o))+                    && boundariesOk (i + 1)+    fast = isValidUtf8Slice arr base (used - base) && boundariesOk 0+    mkSlice o l = if l == 0 then T.empty else Text arr o l+    decodeField o l+        | l == 0 = T.empty+        | isValidUtf8Slice arr o l = Text arr o l+        | otherwise = lenientDecodeSlice arr o l
+ src-internal/DataFrame/Internal/DataFrame.hs view
@@ -0,0 +1,426 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module DataFrame.Internal.DataFrame (+    DataFrame (..),+    forceDataFrame,+    GroupedDataFrame (.., Grouped),+    TruncateConfig (..),+    defaultTruncateConfig,+    ellipsisText,+    toMarkdown,+    toMarkdown',+    asText,+    asTextWith,+    pickColumns,+    insertAt,+    truncateCell,+    empty,+    columnNames,+    insertColumn,+    fromNamedColumns,+    getColumn,+    unsafeGetColumn,+    null,+    toCsv,+    toCsv',+    toSeparated,+    getRowAsText,+    showElement,+    stripJust,+) 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 (+    RenderFormat (..),+    showTable,+ )+import DataFrame.Errors (+    DataFrameException (ColumnsNotFoundException),+ )+import DataFrame.Internal.Column (+    Column (..),+    columnLength,+    columnToTextVec,+    columnTypeString,+    expandColumn,+    forceColumn,+    materializeMerged,+    parBackpermute2Int,+    sliceColumn,+    takeColumn,+ )+import DataFrame.Internal.Column.Bitmap (bitmapTestBit)+import DataFrame.Internal.Data.PackedText (packedIndexText)+import DataFrame.Internal.Expression+import Text.Printf (printf)+import Type.Reflection (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`.++Laziness contract: every field except 'valueIndices' and 'groupRepRows' is+computed eagerly by the grouping paths. 'valueIndices' may be a lazy thunk (the+low-cardinality direct grouping defers the O(n) stable placement pass until a+consumer — grouped median/top-k gathers, set ops, the interpreter's+group-slicing — actually demands the permutation); forcing it always yields the+unique stable counting-sort permutation of 'rowToGroup', so WHAT it evaluates to+is independent of when it is forced. 'groupRepRows' is the per-group+representative row (the first original row of each group, in canonical group+order); aggregation uses it to materialize the key columns without demanding+'valueIndices'. It may also be a thunk; its value always equals+@VU.map (valueIndices !) (VU.init offsets)@.+-}+data GroupedDataFrame = GroupedInternal+    { fullDataframe :: DataFrame+    , groupedColumns :: [T.Text]+    , valueIndices :: VU.Vector Int+    -- ^ Rows sorted by group id (stable); possibly an unevaluated thunk.+    , 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.+    -}+    , groupRepRows :: VU.Vector Int+    -- ^ First original row of each group, length nGroups. See laziness note.+    }++{- | The historical five-field view of 'GroupedDataFrame'. Matching ignores+'groupRepRows'; building derives it lazily from @valueIndices@/@offsets@ (the+thunk only forces them if something actually reads the representative rows).+The direct grouping paths construct 'GroupedInternal' directly instead so the+representative rows never demand the placement pass.+-}+pattern Grouped ::+    DataFrame ->+    [T.Text] ->+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    GroupedDataFrame+pattern Grouped df cols vis offs rtg <- GroupedInternal df cols vis offs rtg _+    where+        Grouped df cols vis offs rtg =+            GroupedInternal+                df+                cols+                vis+                offs+                rtg+                {- Parallel gather (still a deferred thunk; forcing it forces+                vis/offs as before). Indices are grouping-produced and+                in-bounds by construction: offs has nGroups+1 entries and+                offs!g < length vis for every non-empty group. Values are+                identical to the historical @VU.map (vis !) (VU.init offs)@. -}+                (parBackpermute2Int vis offs)++{-# COMPLETE Grouped #-}++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 "" (T.pack . columnTypeString)) survivingCols+        survivingData = map (maybe V.empty columnToTextVec) 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+                    )+     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 (MergedColumn a b) i =+    showElement+        (materializeMerged (MergedColumn (sliceColumn i 1 a) (sliceColumn i 1 b)))+        0+showElement (BoxedColumn bm (c :: V.Vector a)) i = case bm of+    Just b | not (bitmapTestBit b i) -> "null"+    _ -> 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 bm c) i = case bm of+    Just b | not (bitmapTestBit b i) -> "null"+    _ -> 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 "
+ src-internal/DataFrame/Internal/Display/Pretty.hs view
@@ -0,0 +1,129 @@+{- | A minimal Wadler/Leijen-style document combinator and width-aware renderer.+A 'Doc' describes a layout abstractly; 'render' chooses where soft breaks become+newlines to fit a target width. 'Group' lays a region flat when it fits.+-}+module DataFrame.Internal.Display.Pretty (+    Doc,+    text,+    line,+    hardline,+    nest,+    group,+    (<+>),+    hcat,+    punctuate,+    parens,+    parensWhenBroken,+    defaultWidth,+    render,+) where++data Doc+    = Empty+    | Text String+    | Line+    | Cat Doc Doc+    | Nest Int Doc+    | Group Doc+    | Hard+    | Alt Doc Doc++instance Semigroup Doc where+    (<>) = Cat++instance Monoid Doc where+    mempty = Empty++-- | A literal chunk of text. Must not contain newlines (use 'line'/'hardline').+text :: String -> Doc+text = Text++{- | A soft break: a single space when its enclosing 'group' fits the width,+otherwise a newline + current indentation.+-}+line :: Doc+line = Line++-- | A hard break that never flattens; any enclosing 'group' is forced to break.+hardline :: Doc+hardline = Hard++-- | Add @k@ spaces to the indentation applied at line breaks inside @d@.+nest :: Int -> Doc -> Doc+nest = Nest++-- | Lay the document out flat if it fits the remaining width, broken otherwise.+group :: Doc -> Doc+group = Group++-- | Concatenate two documents separated by a single space.+(<+>) :: Doc -> Doc -> Doc+x <+> y = x <> Text " " <> y++infixr 6 <+>++hcat :: [Doc] -> Doc+hcat = mconcat++-- | Append @sep@ after every element but the last.+punctuate :: Doc -> [Doc] -> [Doc]+punctuate _ [] = []+punctuate _ [d] = [d]+punctuate sep (d : ds) = (d <> sep) : punctuate sep ds++parens :: Doc -> Doc+parens d = Text "(" <> d <> Text ")"++{- | Render @d@ bare when it fits flat on the current line, wrapped in parens when+it must break across lines. Keeps operator grouping unambiguous once a+sub-expression wraps, without parenthesis noise on one-line expressions.+-}+parensWhenBroken :: Doc -> Doc+parensWhenBroken d = Group (Alt d (parens d))++defaultWidth :: Int+defaultWidth = 80++data Mode = Flat | Break++-- | Render a document, breaking soft lines so output fits @width@ columns.+render :: Int -> Doc -> String+render width doc = layout 0 [(0, Break, doc)]+  where+    layout :: Int -> [(Int, Mode, Doc)] -> String+    layout _ [] = ""+    layout col ((i, m, d) : rest) = case d of+        Empty -> layout col rest+        Text s -> s ++ layout (col + length s) rest+        Cat x y -> layout col ((i, m, x) : (i, m, y) : rest)+        Nest j x -> layout col ((i + j, m, x) : rest)+        Line -> case m of+            Flat -> ' ' : layout (col + 1) rest+            Break -> '\n' : replicate i ' ' ++ layout i rest+        Hard -> '\n' : replicate i ' ' ++ layout i rest+        Group x ->+            if fits (width - col) ((i, Flat, x) : rest)+                then layout col ((i, Flat, x) : rest)+                else layout col ((i, Break, x) : rest)+        Alt flat broken -> case m of+            Flat -> layout col ((i, Flat, flat) : rest)+            Break -> layout col ((i, Break, broken) : rest)++    fits :: Int -> [(Int, Mode, Doc)] -> Bool+    fits w _ | w < 0 = False+    fits _ [] = True+    fits w ((i, m, d) : rest) = case d of+        Empty -> fits w rest+        Text s -> fits (w - length s) rest+        Cat x y -> fits w ((i, m, x) : (i, m, y) : rest)+        Nest j x -> fits w ((i + j, m, x) : rest)+        Line -> case m of+            Flat -> fits (w - 1) rest+            Break -> True+        Hard -> case m of+            Flat -> False+            Break -> True+        Group x -> fits w ((i, Flat, x) : rest)+        Alt flat broken -> case m of+            Flat -> fits w ((i, Flat, flat) : rest)+            Break -> fits w ((i, Break, broken) : rest)
+ src-internal/DataFrame/Internal/Expression.hs view
@@ -0,0 +1,522 @@+{-# 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.Display.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++toUExpr :: (Columnable a) => Expr a -> UExpr+toUExpr = UExpr++fromUExpr :: forall a. (Columnable a) => UExpr -> Maybe (Expr a)+fromUExpr (UExpr (expr :: Expr b)) = do+    Refl <- testEquality (typeRep @a) (typeRep @b)+    pure expr++type NamedExpr = (T.Text, UExpr)++toNamedExpr :: (Columnable a) => T.Text -> Expr a -> NamedExpr+toNamedExpr exprName expr = (exprName, UExpr expr)++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
+ src-internal/DataFrame/Internal/Expression/Operators.hs view
@@ -0,0 +1,425 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE TypeOperators #-}++module DataFrame.Internal.Expression.Operators where++import Data.Function ((&))+import qualified Data.Text as T+import DataFrame.Internal.Column (Columnable)+import DataFrame.Internal.Column.Types (Promote, PromoteDiv)+import DataFrame.Internal.Expression (+    BinUDF (MkBinaryOp),+    BinaryOp (+        binaryCommutative,+        binaryFn,+        binaryName,+        binaryPrecedence,+        binarySymbol+    ),+    Expr (Binary, Col, If, Lit, Unary),+    NamedExpr,+    UExpr (UExpr),+    UnUDF (MkUnaryOp),+ )+import DataFrame.Internal.Expression.Operators.Nullable (+    BaseType,+    DivWidenOp,+    NullCmpResult,+    NullLift2Op (applyNull2),+    NullableCmpOp (nullCmpOp),+    NumericWidenOp,+    WidenResult,+    WidenResultDiv,+    divArithOp,+    widenArithOp,+    widenCmpOp,+ )++infixr 8 .^^, .^^., .^, .^.+infixl 7 .*, ./, .*., ./.+infixl 6 .+, .-, .+., .-.+infix 4 .==, .==., .<, .<., .<=, .<=., .>=, .>=., .>, .>., ./=, ./=.+infixr 3 .&&, .&&.+infixr 2 .||, .||.+infixr 0 .=++(|>) :: a -> (a -> b) -> b+(|>) = (&)++as :: (Columnable a) => Expr a -> T.Text -> NamedExpr+as expr colName = (colName, UExpr expr)++name :: (Show a) => Expr a -> T.Text+name (Col n) = n+name other =+    error $+        "You must call `name` on a column reference. Not the expression: " ++ show other++col :: (Columnable a) => T.Text -> Expr a+col = Col++ifThenElse :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a+ifThenElse = If++lit :: (Columnable a) => a -> Expr a+lit = Lit++(.=) :: (Columnable a) => T.Text -> Expr a -> NamedExpr+(.=) = flip as++liftDecorated ::+    (Columnable a, Columnable b) =>+    (a -> b) -> T.Text -> Maybe T.Text -> Expr a -> Expr b+liftDecorated f opName rep = Unary (MkUnaryOp f opName rep)++lift2Decorated ::+    (Columnable c, Columnable b, Columnable a) =>+    (c -> b -> a) ->+    T.Text ->+    Maybe T.Text ->+    Bool ->+    Int ->+    Expr c ->+    Expr b ->+    Expr a+lift2Decorated f opName rep comm prec =+    Binary (MkBinaryOp f opName rep comm prec)++data NullEq a b c where+    NullEq ::+        ( NumericWidenOp (BaseType a) (BaseType b)+        , NullLift2Op a b Bool (NullCmpResult a b)+        , Eq (Promote (BaseType a) (BaseType b))+        ) =>+        NullEq a b (NullCmpResult a b)++data NullNeq a b c where+    NullNeq ::+        ( NumericWidenOp (BaseType a) (BaseType b)+        , NullLift2Op a b Bool (NullCmpResult a b)+        , Eq (Promote (BaseType a) (BaseType b))+        ) =>+        NullNeq a b (NullCmpResult a b)++data NullLt a b c where+    NullLt ::+        ( NumericWidenOp (BaseType a) (BaseType b)+        , NullLift2Op a b Bool (NullCmpResult a b)+        , Ord (Promote (BaseType a) (BaseType b))+        ) =>+        NullLt a b (NullCmpResult a b)++data NullGt a b c where+    NullGt ::+        ( NumericWidenOp (BaseType a) (BaseType b)+        , NullLift2Op a b Bool (NullCmpResult a b)+        , Ord (Promote (BaseType a) (BaseType b))+        ) =>+        NullGt a b (NullCmpResult a b)++data NullLeq a b c where+    NullLeq ::+        ( NumericWidenOp (BaseType a) (BaseType b)+        , NullLift2Op a b Bool (NullCmpResult a b)+        , Ord (Promote (BaseType a) (BaseType b))+        ) =>+        NullLeq a b (NullCmpResult a b)++data NullGeq a b c where+    NullGeq ::+        ( NumericWidenOp (BaseType a) (BaseType b)+        , NullLift2Op a b Bool (NullCmpResult a b)+        , Ord (Promote (BaseType a) (BaseType b))+        ) =>+        NullGeq a b (NullCmpResult a b)++data NullAnd a b c where+    NullAnd ::+        (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>+        NullAnd a b (NullCmpResult a b)++data NullOr a b c where+    NullOr ::+        (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>+        NullOr a b (NullCmpResult a b)++instance BinaryOp NullEq where+    binaryFn NullEq = applyNull2 (widenCmpOp (==))+    binaryName NullEq = "eq"+    binarySymbol NullEq = Just ".=="+    binaryCommutative NullEq = True+    binaryPrecedence NullEq = 4+instance BinaryOp NullNeq where+    binaryFn NullNeq = applyNull2 (widenCmpOp (/=))+    binaryName NullNeq = "neq"+    binarySymbol NullNeq = Just "./="+    binaryCommutative NullNeq = True+    binaryPrecedence NullNeq = 4+instance BinaryOp NullLt where+    binaryFn NullLt = applyNull2 (widenCmpOp (<))+    binaryName NullLt = "lt"+    binarySymbol NullLt = Just ".<"+    binaryPrecedence NullLt = 4+instance BinaryOp NullGt where+    binaryFn NullGt = applyNull2 (widenCmpOp (>))+    binaryName NullGt = "gt"+    binarySymbol NullGt = Just ".>"+    binaryPrecedence NullGt = 4+instance BinaryOp NullLeq where+    binaryFn NullLeq = applyNull2 (widenCmpOp (<=))+    binaryName NullLeq = "leq"+    binarySymbol NullLeq = Just ".<="+    binaryPrecedence NullLeq = 4+instance BinaryOp NullGeq where+    binaryFn NullGeq = applyNull2 (widenCmpOp (>=))+    binaryName NullGeq = "geq"+    binarySymbol NullGeq = Just ".>="+    binaryPrecedence NullGeq = 4+instance BinaryOp NullAnd where+    binaryFn NullAnd = nullCmpOp (&&)+    binaryName NullAnd = "nulland"+    binarySymbol NullAnd = Just ".&&"+    binaryCommutative NullAnd = True+    binaryPrecedence NullAnd = 3+instance BinaryOp NullOr where+    binaryFn NullOr = nullCmpOp (||)+    binaryName NullOr = "nullor"+    binarySymbol NullOr = Just ".||"+    binaryCommutative NullOr = True+    binaryPrecedence NullOr = 2++(.==.) ::+    (Columnable a, Eq a) =>+    Expr a ->+    Expr a ->+    Expr Bool+(.==.) = lift2Decorated (==) "eq" (Just ".==.") True 4++(./=.) ::+    (Columnable a, Eq a) =>+    Expr a ->+    Expr a ->+    Expr Bool+(./=.) = lift2Decorated (/=) "neq" (Just "./=.") True 4++(.<.) ::+    (Columnable a, Ord a) =>+    Expr a ->+    Expr a ->+    Expr Bool+(.<.) = lift2Decorated (<) "lt" (Just ".<.") False 4++(.>.) ::+    (Columnable a, Ord a) =>+    Expr a ->+    Expr a ->+    Expr Bool+(.>.) = lift2Decorated (>) "gt" (Just ".>.") False 4++(.<=.) ::+    (Columnable a, Ord a) =>+    Expr a ->+    Expr a ->+    Expr Bool+(.<=.) = lift2Decorated (<=) "leq" (Just ".<=.") False 4++(.>=.) ::+    (Columnable a, Ord a) =>+    Expr a ->+    Expr a ->+    Expr Bool+(.>=.) = lift2Decorated (>=) "geq" (Just ".>=.") False 4++(.+.) :: (Columnable a, Num a) => Expr a -> Expr a -> Expr a+(.+.) = (+)++(.-.) :: (Columnable a, Num a) => Expr a -> Expr a -> Expr a+(.-.) = (-)++(.*.) :: (Columnable a, Num a) => Expr a -> Expr a -> Expr a+(.*.) = (*)++(./.) :: (Columnable a, Fractional a) => Expr a -> Expr a -> Expr a+(./.) = (/)++-- Nullable-aware arithmetic operators++{- | Nullable-aware addition. Works for all combinations of nullable\/non-nullable operands.+@col \@Int "x" .+ col \@(Maybe Int) "y"  -- :: Expr (Maybe Int)@+-}+(.+) ::+    ( NumericWidenOp (BaseType a) (BaseType b)+    , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b)+    , Num (Promote (BaseType a) (BaseType b))+    ) =>+    Expr a ->+    Expr b ->+    Expr (WidenResult a b)+(.+) = lift2Decorated (applyNull2 (widenArithOp (+))) "nulladd" (Just ".+") True 6++-- | Nullable-aware subtraction.+(.-) ::+    ( NumericWidenOp (BaseType a) (BaseType b)+    , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b)+    , Num (Promote (BaseType a) (BaseType b))+    ) =>+    Expr a ->+    Expr b ->+    Expr (WidenResult a b)+(.-) = lift2Decorated (applyNull2 (widenArithOp (-))) "nullsub" (Just ".-") False 6++-- | Nullable-aware multiplication.+(.*) ::+    ( NumericWidenOp (BaseType a) (BaseType b)+    , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b)+    , Num (Promote (BaseType a) (BaseType b))+    ) =>+    Expr a ->+    Expr b ->+    Expr (WidenResult a b)+(.*) = lift2Decorated (applyNull2 (widenArithOp (*))) "nullmul" (Just ".*") True 7++-- | Nullable-aware division. Integral operands are promoted to Double.+(./) ::+    ( DivWidenOp (BaseType a) (BaseType b)+    , NullLift2Op a b (PromoteDiv (BaseType a) (BaseType b)) (WidenResultDiv a b)+    , Fractional (PromoteDiv (BaseType a) (BaseType b))+    ) =>+    Expr a ->+    Expr b ->+    Expr (WidenResultDiv a b)+(./) = lift2Decorated (applyNull2 (divArithOp (/))) "nulldiv" (Just "./") False 7++-- Nullable-aware comparison operators (three-valued logic: Nothing if either operand is Nothing)++{- | Nullable-aware equality. Widens numeric operands to their common type,+so @Expr Double .== Expr Int@ typechecks. Returns @Maybe Bool@ when either+operand is nullable.+-}+(.==) ::+    ( NumericWidenOp (BaseType a) (BaseType b)+    , NullLift2Op a b Bool (NullCmpResult a b)+    , Eq (Promote (BaseType a) (BaseType b))+    ) =>+    Expr a ->+    Expr b ->+    Expr (NullCmpResult a b)+(.==) = Binary NullEq++-- | Nullable-aware inequality. Widens numeric operands to their common type.+(./=) ::+    ( NumericWidenOp (BaseType a) (BaseType b)+    , NullLift2Op a b Bool (NullCmpResult a b)+    , Eq (Promote (BaseType a) (BaseType b))+    ) =>+    Expr a ->+    Expr b ->+    Expr (NullCmpResult a b)+(./=) = Binary NullNeq++-- | Nullable-aware less-than. Widens numeric operands to their common type.+(.<) ::+    ( NumericWidenOp (BaseType a) (BaseType b)+    , NullLift2Op a b Bool (NullCmpResult a b)+    , Ord (Promote (BaseType a) (BaseType b))+    ) =>+    Expr a ->+    Expr b ->+    Expr (NullCmpResult a b)+(.<) = Binary NullLt++-- | Nullable-aware greater-than. Widens numeric operands to their common type.+(.>) ::+    ( NumericWidenOp (BaseType a) (BaseType b)+    , NullLift2Op a b Bool (NullCmpResult a b)+    , Ord (Promote (BaseType a) (BaseType b))+    ) =>+    Expr a ->+    Expr b ->+    Expr (NullCmpResult a b)+(.>) = Binary NullGt++{- | Nullable-aware less-than-or-equal. Widens numeric operands to their+common type, so @Expr Double .<= Expr Int@ typechecks.+-}+(.<=) ::+    ( NumericWidenOp (BaseType a) (BaseType b)+    , NullLift2Op a b Bool (NullCmpResult a b)+    , Ord (Promote (BaseType a) (BaseType b))+    ) =>+    Expr a ->+    Expr b ->+    Expr (NullCmpResult a b)+(.<=) = Binary NullLeq++-- | Nullable-aware greater-than-or-equal. Widens numeric operands to their common type.+(.>=) ::+    ( NumericWidenOp (BaseType a) (BaseType b)+    , NullLift2Op a b Bool (NullCmpResult a b)+    , Ord (Promote (BaseType a) (BaseType b))+    ) =>+    Expr a ->+    Expr b ->+    Expr (NullCmpResult a b)+(.>=) = Binary NullGeq++(.&&.) :: Expr Bool -> Expr Bool -> Expr Bool+(.&&.) = lift2Decorated (&&) "and" (Just ".&&.") True 3++(.||.) :: Expr Bool -> Expr Bool -> Expr Bool+(.||.) = lift2Decorated (||) "or" (Just ".||.") True 2++-- | Nullable-aware logical AND. Returns @Maybe Bool@ when either operand is nullable.+(.&&) ::+    (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>+    Expr a ->+    Expr b ->+    Expr (NullCmpResult a b)+(.&&) = Binary NullAnd++-- | Nullable-aware logical OR. Returns @Maybe Bool@ when either operand is nullable.+(.||) ::+    (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>+    Expr a ->+    Expr b ->+    Expr (NullCmpResult a b)+(.||) = Binary NullOr++(.^^) ::+    ( Columnable (BaseType a)+    , Columnable (BaseType b)+    , Fractional (BaseType a)+    , Integral (BaseType b)+    , NumericWidenOp (BaseType a) (BaseType b)+    , NullLift2Op a b (BaseType a) a+    , Num (Promote (BaseType a) (BaseType b))+    ) =>+    Expr a -> Expr b -> Expr a+(.^^) = lift2Decorated (applyNull2 (^^)) "pow" (Just ".^^") False 8++(.^) ::+    ( Columnable (BaseType a)+    , Columnable (BaseType b)+    , Num (BaseType a)+    , Integral (BaseType b)+    , NumericWidenOp (BaseType a) (BaseType b)+    , NullLift2Op a b (BaseType a) a+    , Num (Promote (BaseType a) (BaseType b))+    ) =>+    Expr a -> Expr b -> Expr a+(.^) = lift2Decorated (applyNull2 (^)) "pow" (Just ".^") False 8++-- Same-type (non-nullable) exponentiation operators++(.^^.) ::+    (Columnable a, Columnable b, Fractional a, Integral b) =>+    Expr a -> Expr b -> Expr a+(.^^.) = lift2Decorated (^^) "pow" (Just ".^^.") False 8++(.^.) ::+    (Columnable a, Columnable b, Num a, Integral b) =>+    Expr a -> Expr b -> Expr a+(.^.) = lift2Decorated (^) "pow" (Just ".^.") False 8
+ src-internal/DataFrame/Internal/Expression/Operators/Nullable.hs view
@@ -0,0 +1,467 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE FunctionalDependencies #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE UndecidableSuperClasses #-}++{- | Nullable-aware arithmetic and comparison operators ('.+', '.==', …) that work+transparently across nullable (@Maybe a@) and non-nullable (@a@) operands.+Functional dependencies infer the result type without annotations.++@+-- Mixing nullable and non-nullable columns:+F.col \@Int \"x\" '.+' F.col \@(Maybe Int) \"y\"  -- :: Expr (Maybe Int)++-- Both non-nullable (existing behaviour preserved):+F.col \@Int \"x\" '.+' F.col \@Int \"y\"           -- :: Expr Int++-- Comparison with three-valued logic:+F.col \@(Maybe Int) \"x\" '.==' F.col \@Int \"y\"  -- :: Expr (Maybe Bool)+@+-}+module DataFrame.Internal.Expression.Operators.Nullable (+    -- * Type family+    BaseType,++    -- * Arithmetic class+    NullableArithOp (..),++    -- * Comparison class+    NullableCmpOp (..),++    -- * Generalized nullable lift classes+    NullLift1Op (..),+    NullLift2Op (..),++    -- * Result-type type families (drive inference in nullLift / nullLift2)+    NullLift1Result,+    NullLift2Result,++    -- * Result-type type family for comparison operators+    NullCmpResult,++    -- * Numeric widening+    NumericWidenOp (..),+    widenArithOp,+    widenCmpOp,+    WidenResult,++    -- * Division widening (integral × integral → Double)+    DivWidenOp (..),+    divArithOp,+    WidenResultDiv,+) where++import Data.Int (Int32, Int64)+import DataFrame.Internal.Column (Columnable)+import DataFrame.Internal.Column.Types (Promote, PromoteDiv)++{- | Strip one layer of 'Maybe'.++@+BaseType (Maybe a) = a+BaseType a         = a   -- for any non-Maybe type+@+-}+type family BaseType a where+    BaseType (Maybe a) = a+    BaseType a = a++{- | Arithmetic binary operations that work over nullable and non-nullable operand+types. The functional dependency @a b -> c@ infers the result; the 'OVERLAPPABLE'+non-nullable instance yields to the specific @(Maybe a, Maybe a)@ one.+-}+class+    ( Columnable a+    , Columnable b+    , Columnable c+    ) =>+    NullableArithOp a b c+        | a b -> c+    where+    {- | Lift an arithmetic function over the inner (non-Maybe) values.+    'Nothing' short-circuits: any 'Nothing' operand produces 'Nothing'.+    -}+    nullArithOp ::+        (BaseType a -> BaseType a -> BaseType a) ->+        a ->+        b ->+        c++{- | Compute the result type of a nullable comparison.++@+NullCmpResult (Maybe a) b = Maybe Bool+NullCmpResult a (Maybe b) = Maybe Bool   -- when a is apart from Maybe+NullCmpResult a b         = Bool+@++Used by the comparison operators ('.==', '.<', etc.) so GHC infers the+return type without an explicit annotation.+-}+type family NullCmpResult a b where+    NullCmpResult (Maybe a) b = Maybe Bool+    NullCmpResult a (Maybe b) = Maybe Bool+    NullCmpResult a b = Bool++{- | Comparison binary operations over nullable and non-nullable operands. No+functional dependency on @e@; overlapping/overlappable instance pragmas pick the+unique most-specific instance from the concrete operand types.+-}+class+    ( Columnable a+    , Columnable b+    , Columnable e+    ) =>+    NullableCmpOp a b e+    where+    {- | Lift a comparison function over the inner values (three-valued logic).+    Returns 'Nothing' when either operand is 'Nothing'.+    -}+    nullCmpOp ::+        (BaseType a -> BaseType a -> Bool) ->+        a ->+        b ->+        e++{- | Non-nullable × Non-nullable: apply directly, no wrapping.+Arithmetic result is @a@; comparison result is @Bool@.+-}+instance+    {-# OVERLAPPABLE #-}+    (Columnable a, a ~ BaseType a) =>+    NullableArithOp a a a+    where+    nullArithOp f = f++instance+    {-# OVERLAPPABLE #-}+    (Columnable a, Columnable Bool, a ~ BaseType a) =>+    NullableCmpOp a a Bool+    where+    nullCmpOp f = f++-- | Nullable × Non-nullable: 'Nothing' short-circuits.+instance+    (Columnable a, Columnable (Maybe a)) =>+    NullableArithOp (Maybe a) a (Maybe a)+    where+    nullArithOp _f Nothing _ = Nothing+    nullArithOp f (Just x) y = Just (f x y)++instance+    (Columnable a, Columnable (Maybe a), Columnable (Maybe Bool)) =>+    NullableCmpOp (Maybe a) a (Maybe Bool)+    where+    nullCmpOp _f Nothing _ = Nothing+    nullCmpOp f (Just x) y = Just (f x y)++-- | Non-nullable × Nullable: 'Nothing' short-circuits.+instance+    ( Columnable a+    , Columnable (Maybe a)+    , a ~ BaseType a+    ) =>+    NullableArithOp a (Maybe a) (Maybe a)+    where+    nullArithOp _f _ Nothing = Nothing+    nullArithOp f x (Just y) = Just (f x y)++instance+    ( Columnable a+    , Columnable (Maybe a)+    , Columnable (Maybe Bool)+    , a ~ BaseType a+    ) =>+    NullableCmpOp a (Maybe a) (Maybe Bool)+    where+    nullCmpOp _f _ Nothing = Nothing+    nullCmpOp f x (Just y) = Just (f x y)++-- | Nullable × Nullable: either 'Nothing' short-circuits.+instance+    {-# OVERLAPPING #-}+    (Columnable a, Columnable (Maybe a)) =>+    NullableArithOp (Maybe a) (Maybe a) (Maybe a)+    where+    nullArithOp _f Nothing _ = Nothing+    nullArithOp _f _ Nothing = Nothing+    nullArithOp f (Just x) (Just y) = Just (f x y)++instance+    {-# OVERLAPPING #-}+    (Columnable a, Columnable (Maybe a), Columnable (Maybe Bool)) =>+    NullableCmpOp (Maybe a) (Maybe a) (Maybe Bool)+    where+    nullCmpOp _f Nothing _ = Nothing+    nullCmpOp _f _ Nothing = Nothing+    nullCmpOp f (Just x) (Just y) = Just (f x y)++-- ---------------------------------------------------------------------------+-- Generalized nullable lift (unary)+-- ---------------------------------------------------------------------------++{- | Lift a unary function over a column expression, propagating 'Nothing' (applied+directly when non-nullable, under 'Just' when @a = Maybe x@). Use via+'DataFrame.Functions.nullLift'.+-}++{- | Compute the result type of a nullable unary lift.++@+NullLift1Result (Maybe a) r = Maybe r+NullLift1Result a         r = r        -- for any non-Maybe a+@++Used by 'DataFrame.Functions.nullLift' so GHC can infer the return type+without an explicit annotation.+-}+type family NullLift1Result a r where+    NullLift1Result (Maybe a) r = Maybe r+    NullLift1Result a r = r++class+    ( Columnable a+    , Columnable r+    , Columnable c+    ) =>+    NullLift1Op a r c+    where+    applyNull1 :: (BaseType a -> r) -> a -> c++-- | Non-nullable: apply directly.+instance+    {-# OVERLAPPABLE #-}+    (Columnable a, Columnable r, a ~ BaseType a) =>+    NullLift1Op a r r+    where+    applyNull1 f = f++-- | Nullable: propagate 'Nothing'.+instance+    {-# OVERLAPPING #-}+    (Columnable a, Columnable r, Columnable (Maybe r)) =>+    NullLift1Op (Maybe a) r (Maybe r)+    where+    applyNull1 _ Nothing = Nothing+    applyNull1 f (Just x) = Just (f x)++-- ---------------------------------------------------------------------------+-- Generalized nullable lift (binary)+-- ---------------------------------------------------------------------------++{- | Lift a binary function over two column expressions, propagating 'Nothing': the+result is @Maybe r@ if either operand is nullable, else @r@. Use via+'DataFrame.Functions.nullLift2'.+-}++{- | Compute the result type of a nullable binary lift.++@+NullLift2Result (Maybe a) b         r = Maybe r+NullLift2Result a         (Maybe b) r = Maybe r   -- when a is apart from Maybe+NullLift2Result a         b         r = r+@++Used by 'DataFrame.Functions.nullLift2' so GHC can infer the return type.+-}+type family NullLift2Result a b r where+    NullLift2Result (Maybe a) b r = Maybe r+    NullLift2Result a (Maybe b) r = Maybe r+    NullLift2Result a b r = r++class+    ( Columnable a+    , Columnable b+    , Columnable r+    , Columnable c+    ) =>+    NullLift2Op a b r c+    where+    applyNull2 :: (BaseType a -> BaseType b -> r) -> a -> b -> c++-- | Both non-nullable: apply directly.+instance+    {-# OVERLAPPABLE #-}+    (Columnable a, Columnable b, Columnable r, a ~ BaseType a, b ~ BaseType b) =>+    NullLift2Op a b r r+    where+    applyNull2 f = f++-- | Left nullable: 'Nothing' short-circuits.+instance+    {-# OVERLAPPABLE #-}+    (Columnable a, Columnable b, Columnable r, Columnable (Maybe r), b ~ BaseType b) =>+    NullLift2Op (Maybe a) b r (Maybe r)+    where+    applyNull2 _ Nothing _ = Nothing+    applyNull2 f (Just x) y = Just (f x y)++-- | Right nullable: 'Nothing' short-circuits.+instance+    {-# OVERLAPPABLE #-}+    (Columnable a, Columnable b, Columnable r, Columnable (Maybe r), a ~ BaseType a) =>+    NullLift2Op a (Maybe b) r (Maybe r)+    where+    applyNull2 _ _ Nothing = Nothing+    applyNull2 f x (Just y) = Just (f x y)++-- | Both nullable: either 'Nothing' short-circuits.+instance+    {-# OVERLAPPING #-}+    (Columnable a, Columnable b, Columnable r, Columnable (Maybe r)) =>+    NullLift2Op (Maybe a) (Maybe b) r (Maybe r)+    where+    applyNull2 _ Nothing _ = Nothing+    applyNull2 _ _ Nothing = Nothing+    applyNull2 f (Just x) (Just y) = Just (f x y)++-- ---------------------------------------------------------------------------+-- Numeric widening+-- ---------------------------------------------------------------------------++{- | Widen two numeric base types to their promoted common type.++When @a ~ b@ the coercions are identity; otherwise one operand is widened+(e.g. 'Int' → 'Double').+-}+class (Columnable (Promote a b)) => NumericWidenOp a b where+    widen1 :: a -> Promote a b+    widen2 :: b -> Promote a b++-- | Same type: identity coercions.+instance {-# OVERLAPPING #-} (Columnable a) => NumericWidenOp a a where+    widen1 = id+    widen2 = id++instance NumericWidenOp Int Double where widen1 = fromIntegral; widen2 = id+instance NumericWidenOp Double Int where+    widen1 = id+    widen2 = fromIntegral+instance NumericWidenOp Float Double where widen1 = realToFrac; widen2 = id+instance NumericWidenOp Double Float where+    widen1 = id+    widen2 = realToFrac+instance NumericWidenOp Int32 Float where widen1 = fromIntegral; widen2 = id+instance NumericWidenOp Float Int32 where+    widen1 = id+    widen2 = fromIntegral+instance NumericWidenOp Int32 Double where widen1 = fromIntegral; widen2 = id+instance NumericWidenOp Double Int32 where+    widen1 = id+    widen2 = fromIntegral+instance NumericWidenOp Int64 Float where widen1 = fromIntegral; widen2 = id+instance NumericWidenOp Float Int64 where+    widen1 = id+    widen2 = fromIntegral+instance NumericWidenOp Int64 Double where widen1 = fromIntegral; widen2 = id+instance NumericWidenOp Double Int64 where+    widen1 = id+    widen2 = fromIntegral++-- | Apply an arithmetic function after widening both operands to their common type.+widenArithOp ::+    forall a b.+    (NumericWidenOp a b) =>+    (Promote a b -> Promote a b -> Promote a b) ->+    a ->+    b ->+    Promote a b+widenArithOp f x y = f (widen1 @a @b x) (widen2 @a @b y)++-- | Apply a comparison function after widening both operands to their common type.+widenCmpOp ::+    forall a b.+    (NumericWidenOp a b) =>+    (Promote a b -> Promote a b -> Bool) ->+    a ->+    b ->+    Bool+widenCmpOp f x y = f (widen1 @a @b x) (widen2 @a @b y)++-- | Result type of a widening binary operator, accounting for nullable wrappers.+type WidenResult a b = NullLift2Result a b (Promote (BaseType a) (BaseType b))++-- ---------------------------------------------------------------------------+-- Division widening (integral × integral → Double)+-- ---------------------------------------------------------------------------++{- | Like 'NumericWidenOp' but uses 'PromoteDiv': integral×integral → Double.+Floating types still dominate (Double > Float), and any two integral types+(same or mixed) are both widened to Double.+-}+class (Columnable (PromoteDiv a b)) => DivWidenOp a b where+    divWiden1 :: a -> PromoteDiv a b+    divWiden2 :: b -> PromoteDiv a b++-- Floating same-type (identity)+instance DivWidenOp Double Double where divWiden1 = id; divWiden2 = id+instance DivWidenOp Float Float where divWiden1 = id; divWiden2 = id++-- Mixed Double/Float+instance DivWidenOp Double Float where divWiden1 = id; divWiden2 = realToFrac+instance DivWidenOp Float Double where divWiden1 = realToFrac; divWiden2 = id++-- Double beats integral+instance DivWidenOp Double Int where divWiden1 = id; divWiden2 = fromIntegral+instance DivWidenOp Int Double where divWiden1 = fromIntegral; divWiden2 = id+instance DivWidenOp Double Int32 where divWiden1 = id; divWiden2 = fromIntegral+instance DivWidenOp Int32 Double where divWiden1 = fromIntegral; divWiden2 = id+instance DivWidenOp Double Int64 where divWiden1 = id; divWiden2 = fromIntegral+instance DivWidenOp Int64 Double where divWiden1 = fromIntegral; divWiden2 = id++-- Float beats integral+instance DivWidenOp Float Int where divWiden1 = id; divWiden2 = fromIntegral+instance DivWidenOp Int Float where divWiden1 = fromIntegral; divWiden2 = id+instance DivWidenOp Float Int32 where divWiden1 = id; divWiden2 = fromIntegral+instance DivWidenOp Int32 Float where divWiden1 = fromIntegral; divWiden2 = id+instance DivWidenOp Float Int64 where divWiden1 = id; divWiden2 = fromIntegral+instance DivWidenOp Int64 Float where divWiden1 = fromIntegral; divWiden2 = id++-- Integral × integral → Double+instance DivWidenOp Int Int where+    divWiden1 = fromIntegral+    divWiden2 = fromIntegral+instance DivWidenOp Int32 Int32 where+    divWiden1 = fromIntegral+    divWiden2 = fromIntegral+instance DivWidenOp Int64 Int64 where+    divWiden1 = fromIntegral+    divWiden2 = fromIntegral+instance DivWidenOp Int Int32 where+    divWiden1 = fromIntegral+    divWiden2 = fromIntegral+instance DivWidenOp Int32 Int where+    divWiden1 = fromIntegral+    divWiden2 = fromIntegral+instance DivWidenOp Int Int64 where+    divWiden1 = fromIntegral+    divWiden2 = fromIntegral+instance DivWidenOp Int64 Int where+    divWiden1 = fromIntegral+    divWiden2 = fromIntegral+instance DivWidenOp Int32 Int64 where+    divWiden1 = fromIntegral+    divWiden2 = fromIntegral+instance DivWidenOp Int64 Int32 where+    divWiden1 = fromIntegral+    divWiden2 = fromIntegral++-- | Apply an arithmetic function after widening both operands via 'PromoteDiv'.+divArithOp ::+    forall a b.+    (DivWidenOp a b) =>+    (PromoteDiv a b -> PromoteDiv a b -> PromoteDiv a b) ->+    a ->+    b ->+    PromoteDiv a b+divArithOp f x y = f (divWiden1 @a @b x) (divWiden2 @a @b y)++-- | Result type of a division-widening binary operator, accounting for nullable wrappers.+type WidenResultDiv a b =+    NullLift2Result a b (PromoteDiv (BaseType a) (BaseType b))
+ src-internal/DataFrame/Internal/Expression/Simplify.hs view
@@ -0,0 +1,417 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE MultiWayIf #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module DataFrame.Internal.Expression.Simplify (+    simplify,+    simplifyPredicatePair,++    -- * Path-condition entailment (for fitted-tree pruning)+    PredFact,+    factTrue,+    factFalse,+    entails,+) where++import Control.Monad (guard)+import Data.Maybe (fromMaybe)+import Data.Type.Equality (testEquality, (:~:) (Refl))+import Type.Reflection (eqTypeRep, typeRep, (:~~:) (HRefl), pattern App)++import DataFrame.Internal.Column (Columnable)+import DataFrame.Internal.Expression (+    BinaryOp,+    Expr (..),+    UnaryOp (unaryName),+    eqExpr,+    normalize,+ )+import DataFrame.Internal.Expression.Operators (+    NullAnd,+    NullEq,+    NullGeq,+    NullGt,+    NullLeq,+    NullLt,+    NullNeq,+    NullOr,+    (.==.),+ )++simplify :: forall a. (Columnable a) => Expr a -> Expr a+simplify e+    | isBoolish @a = fixpoint (10 :: Int) e+    | otherwise = e+  where+    fixpoint 0 x = x+    fixpoint n x = let x' = simplifyB x in if eqExpr x x' then x else fixpoint (n - 1) x'++isBoolish :: forall a. (Columnable a) => Bool+isBoolish =+    case ( testEquality (typeRep @a) (typeRep @Bool)+         , testEquality (typeRep @a) (typeRep @(Maybe Bool))+         ) of+        (Just Refl, _) -> True+        (_, Just Refl) -> True+        _ -> False++data Conn = ConnAnd | ConnOr++connOf :: forall op c b r. (BinaryOp op) => op c b r -> Maybe Conn+connOf _+    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullAnd) = Just ConnAnd+    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullOr) = Just ConnOr+    | otherwise = Nothing++simplifyB :: forall a. (Columnable a) => Expr a -> Expr a+simplifyB expr = case expr of+    Binary (op :: op c b a) l r+        | Just conn <- connOf op+        , Just Refl <- testEquality (typeRep @c) (typeRep @a)+        , Just Refl <- testEquality (typeRep @b) (typeRep @a) ->+            let l' = simplifyB l; r' = simplifyB r+             in fromMaybe (Binary op l' r') (combine conn l' r')+        | otherwise -> expr+    Unary (op :: op b a) inner+        | Just Refl <- testEquality (typeRep @a) (typeRep @Bool)+        , Just Refl <- testEquality (typeRep @b) (typeRep @Bool)+        , unaryName op == "not" ->+            simplifyNot op (simplifyB inner)+        | otherwise -> expr+    If c t f ->+        let c' = simplify c+            t' = simplifyB t+            f' = simplifyB f+         in case asBoolLit c' of+                Just True -> t'+                Just False -> f'+                Nothing+                    | eqExpr t' f' -> t'+                    | Just Refl <- testEquality (typeRep @a) (typeRep @Bool)+                    , asBoolLit t' == Just True+                    , asBoolLit f' == Just False ->+                        c'+                    | otherwise -> If c' t' f'+    _ -> expr++simplifyNot :: (UnaryOp op) => op Bool Bool -> Expr Bool -> Expr Bool+simplifyNot op inner = case asBoolLit inner of+    Just b -> Lit (not b)+    Nothing -> case inner of+        Unary (op2 :: op2 b2 Bool) inner2+            | unaryName op2 == "not"+            , Just Refl <- testEquality (typeRep @b2) (typeRep @Bool) ->+                inner2+        _ -> Unary op inner++combine :: (Columnable a) => Conn -> Expr a -> Expr a -> Maybe (Expr a)+combine ConnAnd = combineAnd+combine ConnOr = combineOr++asBoolLit :: forall a. (Columnable a) => Expr a -> Maybe Bool+asBoolLit (Lit v) =+    case testEquality (typeRep @a) (typeRep @Bool) of+        Just Refl -> Just v+        Nothing -> case testEquality (typeRep @a) (typeRep @(Maybe Bool)) of+            Just Refl -> v+            Nothing -> Nothing+asBoolLit _ = Nothing++{- | Polymorphic boolean literal: @Lit b@ for @Expr Bool@, @Lit (Just b)@ for+@Expr (Maybe Bool)@.+-}+litBoolish :: forall a. (Columnable a) => Bool -> Maybe (Expr a)+litBoolish v =+    case testEquality (typeRep @a) (typeRep @Bool) of+        Just Refl -> Just (Lit v)+        Nothing -> case testEquality (typeRep @a) (typeRep @(Maybe Bool)) of+            Just Refl -> Just (Lit (Just v))+            Nothing -> Nothing++combineAnd :: (Columnable a) => Expr a -> Expr a -> Maybe (Expr a)+combineAnd l r+    | eqExpr l r = Just l+    | asBoolLit l == Just False = litBoolish False+    | asBoolLit r == Just False = litBoolish False+    | asBoolLit l == Just True = Just r+    | asBoolLit r == Just True = Just l+    | absorbs ConnOr l r = Just l+    | absorbs ConnOr r l = Just r+    | otherwise = simplifyPredicatePair True l r++combineOr :: (Columnable a) => Expr a -> Expr a -> Maybe (Expr a)+combineOr l r+    | eqExpr l r = Just l+    | asBoolLit l == Just True = litBoolish True+    | asBoolLit r == Just True = litBoolish True+    | asBoolLit l == Just False = Just r+    | asBoolLit r == Just False = Just l+    | absorbs ConnAnd l r = Just l+    | absorbs ConnAnd r l = Just r+    | otherwise = simplifyPredicatePair False l r++absorbs :: (Columnable a) => Conn -> Expr a -> Expr a -> Bool+absorbs conn x (Binary (op :: op c b a) ya yb)+    | Just c' <- connOf op+    , sameConn conn c'+    , Just Refl <- testEquality (typeRep @c) (typeRep @a)+    , Just Refl <- testEquality (typeRep @b) (typeRep @a) =+        eqExpr x ya || eqExpr x yb+absorbs _ _ _ = False++sameConn :: Conn -> Conn -> Bool+sameConn ConnAnd ConnAnd = True+sameConn ConnOr ConnOr = True+sameConn _ _ = False++data Cmp = CLt | CLeq | CGt | CGeq | CEq | CNeq deriving (Eq)++data NullK = Total | FalseOnNull | UnknownOnNull deriving (Eq)++data Atom = Atom+    { aCmp :: Cmp+    , aThr :: !Double+    , aKey :: String+    , aNull :: NullK+    , aIntegral :: Bool+    }++cmpOf :: forall op c b r. (BinaryOp op) => op c b r -> Maybe Cmp+cmpOf _+    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullLt) = Just CLt+    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullLeq) = Just CLeq+    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullGt) = Just CGt+    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullGeq) = Just CGeq+    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullEq) = Just CEq+    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullNeq) = Just CNeq+    | otherwise = Nothing++isLower, isUpper :: Cmp -> Bool+isLower c = c == CGt || c == CGeq+isUpper c = c == CLt || c == CLeq++-- | True if @x@ is a @Maybe _@ type.+isMaybeTy :: forall x. (Columnable x) => Bool+isMaybeTy = case typeRep @x of+    App con _ -> case eqTypeRep con (typeRep @Maybe) of Just HRefl -> True; _ -> False+    _ -> False++litDouble :: forall b. (Columnable b) => Expr b -> Maybe Double+litDouble (Lit v) =+    case testEquality (typeRep @b) (typeRep @Double) of+        Just Refl -> Just v+        Nothing -> case testEquality (typeRep @b) (typeRep @Int) of+            Just Refl -> Just (fromIntegral v)+            Nothing -> case testEquality (typeRep @b) (typeRep @(Maybe Double)) of+                Just Refl -> v+                Nothing -> case testEquality (typeRep @b) (typeRep @(Maybe Int)) of+                    Just Refl -> fromIntegral <$> v+                    Nothing -> Nothing+litDouble _ = Nothing++{- | True for a column lifted from an integral type (never NaN): @toDouble (col …)@+or a column whose type is itself integral.+-}+integralColE :: forall c. (Columnable c) => Expr c -> Bool+integralColE (Unary op _) = unaryName op == "toDouble"+integralColE _ =+    or+        [ matches @Int+        , matches @(Maybe Int)+        ]+  where+    matches :: forall t. (Columnable t) => Bool+    matches = case testEquality (typeRep @c) (typeRep @t) of Just Refl -> True; _ -> False++atomOf :: forall a. (Columnable a) => Expr a -> Maybe Atom+atomOf (Unary fm (Binary (op :: op c b r) (colE :: Expr c) litE))+    | unaryName fm == "fromMaybe"+    , Just cmp <- cmpOf op+    , Just t <- litDouble litE =+        Just (Atom cmp t (show (normalize colE)) FalseOnNull (integralColE colE))+atomOf (Binary (op :: op c b a) (colE :: Expr c) litE)+    | Just cmp <- cmpOf op+    , Just t <- litDouble litE =+        let nk = if isMaybeTy @c then UnknownOnNull else Total+         in Just (Atom cmp t (show (normalize colE)) nk (integralColE colE))+atomOf _ = Nothing++simplifyPredicatePair ::+    forall a. (Columnable a) => Bool -> Expr a -> Expr a -> Maybe (Expr a)+simplifyPredicatePair isAnd a b = do+    atomA <- atomOf a+    atomB <- atomOf b+    guard (aKey atomA == aKey atomB)+    let nk = aNull atomA+        integral = aIntegral atomA+    if isAnd+        then andAtoms a atomA b atomB nk integral+        else orAtoms a atomA b atomB nk integral++-- | Contradiction folds to a literal False unless null-rows make it unknown.+litFalseGated :: (Columnable a) => NullK -> Maybe (Expr a)+litFalseGated UnknownOnNull = Nothing+litFalseGated _ = litBoolish False++{- | Tautology to literal True is sound only for total (never-null) atoms; the+exhaustive-cover form additionally needs a non-NaN (integral) column.+-}+litTrueTotal :: (Columnable a) => NullK -> Maybe (Expr a)+litTrueTotal Total = litBoolish True+litTrueTotal _ = Nothing++andAtoms ::+    (Columnable a) =>+    Expr a -> Atom -> Expr a -> Atom -> NullK -> Bool -> Maybe (Expr a)+andAtoms a atomA b atomB nk _ =+    let cA = aCmp atomA; tA = aThr atomA; cB = aCmp atomB; tB = aThr atomB+     in if+            | isLower cA, isLower cB, cA == cB -> Just (if tA >= tB then a else b)+            | isUpper cA, isUpper cB, cA == cB -> Just (if tA <= tB then a else b)+            | isLower cA, isUpper cB -> lu cA tA cB tB+            | isUpper cA, isLower cB -> lu cB tB cA tA+            | cA == CEq, cB == CEq -> if tA == tB then Just a else litFalseGated nk+            | cA == CEq, cB == CNeq -> if tA == tB then litFalseGated nk else Just a+            | cA == CNeq, cB == CEq -> if tA == tB then litFalseGated nk else Just b+            | cA == CEq -> if satisfies tA cB tB then Just a else litFalseGated nk+            | cB == CEq -> if satisfies tB cA tA then Just b else litFalseGated nk+            | cA == CNeq, cB == CNeq -> Nothing+            | cA == CNeq -> if outside tA cB tB then Just b else Nothing+            | cB == CNeq -> if outside tB cA tA then Just a else Nothing+            | otherwise -> Nothing+  where+    lu lc lo uc hi+        | lo > hi = litFalseGated nk+        | lo == hi, lc == CGeq, uc == CLeq = pointEq a lo+        | lo == hi = litFalseGated nk+        | otherwise = Nothing++orAtoms ::+    (Columnable a) =>+    Expr a -> Atom -> Expr a -> Atom -> NullK -> Bool -> Maybe (Expr a)+orAtoms a atomA b atomB nk integral =+    let cA = aCmp atomA; tA = aThr atomA; cB = aCmp atomB; tB = aThr atomB+     in if+            | isLower cA, isLower cB, cA == cB -> Just (if tA <= tB then a else b)+            | isUpper cA, isUpper cB, cA == cB -> Just (if tA >= tB then a else b)+            | isUpper cA+            , isLower cB+            , nk == Total+            , integral+            , covers cB tB cA tA ->+                litTrueTotal nk+            | isLower cA+            , isUpper cB+            , nk == Total+            , integral+            , covers cA tA cB tB ->+                litTrueTotal nk+            | cA == CNeq, cB == CNeq -> if tA == tB then Just a else litTrueTotal nk+            | cA == CEq, cB == CNeq -> if tA == tB then litTrueTotal nk else Just b+            | cA == CNeq, cB == CEq -> if tA == tB then litTrueTotal nk else Just a+            | cA == CEq, cB == CEq -> if tA == tB then Just a else Nothing+            | otherwise -> Nothing++{- | Build @col == t@ for the point-collapse rule; only strict @Expr Bool@ over a+@Double@ column (otherwise bail).+-}+pointEq :: forall a. (Columnable a) => Expr a -> Double -> Maybe (Expr a)+pointEq atom lo = case testEquality (typeRep @a) (typeRep @Bool) of+    Just Refl -> (\colE -> colE .==. Lit lo) <$> recoverColD atom+    Nothing -> Nothing++recoverColD :: Expr x -> Maybe (Expr Double)+recoverColD (Binary _ (colE :: Expr c) _) =+    case testEquality (typeRep @c) (typeRep @Double) of+        Just Refl -> Just colE+        _ -> Nothing+recoverColD (Unary _ inner) = recoverColD inner+recoverColD _ = Nothing++covers :: Cmp -> Double -> Cmp -> Double -> Bool+covers lowerCmp lo upperCmp hi =+    lo < hi || (lo == hi && (lowerCmp == CGeq || upperCmp == CLeq))++satisfies :: Double -> Cmp -> Double -> Bool+satisfies t CGt tb = t > tb+satisfies t CGeq tb = t >= tb+satisfies t CLt tb = t < tb+satisfies t CLeq tb = t <= tb+satisfies _ _ _ = False++outside :: Double -> Cmp -> Double -> Bool+outside t CGt tb = t <= tb+outside t CGeq tb = t < tb+outside t CLt tb = t >= tb+outside t CLeq tb = t > tb+outside _ _ _ = False++-- ---------------------------------------------------------------------------+-- Path-condition entailment for fitted-tree pruning.+-- ---------------------------------------------------------------------------++-- | A known same-column threshold fact accumulated along a tree path.+data PredFact = PredFact !String !Cmp !Double++-- | The fact a branch's true edge establishes (the condition holds).+factTrue :: Expr Bool -> Maybe PredFact+factTrue e = (\a -> PredFact (aKey a) (aCmp a) (aThr a)) <$> atomOf e++{- | The fact a branch's false edge establishes (the negated condition). Only+sound for non-NaN (integral) columns — a NaN row takes the false edge too,+so @¬(x>t)@ is not a clean @x<=t@ bound for floats.+-}+factFalse :: Expr Bool -> Maybe PredFact+factFalse e = do+    a <- atomOf e+    guard (aIntegral a && aNull a == Total)+    nc <- negCmp (aCmp a)+    pure (PredFact (aKey a) nc (aThr a))++negCmp :: Cmp -> Maybe Cmp+negCmp CLt = Just CGeq+negCmp CLeq = Just CGt+negCmp CGt = Just CLeq+negCmp CGeq = Just CLt+negCmp _ = Nothing++{- | @entails facts cond@: 'Just' 'True' when the path facts force @cond@ true,+'Just' 'False' when they force it false, 'Nothing' when undecided.+-}+entails :: [PredFact] -> Expr Bool -> Maybe Bool+entails facts cond = do+    a <- atomOf cond+    let decisions =+            [ d+            | PredFact fk fc ft <- facts+            , fk == aKey a+            , Just d <- [factImplies (fc, ft) (aCmp a, aThr a)]+            ]+    case decisions of+        (d : _) -> Just d+        [] -> Nothing++{- | Does the fact's solution set sit inside @cond@ ('Just' 'True'), disjoint+from it ('Just' 'False'), or neither ('Nothing')? Boundary strictness is+honoured: e.g. @x<=t@ does NOT entail @x<t@, and @x>=t ∧ x<=t@ is not empty.+-}+factImplies :: (Cmp, Double) -> (Cmp, Double) -> Maybe Bool+factImplies (fc, ft) (cc, tc)+    | isLower fc, isLower cc, subset = Just True+    | isUpper fc, isUpper cc, subset = Just True+    | isLower fc, isUpper cc, disjointAtEq = Just False+    | isUpper fc, isLower cc, disjointBelow = Just False+    | otherwise = Nothing+  where+    fIncl = fc == CGeq || fc == CLeq+    cIncl = cc == CGeq || cc == CLeq+    subset =+        (if isLower fc then ft > tc else ft < tc)+            || (ft == tc && (not fIncl || cIncl))+    disjointAtEq = ft > tc || (ft == tc && not (fIncl && cIncl))+    disjointBelow = ft < tc || (ft == tc && not (fIncl && cIncl))
+ src-internal/DataFrame/Internal/Grouping.hs view
@@ -0,0 +1,1053 @@+{-# 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.Bits (unsafeShiftR, (.&.))+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))+import DataFrame.Errors+import DataFrame.Internal.Algorithms.Hash+import DataFrame.Internal.Algorithms.Rank.Radix (rankByHash)+import DataFrame.Internal.Column (+    Column (..),+    materializeMerged,+ )+import DataFrame.Internal.Column.Bitmap (+    Bitmap,+    bitmapTestBit,+ )+import DataFrame.Internal.Column.Encode (dictEncodeColumnUpTo)+import DataFrame.Internal.Column.Types+import DataFrame.Internal.Control.Concurrent (+    capabilities,+    chunksFor,+    forkJoin,+    parThreshold,+    shouldParallelize,+    splitChunkRange,+ )+import DataFrame.Internal.Data.HashTable (htInsert, newHashTable)+import DataFrame.Internal.Data.PackedText (+    PackedSel,+    PackedTextData (..),+    offCount,+    packedLength,+    packedSlice,+    selAt,+    sliceEqBytes,+ )+import DataFrame.Internal.DataFrame (DataFrame (..), GroupedDataFrame (..))+import DataFrame.Internal.Grouping.Direct (+    DirectGrouping (..),+    ascendingCodeGroups,+    directGroupThreshold,+    rangeOf,+    tryDirectGroupColumn,+ )+import qualified DataFrame.Internal.Grouping.Direct as GD+import DataFrame.Internal.Grouping.Partitioned (+    parallelAssignGroups,+    rtgFromVisOffs,+ )+import DataFrame.Internal.Row.RowHash (computeRowHashesWithIO)+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 parThreshold n = groupByPar names df+    | otherwise = groupBySeq names df+  where+    !n = nRows df++{- | Low-cardinality direct-indexed grouping fast path+('DataFrame.Internal.Grouping.Direct'): fires for key lists+where every key is a single clean small-range @Int@ column or a canonical+dict-encoded text column, and the product of the key domains stays within+'directGroupThreshold' (the keys fuse into one mixed-radix code; a single key+degenerates to its own code). Returns 'Nothing' on any other key shape, falling+back to the hash path.++Narrow domains build @offsets@/@groupRepRows@ eagerly (histogram-sized work+only) and leave BOTH per-row outputs lazy: @valueIndices@ ('visFromCodes')+only materializes for consumers that gather (median/top-k, set ops,+interpreter slices), and @rowToGroup@ ('rtgFromCodes') only for the streaming+scatter aggregations — each aggregate pays for exactly one O(n) output pass,+not both. Wide domains run the two-level radix engine instead+('DataFrame.Internal.Grouping.Direct.directLayoutLazy'): @rowToGroup@ eager,+@valueIndices@ deferred (see 'fusedDirectGroup').+-}+tryDirectGroup :: [T.Text] -> DataFrame -> Maybe GroupedDataFrame+tryDirectGroup [] _ = Nothing+tryDirectGroup names df = do+    cols <-+        traverse (\nm -> M.lookup nm (columnIndices df) >>= (columns df V.!?)) names+    case traverse fusedKey cols of+        Just keys -> fusedDirectGroup names df keys+        Nothing -> case (names, cols) of+            ([name], [col]) -> tryDictGroup (nRows df) df [name] col+            _ -> Nothing++{- | One key column of a fused multi-key direct grouping: a per-row component+code in @[0, fkDomain)@ (negative marks an invalid/corrupt code, which aborts+the direct path).+-}+data FusedKey = FusedKey+    { fkCode :: Int -> Int+    , fkDomain :: !Int+    }++{- | Classify a key column for the fused multi-key direct path: a clean non-null+unboxed @Int@ of small range, or a non-null canonical dict-encoded text column+of small dictionary. Anything else falls back to the hash group-by.+-}+fusedKey :: Column -> Maybe FusedKey+fusedKey (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 (FusedKey (\i -> VU.unsafeIndex v i - mn) range)+                else Nothing+fusedKey (PackedText Nothing p)+    | Just sel <- ptSel p+    , ptCanonicalSel p =+        let offs = ptOffsets p+            !card = offCount offs - 1+         in if card >= 1 && card <= directGroupThreshold+                then+                    Just+                        ( FusedKey+                            (\i -> let c = selAt sel i in if c >= card then -1 else c)+                            card+                        )+                else Nothing+fusedKey _ = Nothing++{- | Fuse the per-key codes into one mixed-radix code per row+(@((k1*d2)+k2)*d3+...@) and feed the direct counting-sort machinery. Group+order: ascending fused code (lexicographic in key order) — ascending value+order for @Int@ keys (mirroring the order the single-@Int@-key direct path+always had) and ascending dictionary code for dict-encoded text keys (the+dictionary's first-appearance order, a fixed property of the column). The+ascending order keeps @codeToGroup@ an identity map whenever the domain is+fully occupied, so the deferred @rowToGroup@ pass skips its per-row random+remap lookup; ranking dict groups by string hash instead (the historical+order) profiled ~0.6s slower per 1e8 rows at 1e6 groups.++On the narrow-domain engine, @valueIndices@ and @rowToGroup@ are passed to the+constructor as unevaluated applications of 'visFromCodes' / 'rtgFromCodes'+(constructor arguments are not forced even under @-XStrict@, and the fields+are lazy at their definition site), so each per-row output pass is deferred+until a consumer demands it. The wide-domain engine defers only+@valueIndices@ (see the branch comment below).+-}+fusedDirectGroup ::+    [T.Text] -> DataFrame -> [FusedKey] -> Maybe GroupedDataFrame+fusedDirectGroup names df keys = do+    domain <- fusedDomain (map fkDomain keys)+    let n = nRows df+        codeAt' = fusedCodeAt keys+    if GD.useTwoLevel n domain+        then do+            {- Wide domains (> ~1024 codes at parallel scale): the two-level+            radix engine — no pass random-writes a multi-megabyte table per+            worker, unlike the per-chunk direct histograms below (measured ~2x+            on the eager layout at 1e6 codes / 1e8 rows). It builds+            @rowToGroup@ eagerly (the streaming aggregations force it first+            thing anyway); only @valueIndices@ stays deferred, reconstructed+            from @rowToGroup@ by the same engine on demand. -}+            (rtg, offs, reps, nGroups) <-+                GD.directLayoutLazy codeAt' n domain ascendingCodeGroups+            Just+                ( GroupedInternal+                    df+                    names+                    (GD.visFromRowToGroup n nGroups offs rtg)+                    offs+                    rtg+                    reps+                )+        else do+            (offs, reps, counts, ctg, hists, nGroups) <-+                directLayoutLazy codeAt' n domain ascendingCodeGroups+            Just+                ( GroupedInternal+                    df+                    names+                    (visFromCodes codeAt' counts ctg offs hists n domain)+                    offs+                    (rtgFromCodes codeAt' ctg n)+                    reps+                )++{- | Product of the per-key domains, 'Nothing' once it (or any factor) passes+'directGroupThreshold'. Factors are capped before multiplying, so the running+product never exceeds @threshold^2@ and cannot overflow.+-}+fusedDomain :: [Int] -> Maybe Int+fusedDomain = go 1+  where+    go !acc [] = Just acc+    go !acc (d : ds)+        | d < 1 || d > directGroupThreshold = Nothing+        | acc * d > directGroupThreshold = Nothing+        | otherwise = go (acc * d) ds++{- | Per-row fused mixed-radix code; @-1@ when any component code is invalid+(only possible for corrupt dict codes), making 'groupCodesMaybe' bail to the+hash path. Valid components compose to a code in @[0, product of domains)@.+-}+fusedCodeAt :: [FusedKey] -> (Int -> Int)+fusedCodeAt [] = const (-1)+fusedCodeAt (k0 : ks0) = go (fkCode k0) ks0+  where+    go f [] = f+    go f (k : ks) =+        let !d = fkDomain k+            g = fkCode k+         in go+                ( \i ->+                    let a = f i+                     in if a < 0+                            then -1+                            else let b = g i in if b < 0 then -1 else a * d + b+                )+                ks++{- | 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 parThreshold 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.Grouping.Partitioned'):+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+        -- Merged key columns are exotic; hash their eager form.+        selectedCols = map (materializeMerged . (columns df V.!)) indicesToGroup+        !eqRow = eqKeyRow df indicesToGroup+        (vis, os) = unsafePerformIO $ do+            -- Parallel row-hash kernel, bit-identical to 'computeHashes' at the+            -- same dict-code setting (grouping always hashes canonical dict+            -- columns by code; see 'hashPacked').+            hashes <- computeRowHashesWithIO True n selectedCols+            parallelAssignGroups n hashes eqRow+     in -- rowToGroup is passed as an UNFORCED constructor argument (this module+        -- is -XStrict, so it must not be let-bound): gather-style aggregation+        -- over huge group counts never reads it, and the deferred pass writes+        -- values identical to the eager build.+        Grouped df names vis os (rtgFromVisOffs n vis os)+{-# 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+    -- Merged key columns are exotic; hash their eager form.+    let selectedCols = map (materializeMerged . (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+        MergedColumn _ _ ->+            error "computeHashes: MergedColumn is normalized before hashing"+    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 c@(MergedColumn _ _) = colEqRow (materializeMerged c)+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) =+    -- A canonical dictionary selection assigns equal strings the same code,+    -- so two rows are byte-equal iff their codes agree.+    let eqV = case ptSel p of+            Just sel+                | ptCanonicalSel p ->+                    \a b -> selAt sel a == selAt sel b+            _ -> \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)++-------------------------------------------------------------------------------+-- Deferred-placement direct grouping+-------------------------------------------------------------------------------++{- | Contiguous per-worker row ranges: one chunk per capability above the+parallel threshold, a single chunk otherwise.+-}+directRowChunks :: Int -> [(Int, Int)]+directRowChunks = chunksFor parThreshold++{- | Like 'directRowChunks' but over a code/group domain (merge/seed passes),+which pays for a fan-out at a much lower width than the row passes do.+-}+directCodeSlices :: Int -> [(Int, Int)]+directCodeSlices = chunksFor 4096++{- | The eager phases of the direct counting-sort grouping — WITHOUT either+per-row output pass: per-chunk validated code histograms (parallel), per-code+totals and first occurrences (parallel over code slices), the caller-chosen+code->group mapping, the offsets prefix scan and per-group representative+rows. Returns+@(offsets, groupRepRows, counts, codeToGroup, chunkHists, nGroups)@ — the last+three feed the deferred @valueIndices@ placement ('visFromCodes') and+@rowToGroup@ ('rtgFromCodes') thunks, so a consumer pays only for the per-row+output it actually demands. 'Nothing' when any row's code falls outside+@[0, card)@ (fall back to hashing).++Pure w.r.t. its immutable inputs: the fork fan-out is a fixed function of the+row count and capability count, and every merge runs in fixed chunk order, so+the result is deterministic and the 'unsafePerformIO' is safe.+-}+directLayoutLazy ::+    (Int -> Int) ->+    Int ->+    Int ->+    (VU.Vector Int -> (VU.Vector Int, Int)) ->+    Maybe+        ( VU.Vector Int+        , VU.Vector Int+        , VU.Vector Int+        , VU.Vector Int+        , [VU.Vector Int]+        , Int+        )+directLayoutLazy codeAt n card mkGroups+    | n <= 0 || card <= 0 = Nothing+    | n < packedRowLimit = unsafePerformIO $ do+        -- Packed variant: count and first-occurrence row share one word per+        -- code, keeping phase 1 at a single accumulator array per chunk+        -- (measured ~0.15s/1e8 rows cheaper than a second firstOcc array).+        let chunks = directRowChunks n+        parts <- forkJoin [histFirstChunkPacked codeAt card lo hi | (lo, hi) <- chunks]+        if not (all snd parts)+            then pure Nothing+            else finishLayout card mkGroups (map fst parts) $ \histsM totalsM firstAllM lo hi ->+                sumFirstSlicePacked histsM totalsM firstAllM lo hi+    | otherwise = unsafePerformIO $ do+        -- Fallback for gigantic frames where a row index does not fit the+        -- packed word: separate count and firstOcc arrays, same results.+        let chunks = directRowChunks n+        parts <- forkJoin [histFirstChunk codeAt card lo hi | (lo, hi) <- chunks]+        if not (all (\(_, _, ok) -> ok) parts)+            then pure Nothing+            else+                finishLayout+                    card+                    mkGroups+                    (map (\(h, _, _) -> h) parts)+                    ( \histsM totalsM firstAllM lo hi ->+                        sumFirstSlice histsM (map (\(_, f, _) -> f) parts) totalsM firstAllM lo hi+                    )+{-# NOINLINE directLayoutLazy #-}++{- | Shared tail of 'directLayoutLazy': run the totals/first-occurrence merge+(which also normalizes each chunk histogram to plain counts, see+'sumFirstSlicePacked'), derive the group mapping, offsets and representative+rows, and freeze the retained chunk histograms.+-}+finishLayout ::+    Int ->+    (VU.Vector Int -> (VU.Vector Int, Int)) ->+    [VUM.IOVector Int] ->+    ( [VUM.IOVector Int] ->+      VUM.IOVector Int ->+      VUM.IOVector Int ->+      Int ->+      Int ->+      IO ()+    ) ->+    IO+        ( Maybe+            ( VU.Vector Int+            , VU.Vector Int+            , VU.Vector Int+            , VU.Vector Int+            , [VU.Vector Int]+            , Int+            )+        )+finishLayout card mkGroups histsM mergeSlice = do+    totalsM <- VUM.new card+    firstAllM <- VUM.new card+    _ <-+        forkJoin+            [ mergeSlice histsM totalsM firstAllM lo hi+            | (lo, hi) <- directCodeSlices card+            ]+    counts <- VU.unsafeFreeze totalsM+    firstAll <- VU.unsafeFreeze firstAllM+    let (codeToGroup, nGroups) = mkGroups counts+    offs <- scanGroupOffsets counts codeToGroup nGroups+    repsM <- VUM.new nGroups+    _ <-+        forkJoin+            [ scatterRepsSlice counts codeToGroup firstAll repsM lo hi+            | (lo, hi) <- directCodeSlices card+            ]+    reps <- VU.unsafeFreeze repsM+    hists <- mapM VU.unsafeFreeze histsM+    pure (Just (offs, reps, counts, codeToGroup, hists, nGroups))++{- | Rows must satisfy @row + 1 < 2^31@ for the packed count/first-row encoding+(count in the high bits, first row + 1 in the low 31). Above it (a >2e9-row+frame, >17GB per Int column) the unpacked variant runs instead.+-}+packedRowLimit :: Int+packedRowLimit = 0x7FFFFFFF++-- | One unit of count in the packed encoding; also the low-bits mask + 1.+packedCountOne :: Int+packedCountOne = 0x80000000++{- | Whether @codeToGroup@ maps every code to itself (fully occupied ascending+domain — e.g. a dense Int key covering its whole range). The rowToGroup pass+then skips the random remap lookup entirely.+-}+isIdentityMap :: VU.Vector Int -> Bool+isIdentityMap m = go 0+  where+    !k = VU.length m+    go !i+        | i >= k = True+        | VU.unsafeIndex m i /= i = False+        | otherwise = go (i + 1)++{- | Histogram one row chunk with the packed encoding: slot @c@ holds+@count(c) * 2^31 + (firstRow(c) + 1)@ (zero = never seen). One accumulator+array per chunk. Reports 'False' as soon as any code escapes @[0, card)@.+-}+histFirstChunkPacked ::+    (Int -> Int) -> Int -> Int -> Int -> IO (VUM.IOVector Int, Bool)+histFirstChunkPacked codeAt card lo hi = do+    acc <- VUM.replicate card (0 :: Int)+    let go !i+            | i >= hi = pure True+            | otherwise = do+                let !c = codeAt i+                if c < 0 || c >= card+                    then pure False+                    else do+                        x <- VUM.unsafeRead acc c+                        VUM.unsafeWrite+                            acc+                            c+                            (if x == 0 then packedCountOne + (i + 1) else x + packedCountOne)+                        go (i + 1)+    ok <- go lo+    pure (acc, ok)++{- | Per-code totals and overall first occurrences from the PACKED chunk+histograms, rewriting each histogram slot to its plain count in place (the+placement thunk then sees ordinary counts). Chunks are ordered by row range, so+the first chunk with a nonzero slot holds the code's globally first row.+-}+sumFirstSlicePacked ::+    [VUM.IOVector Int] ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    IO ()+sumFirstSlicePacked hists totals firstAll lo hi = go lo+  where+    go !c+        | c >= hi = pure ()+        | otherwise = do+            let sumP [] !acc !firstRow = pure (acc, firstRow)+                sumP (h : hs) !acc !firstRow = do+                    x <- VUM.unsafeRead h c+                    let !cnt = x `unsafeShiftR` 31+                    VUM.unsafeWrite h c cnt+                    if firstRow < 0 && x /= 0+                        then sumP hs (acc + cnt) ((x .&. (packedCountOne - 1)) - 1)+                        else sumP hs (acc + cnt) firstRow+            (s, fo) <- sumP hists 0 (-1)+            VUM.unsafeWrite totals c s+            VUM.unsafeWrite firstAll c fo+            go (c + 1)++{- | Histogram one row chunk into a private @card@-slot count plus the chunk's+first occurrence of each code, reporting 'False' as soon as any code escapes+@[0, card)@ (the counts are then abandoned). Fallback for frames beyond+'packedRowLimit'.+-}+histFirstChunk ::+    (Int -> Int) ->+    Int ->+    Int ->+    Int ->+    IO (VUM.IOVector Int, VUM.IOVector Int, Bool)+histFirstChunk codeAt card lo hi = do+    acc <- VUM.replicate card (0 :: Int)+    firstOcc <- VUM.replicate card (-1 :: Int)+    let go !i+            | i >= hi = pure True+            | otherwise = do+                let !c = codeAt i+                if c < 0 || c >= card+                    then pure False+                    else do+                        x <- VUM.unsafeRead acc c+                        VUM.unsafeWrite acc c (x + 1)+                        when (x == 0) (VUM.unsafeWrite firstOcc c i)+                        go (i + 1)+    ok <- go lo+    pure (acc, firstOcc, ok)++{- | Per-code totals over one code slice, plus the overall first occurrence of+each code: the chunks are ordered by row range, so the first chunk with a+nonzero count for a code holds its globally first row.+-}+sumFirstSlice ::+    [VUM.IOVector Int] ->+    [VUM.IOVector Int] ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    IO ()+sumFirstSlice hists firsts totals firstAll lo hi = go lo+  where+    go !c+        | c >= hi = pure ()+        | otherwise = do+            let sumP [] [] !acc !firstRow = pure (acc, firstRow)+                sumP (h : hs) (f : fs) !acc !firstRow = do+                    x <- VUM.unsafeRead h c+                    if firstRow < 0 && x > 0+                        then do+                            fo <- VUM.unsafeRead f c+                            sumP hs fs (acc + x) fo+                        else sumP hs fs (acc + x) firstRow+                sumP _ _ _ _ = error "sumFirstSlice: mismatched partials"+            (s, fo) <- sumP hists firsts 0 (-1)+            VUM.unsafeWrite totals c s+            VUM.unsafeWrite firstAll c fo+            go (c + 1)++{- | Exclusive prefix scan of per-group counts (gathered through @codeToGroup@)+into the offsets array of length @nGroups + 1@.+-}+scanGroupOffsets :: VU.Vector Int -> VU.Vector Int -> Int -> IO (VU.Vector Int)+scanGroupOffsets counts codeToGroup nGroups = do+    let !card = VU.length counts+    grpCount <- VUM.new nGroups+    let gather !c+            | c >= card = pure ()+            | otherwise = do+                let !cnt = VU.unsafeIndex counts c+                if cnt == 0+                    then gather (c + 1)+                    else do+                        VUM.unsafeWrite grpCount (VU.unsafeIndex codeToGroup c) cnt+                        gather (c + 1)+    gather 0+    offsM <- VUM.new (nGroups + 1)+    let scan !g !acc+            | g >= nGroups = VUM.unsafeWrite offsM nGroups acc+            | otherwise = do+                VUM.unsafeWrite offsM g acc+                c <- VUM.unsafeRead grpCount g+                scan (g + 1) (acc + c)+    scan 0 0+    VU.unsafeFreeze offsM++{- | @reps[codeToGroup c] = firstAll c@ for every occupied code: each group is+exactly one occupied code, so this is a disjoint parallel write and equals+@vis[offs[g]]@ (the group's first row in original order).+-}+scatterRepsSlice ::+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    IO ()+scatterRepsSlice counts codeToGroup firstAll repsM lo hi = go lo+  where+    go !c+        | c >= hi = pure ()+        | VU.unsafeIndex counts c == 0 = go (c + 1)+        | otherwise = do+            VUM.unsafeWrite+                repsM+                (VU.unsafeIndex codeToGroup c)+                (VU.unsafeIndex firstAll c)+            go (c + 1)++{- | Deferred @rowToGroup@: one parallel per-row pass mapping each row's code+through @codeToGroup@ (skipping the lookup entirely when the map is the+identity, i.e. a fully occupied ascending domain). Only the streaming+aggregation paths force this; a purely gather-driven consumer (median, top-k)+never pays for it.++Pure w.r.t. its immutable inputs and deterministic (fixed chunking), so the+'unsafePerformIO' behind a lazy field is safe: whenever and however many times+the thunk is forced it yields the same vector.+-}+rtgFromCodes :: (Int -> Int) -> VU.Vector Int -> Int -> VU.Vector Int+rtgFromCodes codeAt codeToGroup n = unsafePerformIO $ do+    rtgM <- VUM.new n+    let identity = isIdentityMap codeToGroup+    _ <-+        forkJoin+            [ ( if identity+                    then rtgChunkIdentity codeAt rtgM lo hi+                    else rtgChunk codeAt codeToGroup rtgM lo hi+              )+            | (lo, hi) <- directRowChunks n+            ]+    VU.unsafeFreeze rtgM+{-# NOINLINE rtgFromCodes #-}++-- | @rtg[i] = codeToGroup (codeAt i)@ over one row chunk (disjoint writes).+rtgChunk ::+    (Int -> Int) -> VU.Vector Int -> VUM.IOVector Int -> Int -> Int -> IO ()+rtgChunk codeAt codeToGroup rtgM lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            VUM.unsafeWrite rtgM i (VU.unsafeIndex codeToGroup (codeAt i))+            go (i + 1)++-- | 'rtgChunk' without the remap lookup (codeToGroup is the identity).+rtgChunkIdentity ::+    (Int -> Int) -> VUM.IOVector Int -> Int -> Int -> IO ()+rtgChunkIdentity codeAt rtgM lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            VUM.unsafeWrite rtgM i (codeAt i)+            go (i + 1)++{- | Deferred stable placement: build the @valueIndices@ permutation from the+per-row codes and the RETAINED phase-1 chunk histograms — the same+seed-cursors-then-place structure (and cost) the eager path used, minus the+@rowToGroup@ writes. Each chunk's code-indexed cursor starts at the group+offset plus everything earlier chunks (in row order) place there, so rows keep+original order within each group: the result is the unique group-major,+original-row-order permutation, bit-identical to the eager placement at any+chunk count.++Pure w.r.t. its immutable inputs and deterministic (fixed chunking, fixed merge+order), so the 'unsafePerformIO' behind a lazy field is safe: whenever and+however many times the thunk is forced it yields the same vector. The thunk+retains the chunk histograms (capabilities x card words) until forced or the+grouping is dropped.+-}+visFromCodes ::+    (Int -> Int) ->+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    [VU.Vector Int] ->+    Int ->+    Int ->+    VU.Vector Int+visFromCodes codeAt counts codeToGroup offs hists n card = unsafePerformIO $ do+    let chunks = directRowChunks n+    -- Private mutable copies of the retained histograms, rewritten in place+    -- into the per-chunk write cursors.+    cursors <- mapM VU.thaw hists+    _ <-+        forkJoin+            [ seedCursorSlice counts codeToGroup offs cursors lo hi+            | (lo, hi) <- directCodeSlices card+            ]+    vis <- VUM.new n+    _ <-+        forkJoin+            [ placeVisChunk codeAt cursor vis lo hi+            | ((lo, hi), cursor) <- zip chunks cursors+            ]+    VU.unsafeFreeze vis+{-# NOINLINE visFromCodes #-}++{- | Rewrite each chunk's histogram copy in place into its disjoint write+cursor: chunk w's run for code c starts at the offset of c's group plus what+earlier chunks (in row order) place there.+-}+seedCursorSlice ::+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    [VUM.IOVector Int] ->+    Int ->+    Int ->+    IO ()+seedCursorSlice counts codeToGroup offs cursors lo hi = go lo+  where+    go !c+        | c >= hi = pure ()+        | VU.unsafeIndex counts c == 0 = go (c + 1)+        | otherwise = do+            let !g = VU.unsafeIndex codeToGroup c+                loop [] !_ = pure ()+                loop (cur : rest) !acc = do+                    t <- VUM.unsafeRead cur c+                    VUM.unsafeWrite cur c acc+                    loop rest (acc + t)+            loop cursors (VU.unsafeIndex offs g)+            go (c + 1)++{- | Stable placement over one row chunk via the chunk's advancing+code-indexed cursors.+-}+placeVisChunk ::+    (Int -> Int) -> VUM.IOVector Int -> VUM.IOVector Int -> Int -> Int -> IO ()+placeVisChunk codeAt cursor vis lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !c = codeAt i+            pos <- VUM.unsafeRead cursor c+            VUM.unsafeWrite vis pos i+            VUM.unsafeWrite cursor c (pos + 1)+            go (i + 1)++{- | 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. A canonical+dict-encoded column (equal strings share a code) instead mixes its 'Int' code+with one 'mixInt' per row; 'DataFrame.Internal.RowHash.packedRange' applies the+same rule under the grouping setting so 'groupBySeq' and 'groupByPar' bucket+identically.+-}+hashPacked ::+    VUM.MVector s Int -> Maybe Bitmap -> PackedTextData -> ST s ()+hashPacked mh bm p = case ptSel p of+    Just sel | ptCanonicalSel p -> goCodes sel 0+    _ -> 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)+    goCodes !sel !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+                    _ -> mixInt h (selAt sel i)+            VUM.unsafeWrite mh i h'+            goCodes sel (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)
+ src-internal/DataFrame/Internal/Grouping/Direct.hs view
@@ -0,0 +1,859 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | Low-cardinality direct-indexed grouping fast path: when every row's key+reduces to a dense @Int@ code in a small domain, the code itself indexes a dense+accumulator (no hashing/probing). All O(n) passes run chunked across+capabilities — per-chunk histograms feed prefix-summed disjoint write cursors —+so the stable within-group row order of the sequential counting sort is+reproduced exactly.++Two engines cover the code-domain spectrum with bit-identical results:++* narrow domains (@card <= 'twoLevelCardThreshold'@) index per-worker+  histogram\/cursor tables directly — they stay cache-resident;++* wide domains use a two-level radix split (top code bits pick one of ~@2^10@+  buckets, cursors cache-resident) so no pass ever random-writes a+  multi-megabyte table per worker.++'directLayoutLazy' is the aggregation entry point: it builds @rowToGroup@,+@offsets@ and the per-group representative rows eagerly but skips the O(n)+stable placement entirely; 'visFromRowToGroup' reconstructs @valueIndices@+on demand (its value is the unique stable counting-sort permutation, so WHEN it+runs is unobservable).+-}+module DataFrame.Internal.Grouping.Direct (+    directGroupThreshold,+    tryDirectGroupColumn,+    groupCodesMaybe,+    directLayoutLazy,+    visFromRowToGroup,+    ascendingCodeGroups,+    rangeOf,+    useTwoLevel,+    DirectGrouping (..),+) where++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 Control.Monad (when)+import Control.Monad.ST (runST)+import Data.Bits (countLeadingZeros, unsafeShiftL, unsafeShiftR, (.&.))+import DataFrame.Internal.Column (Column (..))+import DataFrame.Internal.Control.Concurrent (+    capabilities,+    chunksFor,+    forkJoin,+    parThreshold,+    parallelChunks,+    pooledIndices,+    shouldParallelize,+    splitChunkRange,+ )++{- | Largest key code DOMAIN (single-key value range, or the product of per-key+domains for a fused multi-key code) 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 domains 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+    }++{- | 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+                    groupCodesMaybe+                        (\i -> VU.unsafeIndex v i - mn)+                        (VU.length v)+                        range+                        ascendingCodeGroups+                else Nothing+tryDirectGroupColumn _ = Nothing++-- | Parallel min/max reduce (order-independent).+rangeOf :: VU.Vector Int -> (Int, Int)+rangeOf v+    | not (shouldParallelize parThreshold n) = rangeChunk v 0 n+    | otherwise = unsafePerformIO $ do+        rs <- parallelChunks parThreshold n (\lo hi -> pure $! rangeChunk v lo hi)+        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++-- | Whether to fan out at this row count.+shouldPar :: Int -> Bool+shouldPar = shouldParallelize parThreshold++{- | Contiguous per-worker row ranges: one chunk per capability above the+parallel threshold, a single chunk otherwise (the sequential fallback runs the+same code on the calling thread).+-}+rowChunks :: Int -> [(Int, Int)]+rowChunks = chunksFor parThreshold++{- | Like 'rowChunks' but over the code domain (parallel merge/seed passes),+which pays for a fan-out at a much lower width than the row passes do.+-}+codeSlices :: Int -> [(Int, Int)]+codeSlices = chunksFor 4096++{- | Run each action on its own thread and collect the results in order;+rethrow the first failure. A single action runs on the calling thread.+-}++-------------------------------------------------------------------------------+-- Full grouping (eager valueIndices): compatibility entry point+-------------------------------------------------------------------------------++{- | Build the grouping by counting sort on a per-row code in @[0, card)@.+Returns 'Nothing' when any row's code falls outside @[0, card)@ (fall back to+hashing).++@mkGroups counts@ must return a dense group id for every code with a nonzero+count (other slots are never read; occupied codes must get distinct ids) and+the group count; it decides group order.++Equivalent to 'directLayoutLazy' plus a forced 'visFromRowToGroup'; kept for+callers that want the whole layout eagerly.+-}+groupCodesMaybe ::+    (Int -> Int) ->+    Int ->+    Int ->+    (VU.Vector Int -> (VU.Vector Int, Int)) ->+    Maybe DirectGrouping+groupCodesMaybe codeAt n card mkGroups = do+    (rtg, offs, _reps, nGroups) <- directLayoutLazy codeAt n card mkGroups+    let !vis = visFromRowToGroup n nGroups offs rtg+    Just (DirectGrouping rtg vis offs nGroups)++-------------------------------------------------------------------------------+-- Eager layout without placement: rowToGroup + offsets + group rep rows+-------------------------------------------------------------------------------++{- | The layout every aggregation needs, WITHOUT the O(n) stable placement:+@(rowToGroup, offsets, groupRepRows, nGroups)@, all four computed eagerly.+@groupRepRows[g]@ is the first original row of group @g@ (what+@valueIndices[offsets[g]]@ evaluates to). Pair with 'visFromRowToGroup' for a+deferred @valueIndices@. Returns 'Nothing' when any row's code falls outside+@[0, card)@.++@mkGroups@ contract as in 'groupCodesMaybe'.+-}+directLayoutLazy ::+    (Int -> Int) ->+    Int ->+    Int ->+    (VU.Vector Int -> (VU.Vector Int, Int)) ->+    Maybe (VU.Vector Int, VU.Vector Int, VU.Vector Int, Int)+directLayoutLazy codeAt n card mkGroups+    | n <= 0 || card <= 0 = Nothing+    | useTwoLevel n card = unsafePerformIO (layoutWide codeAt n card mkGroups)+    | otherwise = unsafePerformIO (layoutNarrow codeAt n card mkGroups)+{-# NOINLINE directLayoutLazy #-}++{- | Narrow domains (and the sequential small-@n@ fallback): per-chunk direct+histograms with first-occurrence tracking, merged over code slices.+-}+layoutNarrow ::+    (Int -> Int) ->+    Int ->+    Int ->+    (VU.Vector Int -> (VU.Vector Int, Int)) ->+    IO (Maybe (VU.Vector Int, VU.Vector Int, VU.Vector Int, Int))+layoutNarrow codeAt n card mkGroups = do+    let chunks = rowChunks n+    parts <- forkJoin [histFirstChunk codeAt card lo hi | (lo, hi) <- chunks]+    if not (all (\(_, _, ok) -> ok) parts)+        then pure Nothing+        else do+            let partials = [(cs, fs) | (cs, fs, _) <- parts]+            totalsM <- VUM.unsafeNew card+            firstRowM <- VUM.unsafeNew card+            _ <-+                forkJoin+                    [mergeSlice partials totalsM firstRowM lo hi | (lo, hi) <- codeSlices card]+            counts <- VU.unsafeFreeze totalsM+            firstRow <- VU.unsafeFreeze firstRowM+            finishLayout codeAt n card mkGroups counts firstRow++{- | Histogram one row chunk into a private @card@-slot count, recording the+chunk's first row of each code, and reporting invalid codes (third component+'False'; the arrays are then abandoned).+-}+histFirstChunk ::+    (Int -> Int) ->+    Int ->+    Int ->+    Int ->+    IO (VUM.IOVector Int, VUM.IOVector Int, Bool)+histFirstChunk codeAt card lo hi = do+    acc <- VUM.replicate card (0 :: Int)+    firstV <- VUM.unsafeNew card+    let go !i+            | i >= hi = pure True+            | otherwise = do+                let !c = codeAt i+                if c < 0 || c >= card+                    then pure False+                    else do+                        x <- VUM.unsafeRead acc c+                        when (x == 0) (VUM.unsafeWrite firstV c i)+                        VUM.unsafeWrite acc c (x + 1)+                        go (i + 1)+    ok <- go lo+    pure (acc, firstV, ok)++{- | @totals[c] = Σ_w counts_w[c]@ and @firstRow[c]@ = the first chunk's first+occurrence (chunks are in row order, so that IS the global first row of @c@),+over one code slice. Every @totals@ slot is written; @firstRow[c]@ only where+the count is nonzero (never read otherwise).+-}+mergeSlice ::+    [(VUM.IOVector Int, VUM.IOVector Int)] ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    IO ()+mergeSlice partials totals firstRow lo hi = go lo+  where+    go !c+        | c >= hi = pure ()+        | otherwise = do+            let sumP [] !acc = pure acc+                sumP ((cs, fs) : ps) !acc = do+                    x <- VUM.unsafeRead cs c+                    when (acc == 0 && x > 0) $+                        VUM.unsafeRead fs c >>= VUM.unsafeWrite firstRow c+                    sumP ps (acc + x)+            s <- sumP partials 0+            VUM.unsafeWrite totals c s+            go (c + 1)++{- | Wide domains: two-level radix. Rows are bucketed by the top code bits+(cache-resident cursors) as packed @(code, row)@ words; each bucket's slice —+in original row order — yields its exact per-code counts and first rows from an+L1-resident table. No placement pass runs here.+-}+layoutWide ::+    (Int -> Int) ->+    Int ->+    Int ->+    (VU.Vector Int -> (VU.Vector Int, Int)) ->+    IO (Maybe (VU.Vector Int, VU.Vector Int, VU.Vector Int, Int))+layoutWide codeAt n card mkGroups = do+    mPacked <- packBucketed True codeAt n card+    case mPacked of+        Nothing -> pure Nothing+        Just (shift, bucketStart, packed) -> do+            countsM <- VUM.unsafeNew card+            firstRowM <- VUM.unsafeNew card+            overBuckets+                bucketStart+                n+                (countFirstBucket shift card bucketStart packed countsM firstRowM)+                (countFirstBucketPar shift card bucketStart packed countsM firstRowM)+            counts <- VU.unsafeFreeze countsM+            firstRow <- VU.unsafeFreeze firstRowM+            finishLayout codeAt n card mkGroups counts firstRow++{- | Shared tail of 'directLayoutLazy': group mapping, offsets, representative+rows gathered through the code-to-group table, and the parallel @rowToGroup@+pass (one sequential read of the codes, one table lookup each).+-}+finishLayout ::+    (Int -> Int) ->+    Int ->+    Int ->+    (VU.Vector Int -> (VU.Vector Int, Int)) ->+    VU.Vector Int ->+    VU.Vector Int ->+    IO (Maybe (VU.Vector Int, VU.Vector Int, VU.Vector Int, Int))+finishLayout codeAt n card mkGroups counts firstRow = do+    let (!codeToGroup, !nGroups) = mkGroups counts+    offs <- scanOffsets counts codeToGroup nGroups+    repsM <- VUM.unsafeNew nGroups+    let repLoop !c+            | c >= card = pure ()+            | VU.unsafeIndex counts c == 0 = repLoop (c + 1)+            | otherwise = do+                VUM.unsafeWrite+                    repsM+                    (VU.unsafeIndex codeToGroup c)+                    (VU.unsafeIndex firstRow c)+                repLoop (c + 1)+    repLoop 0+    reps <- VU.unsafeFreeze repsM+    rtgM <- VUM.unsafeNew n+    -- A fully occupied ascending domain maps every code to itself; skipping+    -- the per-row random table lookup then leaves one sequential read+write.+    let identity = isIdentityMap codeToGroup+    _ <-+        forkJoin+            [ ( if identity+                    then rtgChunkIdentity codeAt rtgM lo hi+                    else rtgChunk codeAt codeToGroup rtgM lo hi+              )+            | (lo, hi) <- rowChunks n+            ]+    rtg <- VU.unsafeFreeze rtgM+    pure (Just (rtg, offs, reps, nGroups))++-- | Whether @codeToGroup@ maps every code to itself (fully occupied domain).+isIdentityMap :: VU.Vector Int -> Bool+isIdentityMap m = go 0+  where+    !k = VU.length m+    go !i+        | i >= k = True+        | VU.unsafeIndex m i /= i = False+        | otherwise = go (i + 1)++{- | Exclusive prefix scan of per-group counts (gathered through @codeToGroup@)+into the offsets array of length @nGroups + 1@.+-}+scanOffsets :: VU.Vector Int -> VU.Vector Int -> Int -> IO (VU.Vector Int)+scanOffsets counts codeToGroup nGroups = do+    let !card = VU.length counts+    grpCount <- VUM.new nGroups+    let gather !c+            | c >= card = pure ()+            | otherwise = do+                let !cnt = VU.unsafeIndex counts c+                if cnt == 0+                    then gather (c + 1)+                    else do+                        VUM.unsafeWrite grpCount (VU.unsafeIndex codeToGroup c) cnt+                        gather (c + 1)+    gather 0+    offsM <- VUM.new (nGroups + 1)+    let scan !g !acc+            | g >= nGroups = VUM.unsafeWrite offsM nGroups acc+            | otherwise = do+                VUM.unsafeWrite offsM g acc+                c <- VUM.unsafeRead grpCount g+                scan (g + 1) (acc + c)+    scan 0 0+    VU.unsafeFreeze offsM++-- | @rowToGroup@ for one row chunk: remap each row's code through the table.+rtgChunk ::+    (Int -> Int) -> VU.Vector Int -> VUM.IOVector Int -> Int -> Int -> IO ()+rtgChunk codeAt codeToGroup rtgM lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !c = codeAt i+            VUM.unsafeWrite rtgM i (VU.unsafeIndex codeToGroup c)+            go (i + 1)++-- | 'rtgChunk' without the remap lookup (codeToGroup is the identity).+rtgChunkIdentity :: (Int -> Int) -> VUM.IOVector Int -> Int -> Int -> IO ()+rtgChunkIdentity codeAt rtgM lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            VUM.unsafeWrite rtgM i (codeAt i)+            go (i + 1)++-------------------------------------------------------------------------------+-- Deferred stable placement: valueIndices from rowToGroup+-------------------------------------------------------------------------------++{- | The unique stable counting-sort permutation of @rowToGroup@: rows sorted by+group id, original order within each group — exactly what the eager engines'+placement pass produces. Pure (integer bookkeeping only, deterministic), so it+can sit under a lazy 'DataFrame.Internal.DataFrame.valueIndices' field.++Preconditions (all guaranteed by the grouping paths): @rtg@ has length @n@ with+every value in @[0, nGroups)@, and @offs@ is the group-count prefix array of+length @nGroups + 1@ with @offs[nGroups] == n@.+-}+visFromRowToGroup ::+    Int -> Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int+visFromRowToGroup n nGroups offs rtg+    | n <= 0 = VU.empty+    | useTwoLevel n nGroups = unsafePerformIO (visWide n nGroups offs rtg)+    | otherwise = unsafePerformIO (visNarrow n nGroups offs rtg)+{-# NOINLINE visFromRowToGroup #-}++{- | Narrow group domains: per-chunk group histograms prefix-summed (chunk+order) into disjoint cursors seeded from @offs@, then parallel stable placement.+-}+visNarrow :: Int -> Int -> VU.Vector Int -> VU.Vector Int -> IO (VU.Vector Int)+visNarrow n nGroups offs rtg = do+    let chunks = rowChunks n+    hists <-+        forkJoin+            [ do+                acc <- VUM.replicate nGroups (0 :: Int)+                let go !i+                        | i >= hi = pure acc+                        | otherwise = do+                            let !g = VU.unsafeIndex rtg i+                            x <- VUM.unsafeRead acc g+                            VUM.unsafeWrite acc g (x + 1)+                            go (i + 1)+                go lo+            | (lo, hi) <- chunks+            ]+    -- Rewrite each chunk histogram into its write cursor:+    -- cursor_w[g] = offs[g] + Σ_{w'<w} hist_w'[g].+    _ <-+        forkJoin+            [ let seed !g+                    | g >= hi = pure ()+                    | otherwise = do+                        let inner [] !_ = pure ()+                            inner (h : hs) !a = do+                                t <- VUM.unsafeRead h g+                                VUM.unsafeWrite h g a+                                inner hs (a + t)+                        inner hists (VU.unsafeIndex offs g)+                        seed (g + 1)+               in seed lo+            | (lo, hi) <- codeSlices nGroups+            ]+    visM <- VUM.unsafeNew n+    _ <-+        forkJoin+            [ let place !i+                    | i >= hi = pure ()+                    | otherwise = do+                        let !g = VU.unsafeIndex rtg i+                        p <- VUM.unsafeRead cursor g+                        VUM.unsafeWrite visM p i+                        VUM.unsafeWrite cursor g (p + 1)+                        place (i + 1)+               in place lo+            | ((lo, hi), cursor) <- zip chunks hists+            ]+    VU.unsafeFreeze visM++{- | Wide group domains: two-level radix. Group ids ascend with buckets, so the+bucket-sorted layout written at @offs@-seeded cursors IS @valueIndices@ — each+bucket writes one contiguous region.+-}+visWide :: Int -> Int -> VU.Vector Int -> VU.Vector Int -> IO (VU.Vector Int)+visWide n nGroups offs rtg = do+    mPacked <- packBucketed False (VU.unsafeIndex rtg) n nGroups+    case mPacked of+        Nothing -> visNarrow n nGroups offs rtg -- unreachable: no validation+        Just (shift, bucketStart, packed) -> do+            visM <- VUM.unsafeNew n+            overBuckets+                bucketStart+                n+                (placeBucketOffs shift nGroups bucketStart offs packed visM)+                (placeBucketOffsPar shift nGroups bucketStart offs packed visM)+            VU.unsafeFreeze visM++{- | Stable placement of one bucket's packed slice at cursors seeded straight+from the group offsets (the bucket's groups own a contiguous @vis@ region).+-}+placeBucketOffs ::+    Int ->+    Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    PackedBuckets ->+    VUM.IOVector Int ->+    Int ->+    IO ()+placeBucketOffs shift card bucketStart offs (PackedBuckets packed) visM b = do+    let !s = VU.unsafeIndex bucketStart b+        !e = VU.unsafeIndex bucketStart (b + 1)+        !base = b `unsafeShiftL` shift+        !range = min (1 `unsafeShiftL` shift) (card - base)+    cursor <- VUM.unsafeNew range+    let initC !j+            | j >= range = pure ()+            | otherwise = do+                VUM.unsafeWrite cursor j (VU.unsafeIndex offs (base + j))+                initC (j + 1)+    initC 0+    let place !pos+            | pos >= e = pure ()+            | otherwise = do+                pc <- VUM.unsafeRead packed pos+                let !j = (pc `unsafeShiftR` packShift) - base+                p <- VUM.unsafeRead cursor j+                VUM.unsafeWrite visM p (pc .&. packRowMask)+                VUM.unsafeWrite cursor j (p + 1)+                place (pos + 1)+    place s++{- | 'placeBucketOffs' for an oversized bucket: per-sub-chunk histograms+prefix-summed (sub-chunks in row order) onto the offset-seeded cursors keep the+placement identical to the serial walk.+-}+placeBucketOffsPar ::+    Int ->+    Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    PackedBuckets ->+    VUM.IOVector Int ->+    Int ->+    IO ()+placeBucketOffsPar shift card bucketStart offs pb visM b = do+    let !s = VU.unsafeIndex bucketStart b+        !e = VU.unsafeIndex bucketStart (b + 1)+        !base = b `unsafeShiftL` shift+        !range = min (1 `unsafeShiftL` shift) (card - base)+        packed = packedVec pb+        subChunks = [(s + lo, s + hi) | (lo, hi) <- splitChunkRange capabilities (e - s)]+    hists <-+        forkJoin+            [subHist pb base range lo hi | (lo, hi) <- subChunks]+    let seed !j+            | j >= range = pure ()+            | otherwise = do+                let inner [] !_ = pure ()+                    inner (h : hs) !a = do+                        t <- VUM.unsafeRead h j+                        VUM.unsafeWrite h j a+                        inner hs (a + t)+                inner hists (VU.unsafeIndex offs (base + j))+                seed (j + 1)+    seed 0+    _ <-+        forkJoin+            [ let place !pos+                    | pos >= hi = pure ()+                    | otherwise = do+                        pc <- VUM.unsafeRead packed pos+                        let !j = (pc `unsafeShiftR` packShift) - base+                        p <- VUM.unsafeRead cursor j+                        VUM.unsafeWrite visM p (pc .&. packRowMask)+                        VUM.unsafeWrite cursor j (p + 1)+                        place (pos + 1)+               in place lo+            | ((lo, hi), cursor) <- zip subChunks hists+            ]+    pure ()++-------------------------------------------------------------------------------+-- Two-level radix plumbing+-------------------------------------------------------------------------------++{- | Above this code-domain size the parallel passes switch to the two-level+radix engine; below it, per-worker direct tables stay cache-resident and are+faster (no bucket store).+-}+twoLevelCardThreshold :: Int+twoLevelCardThreshold = 1024++{- | The two-level engine packs @(code, row)@ into one machine word: row in the+low 'packShift' bits, code above them. Codes are capped at 'directGroupThreshold'+(@2^20@) by every caller, so the packed value stays well within 63 bits; the+guards in 'useTwoLevel' keep the narrow engine for anything larger.+-}+packShift :: Int+packShift = 40++packRowMask :: Int+packRowMask = (1 `unsafeShiftL` packShift) - 1++-- | Bucket-count target of the two-level engine (@2^10@ buckets).+bucketBits :: Int+bucketBits = 10++-- | Use the two-level engine? (Parallel-scale @n@, wide but packable domain.)+useTwoLevel :: Int -> Int -> Bool+useTwoLevel n card =+    shouldPar n+        && card > twoLevelCardThreshold+        && card <= (1 `unsafeShiftL` 22)+        && n <= packRowMask++-- | @ceilLog2 x@: smallest @s@ with @2^s >= x@ (for @x >= 1@).+ceilLog2 :: Int -> Int+ceilLog2 x+    | x <= 1 = 0+    | otherwise = 64 - countLeadingZeros (x - 1)++{- | The bucket store: rows partitioned by the top code bits, each bucket+holding packed @(code, row)@ words in original row order.+-}+newtype PackedBuckets = PackedBuckets (VUM.IOVector Int)++packedVec :: PackedBuckets -> VUM.IOVector Int+packedVec (PackedBuckets v) = v++{- | Partition rows into ~@2^'bucketBits'@ buckets by the top bits of their+code, as packed @(code, row)@ words: per-chunk bucket histograms (validating+every code when asked), prefix-summed in (bucket, chunk) order into disjoint+cursors, then a parallel scatter. Chunks are processed in row order, so every+bucket keeps its rows in ascending original row order. Returns 'Nothing' iff+validation was requested and some code fell outside @[0, card)@.+-}+packBucketed ::+    Bool ->+    (Int -> Int) ->+    Int ->+    Int ->+    IO (Maybe (Int, VU.Vector Int, PackedBuckets))+packBucketed validate codeAt n card = do+    let !shift = max 0 (ceilLog2 card - bucketBits)+        !nBuckets = ((card - 1) `unsafeShiftR` shift) + 1+        chunks = rowChunks n+    parts <-+        forkJoin+            [bucketHist validate codeAt card shift nBuckets lo hi | (lo, hi) <- chunks]+    if not (all snd parts)+        then pure Nothing+        else do+            let cursors = map fst parts+            bucketStartM <- VUM.unsafeNew (nBuckets + 1)+            let seed !b !acc+                    | b >= nBuckets = VUM.unsafeWrite bucketStartM nBuckets acc+                    | otherwise = do+                        VUM.unsafeWrite bucketStartM b acc+                        let inner [] !a = pure a+                            inner (cur : rest) !a = do+                                t <- VUM.unsafeRead cur b+                                VUM.unsafeWrite cur b a+                                inner rest (a + t)+                        acc' <- inner cursors acc+                        seed (b + 1) acc'+            seed 0 0+            bucketStart <- VU.unsafeFreeze bucketStartM+            packed <- VUM.unsafeNew n+            _ <-+                forkJoin+                    [ scatterPacked codeAt shift cur packed lo hi+                    | ((lo, hi), cur) <- zip chunks cursors+                    ]+            pure (Just (shift, bucketStart, PackedBuckets packed))++{- | Histogram one row chunk by bucket (top code bits) into a private+@nBuckets@-slot count; with @validate@, report 'False' as soon as any code+escapes @[0, card)@ (the counts are then abandoned).+-}+bucketHist ::+    Bool ->+    (Int -> Int) ->+    Int ->+    Int ->+    Int ->+    Int ->+    Int ->+    IO (VUM.IOVector Int, Bool)+bucketHist validate codeAt card shift nBuckets lo hi = do+    acc <- VUM.replicate nBuckets (0 :: Int)+    let bump !c !i = do+            let !b = c `unsafeShiftR` shift+            x <- VUM.unsafeRead acc b+            VUM.unsafeWrite acc b (x + 1)+            go (i + 1)+        go !i+            | i >= hi = pure True+            | otherwise = do+                let !c = codeAt i+                if validate && (c < 0 || c >= card)+                    then pure False+                    else bump c i+    ok <- go lo+    pure (acc, ok)++-- | Scatter one row chunk's packed @(code, row)@ words through its bucket cursor.+scatterPacked ::+    (Int -> Int) ->+    Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    IO ()+scatterPacked codeAt shift cursor packed lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !c = codeAt i+                !b = c `unsafeShiftR` shift+            pos <- VUM.unsafeRead cursor b+            VUM.unsafeWrite packed pos ((c `unsafeShiftL` packShift) + i)+            VUM.unsafeWrite cursor b (pos + 1)+            go (i + 1)++{- | Drive one action per bucket: buckets far above the fair per-worker share+run first through @big@ (internally parallel, one at a time), the rest are+pulled off a shared counter by one worker per capability. Every bucket —+including empty ones — is visited exactly once, so per-bucket passes may rely+on covering their whole output slice.+-}+overBuckets :: VU.Vector Int -> Int -> (Int -> IO ()) -> (Int -> IO ()) -> IO ()+overBuckets bucketStart n small big = do+    let !nBuckets = VU.length bucketStart - 1+        !bigCut = max parThreshold (2 * (n `div` max 1 capabilities))+        size b = VU.unsafeIndex bucketStart (b + 1) - VU.unsafeIndex bucketStart b+    mapM_ big [b | b <- [0 .. nBuckets - 1], size b >= bigCut]+    pooledIndices capabilities nBuckets $ \b ->+        when (size b < bigCut) (small b)++{- | One bucket's exact per-code counts and first rows from its (row-ordered)+packed slice, via an L1-resident table spanning only the bucket's code range.+Writes the bucket's whole slice of @counts@ (zeros included); @firstRow@ only+where the count is nonzero.+-}+countFirstBucket ::+    Int ->+    Int ->+    VU.Vector Int ->+    PackedBuckets ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    Int ->+    IO ()+countFirstBucket shift card bucketStart (PackedBuckets packed) countsM firstRowM b = do+    let !s = VU.unsafeIndex bucketStart b+        !e = VU.unsafeIndex bucketStart (b + 1)+        !base = b `unsafeShiftL` shift+        !range = min (1 `unsafeShiftL` shift) (card - base)+    local <- VUM.replicate range (0 :: Int)+    let hist !pos+            | pos >= e = pure ()+            | otherwise = do+                pc <- VUM.unsafeRead packed pos+                let !j = (pc `unsafeShiftR` packShift) - base+                x <- VUM.unsafeRead local j+                when (x == 0) $+                    VUM.unsafeWrite firstRowM (base + j) (pc .&. packRowMask)+                VUM.unsafeWrite local j (x + 1)+                hist (pos + 1)+    hist s+    let flush !j+            | j >= range = pure ()+            | otherwise = do+                t <- VUM.unsafeRead local j+                VUM.unsafeWrite countsM (base + j) t+                flush (j + 1)+    flush 0++-- | 'countFirstBucket' for an oversized bucket, chunked across capabilities.+countFirstBucketPar ::+    Int ->+    Int ->+    VU.Vector Int ->+    PackedBuckets ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    Int ->+    IO ()+countFirstBucketPar shift card bucketStart pb countsM firstRowM b = do+    let !s = VU.unsafeIndex bucketStart b+        !e = VU.unsafeIndex bucketStart (b + 1)+        !base = b `unsafeShiftL` shift+        !range = min (1 `unsafeShiftL` shift) (card - base)+        packed = packedVec pb+        subChunks = [(s + lo, s + hi) | (lo, hi) <- splitChunkRange capabilities (e - s)]+    parts <-+        forkJoin+            [ do+                local <- VUM.replicate range (0 :: Int)+                firstL <- VUM.unsafeNew range+                let hist !pos+                        | pos >= hi = pure (local, firstL)+                        | otherwise = do+                            pc <- VUM.unsafeRead packed pos+                            let !j = (pc `unsafeShiftR` packShift) - base+                            x <- VUM.unsafeRead local j+                            when (x == 0) $+                                VUM.unsafeWrite firstL j (pc .&. packRowMask)+                            VUM.unsafeWrite local j (x + 1)+                            hist (pos + 1)+                hist lo+            | (lo, hi) <- subChunks+            ]+    -- Merge in sub-chunk (= row) order: totals and global first occurrence.+    let merge !j+            | j >= range = pure ()+            | otherwise = do+                let inner [] !acc = pure acc+                    inner ((cs, fs) : ps) !acc = do+                        x <- VUM.unsafeRead cs j+                        when (acc == 0 && x > 0) $+                            VUM.unsafeRead fs j >>= VUM.unsafeWrite firstRowM (base + j)+                        inner ps (acc + x)+                t <- inner parts 0+                VUM.unsafeWrite countsM (base + j) t+                merge (j + 1)+    merge 0++-- | Per-sub-chunk histogram of one bucket's packed slice (codes only).+subHist :: PackedBuckets -> Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)+subHist (PackedBuckets packed) base range lo hi = do+    acc <- VUM.replicate range (0 :: Int)+    let go !pos+            | pos >= hi = pure acc+            | otherwise = do+                pc <- VUM.unsafeRead packed pos+                let !j = (pc `unsafeShiftR` packShift) - base+                x <- VUM.unsafeRead acc j+                VUM.unsafeWrite acc j (x + 1)+                go (pos + 1)+    go lo++{- | The ascending-code group order: walk the counts in code order, assigning a+dense group id to each non-empty code (empty codes get no id and no output+group). The single-Int-key path keeps its groups in ascending value order.+-}+ascendingCodeGroups :: VU.Vector Int -> (VU.Vector Int, Int)+ascendingCodeGroups counts = runST $ do+    let !card = VU.length counts+    m <- VUM.new card+    let go !c !next+            | c >= card = pure next+            | VU.unsafeIndex counts c > 0 = do+                VUM.unsafeWrite m c next+                go (c + 1) (next + 1)+            | otherwise = go (c + 1) next+    nGroups <- go 0 0+    frozen <- VU.unsafeFreeze m+    pure (frozen, nGroups)
+ src-internal/DataFrame/Internal/Grouping/Partitioned.hs view
@@ -0,0 +1,406 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE Strict #-}++{- | Partitioned group-by: rows are counting-sorted into partitions by the top+hash bits, then one task per capability groups its partitions independently,+after which the group numbering is canonicalized to first-appearance order.+Output is bit-for-bit identical to the sequential+'DataFrame.Internal.Grouping.groupBy'.++The name is the mechanism, not the threading: this is a genuinely different+algorithm from the sequential single-hash-table path, not that path with a+fork\/join wrapped around it. (Its sibling+"DataFrame.Internal.Grouping.Direct" is also internally parallel.) Whether to+take this path is 'DataFrame.Internal.Grouping.groupBy''s decision, not this+module's.+-}+module DataFrame.Internal.Grouping.Partitioned (+    parallelAssignGroups,+    rtgFromVisOffs,+    numPartitionsFor,+) where++import Control.Concurrent (getNumCapabilities)+import Control.Monad (forM_, when)+import Data.Bits (countLeadingZeros, unsafeShiftR)+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.Algorithms.Rank.Radix (rankByHash)+import DataFrame.Internal.Control.Concurrent (+    forkJoin,+    forkJoin_,+    parThreshold,+    parallelChunks_,+    pooledIndices,+ )+import DataFrame.Internal.Data.HashTable (+    htInsert,+    newHashTable,+ )+import System.IO.Unsafe (unsafePerformIO)++{- | 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 — and grown with the row count (up to+'maxPartitions') so a partition's worst-case hash table (every row a distinct+group: @nextPow2Above (2 * rows/p)@ slots x 3 arrays) stays cache-resident+instead of thrashing DRAM. Partitioning is by the top hash bits and canonical+ranking is ascending unsigned key both across and within partitions, so the+grouping output is bit-for-bit identical at ANY partition count; only the+constant matters for speed.+-}+numPartitionsFor :: Int -> Int -> Int+numPartitionsFor caps n = go 1+  where+    base = max 256 (4 * caps)+    go p+        | p < base = go (p * 2)+        | p < maxPartitions && n > p * partRowTarget = go (p * 2)+        | otherwise = p++-- | Cap on partition count (scatter-pass stream count stays manageable).+maxPartitions :: Int+maxPartitions = 4096++-- | Target rows per partition (~24k rows -> 64k-slot table, ~1.5MB).+partRowTarget :: Int+partRowTarget = 24576++-- | @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+@(valueIndices, offsets)@ in canonical group order. @eqRow a b@ must report+whether rows @a@ and @b@ share all key columns (null-aware). @rowToGroup@ is+NOT built here any more: gather-style aggregation over huge group counts never+reads it, so callers derive it on demand with 'rtgFromVisOffs'.+-}+parallelAssignGroups ::+    Int ->+    VU.Vector Int ->+    (Int -> Int -> Bool) ->+    IO (VU.Vector Int, VU.Vector Int)+parallelAssignGroups n hashes eqRow = do+    caps <- getNumCapabilities+    let !p = numPartitionsFor caps n+        !shift = 64 - intLog2 p+    (partStart, sortedRows, sortedHash) <- 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+        sortedHash+        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@), the row+indices laid out partition-by-partition in @sortedRows@, and each sorted+position's hash in @sortedHash@ (same layout) so the grouping loop reads its+hashes sequentially instead of a random @hashes[row]@ per row.++Runs chunked across capabilities: per-chunk partition histograms are prefix+summed (in chunk order) into disjoint per-chunk write cursors, so the scatter+threads never contend and each partition keeps its rows in ascending original+row order — bit-for-bit the sequential counting sort's layout.+-}+partitionRows ::+    Int ->+    VU.Vector Int ->+    Int ->+    Int ->+    IO (VU.Vector Int, VU.Vector Int, VU.Vector Int)+partitionRows n hashes p shift = do+    caps <- getNumCapabilities+    let chunks = rowChunks caps n+    cursors <- forkJoin [histChunk hashes p shift lo hi | (lo, hi) <- chunks]+    -- Exclusive prefix over partitions (outer) and chunks (inner): partStart+    -- from the totals, and each chunk's histogram rewritten into its cursor.+    partStartM <- VUM.new (p + 1)+    let seed !pp !acc+            | pp >= p = VUM.unsafeWrite partStartM p acc+            | otherwise = do+                VUM.unsafeWrite partStartM pp acc+                let inner [] !a = pure a+                    inner (cur : rest) !a = do+                        t <- VUM.unsafeRead cur pp+                        VUM.unsafeWrite cur pp a+                        inner rest (a + t)+                acc' <- inner cursors acc+                seed (pp + 1) acc'+    seed 0 0+    sortedM <- VUM.new (max 1 n)+    sortedHashM <- VUM.new (max 1 n)+    forkJoin_+        [ scatterChunk hashes shift cur sortedM sortedHashM lo hi+        | ((lo, hi), cur) <- zip chunks cursors+        ]+    partStart <- VU.unsafeFreeze partStartM+    sortedRows <- VU.unsafeFreeze sortedM+    sortedHash <- VU.unsafeFreeze sortedHashM+    pure (partStart, sortedRows, sortedHash)++-- | Contiguous near-equal row chunks, one per capability; empties dropped.+rowChunks :: Int -> Int -> [(Int, Int)]+rowChunks caps n =+    [ (lo, hi)+    | w <- [0 .. caps - 1]+    , let lo = min n (w * per)+    , let hi = min n (lo + per)+    , lo < hi+    ]+  where+    !per = (n + max 1 caps - 1) `div` max 1 caps++-- | Per-partition counts of one row chunk.+histChunk :: VU.Vector Int -> Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)+histChunk hashes p shift lo hi = do+    acc <- VUM.replicate p (0 :: Int)+    let go !i+            | i >= hi = pure acc+            | otherwise = do+                let !pp = partIx shift (VU.unsafeIndex hashes i)+                c <- VUM.unsafeRead acc pp+                VUM.unsafeWrite acc pp (c + 1)+                go (i + 1)+    go lo++{- | Scatter one row chunk into @sortedM@/@sortedHashM@ through the chunk's+private cursor.+-}+scatterChunk ::+    VU.Vector Int ->+    Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    IO ()+scatterChunk hashes shift cursor sortedM sortedHashM lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            let !h = VU.unsafeIndex hashes i+                !pp = partIx shift h+            pos <- VUM.unsafeRead cursor pp+            VUM.unsafeWrite sortedM pos i+            VUM.unsafeWrite sortedHashM pos h+            VUM.unsafeWrite cursor pp (pos + 1)+            go (i + 1)++-------------------------------------------------------------------------------+-- 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 sortedHash eqRow localGid canonBoxes nLocalGroups =+    pooledIndices caps p groupPartition+  where+    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 sortedHash pos+                        (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++-------------------------------------------------------------------------------+-- 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 @(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 per-original-row+inverse, is no longer built here — 'rtgFromVisOffs' derives it on demand, so+aggregations that never read it skip its full random-write pass.)++Each partition owns a disjoint @sortedRows@ range and a disjoint global group-id+range, and its rows are exactly its groups' rows — so its first group's offset is+its own @partStart@ and every pass (group ids, offsets, placement) runs per+partition on parallel workers with no shared writes. @sortedRows@ keeps ascending+original row order inside a partition, so per-group row order matches the+sequential pass exactly.+-}+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)+assemble n p partStart sortedRows localGid globalBase canonOf nGroups = do+    caps <- getNumCapabilities+    gidAt <- VUM.new (max 1 n)+    counts <- VUM.new (max 1 nGroups)+    offsM <- VUM.new (nGroups + 1)+    visM <- VUM.new (max 1 n)+    let doPartition !pp = do+            let !s = VU.unsafeIndex partStart pp+                !e = VU.unsafeIndex partStart (pp + 1)+                !base = VU.unsafeIndex globalBase pp+                !gEnd = VU.unsafeIndex globalBase (pp + 1)+                !canon = V.unsafeIndex canonOf pp+            let zero !g+                    | g >= gEnd = pure ()+                    | otherwise = VUM.unsafeWrite counts g 0 >> zero (g + 1)+            zero base+            -- Pass 1: global group ids and per-group counts.+            let pass1 !pos+                    | pos >= e = pure ()+                    | otherwise = do+                        lg <- VUM.unsafeRead localGid pos+                        let !g = base + VU.unsafeIndex canon lg+                        VUM.unsafeWrite gidAt pos g+                        c <- VUM.unsafeRead counts g+                        VUM.unsafeWrite counts g (c + 1)+                        pass1 (pos + 1)+            pass1 s+            -- Offsets for our group range (they start at our partStart);+            -- counts becomes the per-group write cursor.+            let offsLoop !g !acc+                    | g >= gEnd = pure ()+                    | otherwise = do+                        VUM.unsafeWrite offsM g acc+                        c <- VUM.unsafeRead counts g+                        VUM.unsafeWrite counts g acc+                        offsLoop (g + 1) (acc + c)+            offsLoop base s+            -- Pass 2: stable placement into valueIndices.+            let pass2 !pos+                    | pos >= e = pure ()+                    | otherwise = do+                        g <- VUM.unsafeRead gidAt pos+                        let !row = VU.unsafeIndex sortedRows pos+                        c <- VUM.unsafeRead counts g+                        VUM.unsafeWrite visM c row+                        VUM.unsafeWrite counts g (c + 1)+                        pass2 (pos + 1)+            pass2 s+    pooledIndices caps p doPartition+    VUM.unsafeWrite offsM nGroups n+    offs <- VU.unsafeFreeze offsM+    vis <- VU.unsafeFreeze visM+    pure (vis, offs)++{- | Deferred @rowToGroup@ from @(valueIndices, offsets)@:+@rtg[vis[i]] = g@ for every @i@ in group @g@'s range. @vis@ is a permutation,+so any split of the position space writes disjoint slots; each worker binary+searches its first group and then walks group ranges. Values are identical to+the @rowToGroup@ the assembly pass used to build inline. Pure w.r.t. its+immutable inputs, so the 'unsafePerformIO' is safe.+-}+rtgFromVisOffs :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int+rtgFromVisOffs n vis offs = unsafePerformIO $ do+    let !nGroups = VU.length offs - 1+    rtgM <- VUM.new (max 1 n)+    let+        -- Largest g with offs[g] <= i (offsets are non-decreasing).+        findGroup !i = go2 0 nGroups+          where+            go2 !lo !hi+                | lo >= hi = lo - 1+                | otherwise =+                    let !mid = (lo + hi) `div` 2+                     in if VU.unsafeIndex offs mid <= i+                            then go2 (mid + 1) hi+                            else go2 lo mid+        fill !i !hi !g+            | i >= hi = pure ()+            | otherwise = do+                let !g' = advance g+                    advance !gg =+                        if VU.unsafeIndex offs (gg + 1) <= i+                            then advance (gg + 1)+                            else gg+                VUM.unsafeWrite rtgM (VU.unsafeIndex vis i) g'+                fill (i + 1) hi g'+    parallelChunks_ parThreshold n (\lo hi -> fill lo hi (findGroup lo))+    VU.unsafeFreeze rtgM+{-# NOINLINE rtgFromVisOffs #-}
+ src-internal/DataFrame/Internal/Interpreter.hs view
@@ -0,0 +1,1098 @@+{-# 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.Column.Bitmap+import DataFrame.Internal.DataFrame+import DataFrame.Internal.Expression+import qualified DataFrame.Internal.Grouping as G+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+    #-}+-- toDouble on an Int column (hot path for derived arithmetic)+{-# SPECIALIZE mapColumn ::+    (Int -> Double) -> 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+    MergedColumn _ _ -> sliceGroups (materializeMerged col) os indices+    BoxedColumn bm vec ->+        let !sorted =+                V.generate+                    (VU.length indices)+                    ((vec `V.unsafeIndex`) . (indices `VU.unsafeIndex`))+            !sortedBm = permuteBitmap bm+         in V.generate nGroups $ \i ->+                BoxedColumn+                    (fmap (bitmapSlice (start i) (len i)) sortedBm)+                    (V.unsafeSlice (start i) (len i) sorted)+    UnboxedColumn bm vec ->+        let !sorted = VU.unsafeBackpermute vec indices+            !sortedBm = permuteBitmap bm+         in V.generate nGroups $ \i ->+                UnboxedColumn+                    (fmap (bitmapSlice (start i) (len i)) sortedBm)+                    (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+    permuteBitmap = fmap $ \bm ->+        buildBitmapFromValid $+            VU.map+                (\r -> if bitmapTestBit bm r then 1 else 0)+                indices+{-# 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)+    MergedColumn _ _ -> promoteToDoubleWith onResult (materializeMerged 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)+    MergedColumn _ _ -> promoteToFloatWith onResult (materializeMerged 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)+    MergedColumn _ _ -> promoteToIntWith onResult (materializeMerged 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)+    MergedColumn _ _ -> tryParseWith onResult (materializeMerged 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 (fastUnaryFn @b @a (unaryName op) (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' mappendColumns 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++{- | The op's element function, with a fast path for @toDouble@ (matched by+'unaryName', like the @toDouble@ peeling in "DataFrame.Internal.Simplify").+The closure captured at 'Expr'-construction time is @realToFrac@, which at an+integral source type without a fired rewrite rule lowers to+@fromRational . toRational@ — a 'Rational' allocation plus 'fromRat' per+element. 'fromIntegral' at the concrete type is the same correctly-rounded+conversion, so the swap is bit-identical; only the constant factor changes.+Non-integral sources and other ops keep the stored function.+-}+fastUnaryFn ::+    forall b a.+    (Columnable b, Columnable a) =>+    T.Text -> (b -> a) -> (b -> a)+fastUnaryFn name f+    | name == "toDouble"+    , Just Refl <- testEquality (typeRep @a) (typeRep @Double) =+        integralToDouble @b f+    | otherwise = f++integralToDouble :: forall b. (Columnable b) => (b -> Double) -> b -> Double+integralToDouble f+    | Just Refl <- testEquality rb (typeRep @Int) = fromIntegral+    | Just Refl <- testEquality rb (typeRep @Int8) = fromIntegral+    | Just Refl <- testEquality rb (typeRep @Int16) = fromIntegral+    | Just Refl <- testEquality rb (typeRep @Int32) = fromIntegral+    | Just Refl <- testEquality rb (typeRep @Int64) = fromIntegral+    | Just Refl <- testEquality rb (typeRep @Word) = fromIntegral+    | otherwise = f+  where+    rb = typeRep @b++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
+ src-internal/DataFrame/Internal/Row.hs view
@@ -0,0 +1,217 @@+{-# LANGUAGE ExistentialQuantification #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module DataFrame.Internal.Row 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 Control.Exception (throw)+import Data.Function (on)+import Data.Maybe (catMaybes, fromMaybe, isNothing)+import Data.Type.Equality (TestEquality (..))+import Data.Typeable (Typeable, type (:~:) (..))+import DataFrame.Errors (DataFrameException (..), TypeErrorContext (..))+import DataFrame.Internal.Column (+    Column (..),+    Columnable,+    columnLength,+    fromList,+    fromVector,+    materializeMerged,+    sliceColumn,+ )+import DataFrame.Internal.Column.Bitmap (Bitmap, bitmapTestBit)+import DataFrame.Internal.Data.PackedText (packedIndexText, packedLength)+import DataFrame.Internal.DataFrame (+    DataFrame (columnIndices, dataframeDimensions),+    getColumn,+ )+import DataFrame.Internal.Expression (Expr (..))+import Type.Reflection (TypeRep, typeOf, typeRep)++data Any where+    Value :: (Columnable a) => a -> Any+    -- Saves us the extra indirection we get from making Value (Maybe a)+    -- and having to unpack it again to check for nulls.+    -- Instead, we just have Null as a separate constructor.+    Null :: Any++instance Eq Any where+    (==) :: Any -> Any -> Bool+    (Value a) == (Value b) = fromMaybe False $ do+        Refl <- testEquality (typeOf a) (typeOf b)+        return $ a == b+    Null == Null = True+    _ == _ = False++instance Show Any where+    show :: Any -> String+    show (Value a) = T.unpack (showValue a)+    show Null = "null"++showValue :: forall a. (Columnable a) => a -> T.Text+showValue v = case testEquality (typeRep @a) (typeRep @T.Text) of+    Just Refl -> v+    Nothing -> case testEquality (typeRep @a) (typeRep @String) of+        Just Refl -> T.pack v+        Nothing -> (T.pack . show) v++-- | Wraps a value into an \Any\ type. This helps up represent rows as heterogenous lists.+toAny :: forall a. (Columnable a) => a -> Any+toAny = Value++-- | Unwraps a value from an \Any\ type. A 'Null' cell yields 'Nothing'.+fromAny :: forall a. (Columnable a) => Any -> Maybe a+fromAny Null = Nothing+fromAny (Value (v :: b)) = do+    Refl <- testEquality (typeRep @a) (typeRep @b)+    pure v++{- | Wrap a column cell into an 'Any', honouring the column's null bitmap: a slot+marked invalid becomes 'Null', any other slot becomes a 'Value'. Only needs+@Columnable a@ (the stored element type), not @Columnable (Maybe a)@.+-}+cellAny :: (Columnable a) => Maybe Bitmap -> Int -> a -> Any+cellAny Nothing _ x = Value x+cellAny (Just bm) i x = if bitmapTestBit bm i then Value x else Null++type Row = V.Vector Any++(!?) :: [a] -> Int -> Maybe a+(!?) [] _ = Nothing+(!?) (x : _) 0 = Just x+(!?) (_x : xs) n = (!?) xs (n - 1)++mkColumnFromRow :: T.Text -> Int -> [[Any]] -> Column+mkColumnFromRow name i rows =+    case L.find isValue cells of+        Just (Value (_ :: a)) ->+            let collect _ Null = Nothing :: Maybe a+                collect r (Value (v' :: b)) =+                    case testEquality (typeRep @a) (typeRep @b) of+                        Just Refl -> Just v'+                        Nothing -> throw (mismatchAt r (typeRep @b) (typeRep @a))+                maybes = zipWith collect [0 :: Int ..] cells+             in if any isNothing maybes+                    then fromVector (V.fromList maybes)+                    else fromList (catMaybes maybes)+        _ -> fromVector (V.fromList (map (const (Nothing :: Maybe T.Text)) cells))+  where+    cells = zipWith cellAt [0 :: Int ..] rows+    cellAt r row = fromMaybe (throw (missingCellAt r)) (row !? i)+    isValue (Value _) = True+    isValue Null = False+    mismatchAt ::+        forall x y.+        (Typeable x, Typeable y) => Int -> TypeRep x -> TypeRep y -> DataFrameException+    mismatchAt r actual expected =+        TypeMismatchException+            MkTypeErrorContext+                { userType = Right actual+                , expectedType = Right expected+                , errorColumnName = Just (T.unpack name ++ ", row " ++ show r)+                , callingFunctionName = Just "fromRows"+                }+    missingCellAt r =+        InternalException+            ( "fromRows: row "+                <> T.pack (show r)+                <> " has no cell for column "+                <> name+            )++{- | Convert the whole dataframe to a list of rows, one per row index in natural+order; each row lists all columns ordered by column index. Materializes every+row, so prefer 'toRowVector' for large frames.++>>> toRowList df+[[("name", "Alice"), ("age", 25), ...], [("name", "Bob"), ("age", 30), ...], ...]+-}+toRowList :: DataFrame -> [[(T.Text, Any)]]+toRowList df =+    let+        names = map fst (L.sortBy (compare `on` snd) $ M.toList (columnIndices df))+     in+        map+            (zip names . V.toList . mkRowRep df names)+            [0 .. (fst (dataframeDimensions df) - 1)]++{- | Convert the dataframe to a vector of rows containing only the named columns,+in the given order. An empty name list yields one empty row per dataframe row.++>>> toRowVector ["name", "age"] df+Vector of rows with only name and age fields+-}+toRowVector :: [T.Text] -> DataFrame -> V.Vector Row+toRowVector names df = V.generate (fst (dataframeDimensions df)) (mkRowRep df names)++{- | Given a row gets the value associated with a field.++==== __Examples__++>>> map (rowValue (F.col @Int "age")) (toRowList df)+[25,30, ...]+-}+rowValue :: forall a. Expr a -> [(T.Text, Any)] -> Maybe a+rowValue (Col name) row = lookup name row >>= fromAny @a+rowValue _ _ = error "Can only get rowValue of column reference"++mkRowFromArgs :: [T.Text] -> DataFrame -> Int -> Row+mkRowFromArgs names df i = V.map get (V.fromList names)+  where+    get name = case getColumn name df of+        Nothing ->+            throw $+                ColumnsNotFoundException+                    [name]+                    "[INTERNAL] mkRowFromArgs"+                    (M.keys $ columnIndices df)+        Just (BoxedColumn bm column) -> cellAny bm i (column V.! i)+        Just (UnboxedColumn bm column) -> cellAny bm i (column VU.! i)+        Just (PackedText bm p) -> cellAny bm i (packedIndexText p i)+        Just c@(MergedColumn _ _) ->+            case materializeMerged (sliceColumn i 1 c) of+                BoxedColumn bm column -> cellAny bm 0 (column V.! 0)+                _ -> error "mkRowFromArgs: materializeMerged is boxed"++-- Returns row values in the caller's requested column order, not the+-- dataframe's storage order.+mkRowRep :: DataFrame -> [T.Text] -> Int -> Row+mkRowRep df names i = V.generate (L.length names) (\index -> get (names' V.! index))+  where+    names' = V.fromList names+    throwError name =+        error $+            "Column "+                ++ T.unpack name+                ++ " has less items than "+                ++ "the other columns at index "+                ++ show i+    get name = case getColumn name df of+        Just (BoxedColumn bm c) -> case c V.!? i of+            Just e -> cellAny bm i e+            Nothing -> throwError name+        Just (UnboxedColumn bm c) -> case c VU.!? i of+            Just e -> cellAny bm i e+            Nothing -> throwError name+        Just (PackedText bm p)+            | i < packedLength p -> cellAny bm i (packedIndexText p i)+            | otherwise -> throwError name+        Just c@(MergedColumn _ _)+            | i < columnLength c ->+                case materializeMerged (sliceColumn i 1 c) of+                    BoxedColumn bm column -> cellAny bm 0 (column V.! 0)+                    _ -> error "mkRowRep: materializeMerged is boxed"+            | otherwise -> throwError name+        Nothing ->+            throw $ ColumnsNotFoundException [name] "mkRowRep" (M.keys $ columnIndices df)
+ src-internal/DataFrame/Internal/Row/RowHash.hs view
@@ -0,0 +1,240 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | Row-hash kernels with a parallel driver, feeding grouping and the join+build/probe. Each row's hash depends only on its own bytes, so hashing disjoint+ranges in parallel is race-free and bit-identical to the sequential pass.+-}+module DataFrame.Internal.Row.RowHash (+    computeRowHashesIO,+    computeRowHashesWithIO,+    hashRowRange,+    parRowHashThreshold,+) where++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.Algorithms.Hash (+    fnvOffset,+    mixBytes,+    mixDouble,+    mixInt,+    mixShow,+    mixText,+    nullSalt,+ )+import DataFrame.Internal.Column (+    Column (..),+    materializeMerged,+ )+import DataFrame.Internal.Column.Bitmap (+    Bitmap,+    bitmapTestBit,+ )+import DataFrame.Internal.Column.Types (+    SBool (..),+    sFloating,+    sIntegral,+ )+import DataFrame.Internal.Control.Concurrent (parallelChunks_)+import DataFrame.Internal.Data.PackedText (+    PackedSel,+    PackedTextData (..),+    offAt,+    packedSlice,+    selAt,+ )++{- | At least this many rows make the fork/coordination overhead of the parallel+hash worth it. Below it the sequential single range is used. Matches the+grouping/join parallel thresholds so the whole pipeline switches together.+-}+parRowHashThreshold :: Int+parRowHashThreshold = 200000++{- | Compute the per-row key hash over the selected key columns of an @n@-row+frame. Forks one worker per capability over disjoint row ranges when the row+count justifies it, else hashes the single full range; output is capability-independent.++Dictionary-code hashing is disabled: this is the join entry point, and joins+hash each side by its own representation, so a canonical dict column on one+side of a join against a plain 'T.Text' (or non-canonical packed) column on+the other must byte-hash to keep both sides bucketing identically.+-}+computeRowHashesIO :: Int -> [Column] -> IO (VU.Vector Int)+computeRowHashesIO = computeRowHashesWithIO False++{- | 'computeRowHashesIO' with an explicit dictionary-code switch. When+@useDictCodes@ is 'True', a canonical dict-encoded 'PackedText' column mixes+its 'Int' code per row instead of its byte slice (equal strings share a code,+so bucketing within one frame is preserved). Only sound when every consumer of+the hashes uses the same rule — the grouping path passes 'True', joins 'False'.+-}+computeRowHashesWithIO :: Bool -> Int -> [Column] -> IO (VU.Vector Int)+computeRowHashesWithIO useDictCodes n selected = do+    mv <- VUM.unsafeNew (max 1 n)+    let runRange lo hi = hashRowRange useDictCodes mv lo hi selected+    parallelChunks_ parRowHashThreshold n runRange+    VU.unsafeFreeze (VUM.slice 0 n mv)++{- | Mix every selected column over the row range @[lo, hi)@ into @mv@, seeding+each slot with 'fnvOffset'. Must match the sequential grouping hash byte-for-byte+(at the same @useDictCodes@ setting) so grouping and joins bucket identically.+-}+hashRowRange :: Bool -> VUM.IOVector Int -> Int -> Int -> [Column] -> IO ()+hashRowRange useDictCodes mv lo hi cols = do+    seedRange mv lo hi+    mapM_ (mixColumnRange useDictCodes mv lo hi) cols++seedRange :: VUM.IOVector Int -> Int -> Int -> IO ()+seedRange mv lo hi = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = VUM.unsafeWrite mv i fnvOffset >> go (i + 1)++{- | Fold one column's values over @[lo, hi)@ into the running hashes. The branch+structure mirrors the sequential grouping hash: typed unboxed fast paths, then a+'mixShow' fallback, with the null bitmap mixing 'nullSalt'.+-}+mixColumnRange :: Bool -> VUM.IOVector Int -> Int -> Int -> Column -> IO ()+mixColumnRange useDictCodes mv lo hi = \case+    c@(MergedColumn _ _) -> mixColumnRange useDictCodes mv lo hi (materializeMerged c)+    UnboxedColumn ubm (v :: VU.Vector a) ->+        case testEquality (typeRep @a) (typeRep @Int) of+            Just Refl -> unboxedRange mv lo hi ubm mixInt v+            Nothing ->+                case testEquality (typeRep @a) (typeRep @Double) of+                    Just Refl -> unboxedRange mv lo hi ubm mixDouble v+                    Nothing ->+                        case sIntegral @a of+                            STrue ->+                                unboxedRange mv lo hi ubm (\h d -> mixInt h (fromIntegral @a @Int d)) v+                            SFalse ->+                                case sFloating @a of+                                    STrue ->+                                        unboxedRange mv lo hi ubm (\h d -> mixDouble h (realToFrac d :: Double)) v+                                    SFalse ->+                                        unboxedRange mv lo hi ubm mixShow v+    BoxedColumn bm (v :: V.Vector a) ->+        case testEquality (typeRep @a) (typeRep @T.Text) of+            Just Refl -> boxedRange mv lo hi bm mixText v+            Nothing -> boxedRange mv lo hi bm mixShow v+    PackedText bm p -> packedRange useDictCodes mv lo hi bm p++{- | Mix an unboxed column's range, mixing 'nullSalt' at null slots. @INLINE@d to+specialise on the element type and mixing function per call site.+-}+unboxedRange ::+    (VU.Unbox a) =>+    VUM.IOVector Int ->+    Int ->+    Int ->+    Maybe Bitmap ->+    (Int -> a -> Int) ->+    VU.Vector a ->+    IO ()+unboxedRange mv lo hi ubm mix v = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            h <- VUM.unsafeRead mv i+            let !h' = case ubm of+                    Just bm | not (bitmapTestBit bm i) -> mixInt h nullSalt+                    _ -> mix h (VU.unsafeIndex v i)+            VUM.unsafeWrite mv i h'+            go (i + 1)+{-# INLINE unboxedRange #-}++boxedRange ::+    VUM.IOVector Int ->+    Int ->+    Int ->+    Maybe Bitmap ->+    (Int -> a -> Int) ->+    V.Vector a ->+    IO ()+boxedRange mv lo hi bm mix v = go lo+  where+    go !i+        | i >= hi = pure ()+        | otherwise = do+            h <- VUM.unsafeRead mv i+            let !h' = case bm of+                    Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt+                    _ -> mix h (V.unsafeIndex v i)+            VUM.unsafeWrite mv i h'+            go (i + 1)+{-# INLINE boxedRange #-}++{- | Mix a packed-text column's range over its raw UTF-8 byte slices. The+unselected payload is the hot path (indexes the offset vector directly); a+selected payload (a gather/join result) falls back to 'packedSlice'. When+@useDictCodes@ holds and the selection is a canonical dictionary encoding+(equal strings share a code), each row mixes its 'Int' code with one 'mixInt'+instead of walking the string bytes.+-}+packedRange ::+    Bool ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    Maybe Bitmap ->+    PackedTextData ->+    IO ()+packedRange useDictCodes mv lo hi bm p =+    case ptSel p of+        Just sel | useDictCodes && ptCanonicalSel p -> codes sel+        Nothing -> contiguous (ptBytes p) (ptOffsets p)+        Just _ -> selected+  where+    valid i = case bm of+        Just bm' -> bitmapTestBit bm' i+        Nothing -> True+    contiguous !arr !offs = go lo+      where+        go !i+            | i >= hi = pure ()+            | otherwise = do+                h <- VUM.unsafeRead mv i+                let !o = offAt offs i+                    !l = offAt offs (i + 1) - o+                    !h' = if valid i then mixBytes h arr o l else mixInt h nullSalt+                VUM.unsafeWrite mv i h'+                go (i + 1)+    selected = go lo+      where+        go !i+            | i >= hi = pure ()+            | otherwise = do+                h <- VUM.unsafeRead mv i+                let !h' =+                        if valid i+                            then let (arr, o, l) = packedSlice p i in mixBytes h arr o l+                            else mixInt h nullSalt+                VUM.unsafeWrite mv i h'+                go (i + 1)+    codes :: PackedSel -> IO ()+    codes !sel = go lo+      where+        go !i+            | i >= hi = pure ()+            | otherwise = do+                h <- VUM.unsafeRead mv i+                let !h' =+                        if valid i+                            then mixInt h (selAt sel i)+                            else mixInt h nullSalt+                VUM.unsafeWrite mv i h'+                go (i + 1)+{-# INLINE packedRange #-}
+ src/DataFrame/Core.hs view
@@ -0,0 +1,110 @@+{- | The curated public surface of @dataframe-core@: the interchange types+('DataFrame', 'Column', 'Row', 'Expr'), element constraints, and the+rendering/serialization verbs. Internal plumbing stays in @dataframe-core:internal@.+-}+module DataFrame.Core (+    -- * The DataFrame+    DataFrame,+    GroupedDataFrame,+    empty,+    fromNamedColumns,+    insertColumn,+    columnNames,+    null,++    -- * Columns+    Column,+    fromList,+    fromVector,+    fromUnboxedVector,+    mkRandom,+    toList,+    toVector,+    hasElemType,+    hasMissing,+    isNumeric,++    -- * Element constraints+    Columnable,+    Columnable',++    -- * Rows+    Row,+    Any,+    toAny,+    fromAny,+    rowValue,+    toRowList,+    toRowVector,++    -- * Expressions+    Expr,+    NamedExpr,+    toNamedExpr,+    toUExpr,+    fromUExpr,+    eSize,+    prettyPrint,+    prettyPrintWidth,++    -- * Rendering & serialization+    TruncateConfig (..),+    defaultTruncateConfig,+    toCsv,+    toCsv',+    toSeparated,+    toMarkdown,+    toMarkdown',+) where++import Prelude hiding (null)++import DataFrame.Internal.Column (+    Column,+    Columnable,+    fromList,+    fromUnboxedVector,+    fromVector,+    hasElemType,+    hasMissing,+    isNumeric,+    mkRandom,+    toList,+    toVector,+ )+import DataFrame.Internal.Column.Types (Columnable')+import DataFrame.Internal.DataFrame (+    DataFrame,+    GroupedDataFrame,+    TruncateConfig (..),+    columnNames,+    defaultTruncateConfig,+    empty,+    fromNamedColumns,+    insertColumn,+    null,+    toCsv,+    toCsv',+    toMarkdown,+    toMarkdown',+    toSeparated,+ )+import DataFrame.Internal.Expression (+    Expr,+    NamedExpr,+    eSize,+    fromUExpr,+    prettyPrint,+    prettyPrintWidth,+    toNamedExpr,+    toUExpr,+ )+import DataFrame.Internal.Row (+    Any,+    Row,+    fromAny,+    rowValue,+    toAny,+    toRowList,+    toRowVector,+ )
− src/DataFrame/Display/Terminal/Colours.hs
@@ -1,7 +0,0 @@-module DataFrame.Display.Terminal.Colours where--red, green, brightGreen, brightBlue :: String -> String-red = id-green = id-brightGreen = id-brightBlue = id
− src/DataFrame/Display/Terminal/PrettyPrint.hs
@@ -1,118 +0,0 @@-{-# 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.---- Adapted from: https://stackoverflow.com/questions/5929377/format-list-output-in-haskell--- a type for fill functions-type Filler = Int -> T.Text -> T.Text---- a type for describing table columns-data ColDesc t = ColDesc-    { colTitleFill :: Filler-    , colTitle :: T.Text-    , colValueFill :: Filler-    }---- functions that fill a string (s) to a given width (n) by adding pad--- character (c) to align 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---- functions that 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-        -- In a GitHub pipe table a literal '|' ends the cell and a newline ends-        -- the row, so a value like the operator name <|> would split the row into-        -- the wrong columns. Escape both for the Markdown format only.-        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 "|" "\\|"
− src/DataFrame/Errors.hs
@@ -1,187 +0,0 @@-{-# 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)-            ]
+ src/DataFrame/Expression/Operators.hs view
@@ -0,0 +1,12 @@+{- |+The public name for dataframe's expression operators.++This is a re-export shim: the implementation lives in+"DataFrame.Internal.Expression.Operators", which is internal and may be+reorganised without notice. Depend on this module instead.+-}+module DataFrame.Expression.Operators (+    module DataFrame.Internal.Expression.Operators,+) where++import DataFrame.Internal.Expression.Operators
− src/DataFrame/Internal/AggKernel.hs
@@ -1,308 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE LambdaCase #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--{- | Vectorized scatter-accumulate aggregation kernel.--The grouping layer ('DataFrame.Internal.Grouping') hands us a dense-@rowToGroup@ vector (group id per row, canonical order) plus the number of-groups. For the common reductions this kernel replaces the per-group boxed-expression-interpreter fold with a single unboxed linear pass that scatters-each row's value into primitive per-group accumulator arrays indexed by the-group id. No boxed accumulator record, no per-element dictionary closure: the-element type is resolved once per column by a 'typeRep' switch and the inner-loop is a monomorphic primop on a 'VU.Vector'.--Result columns are length @nGroups@ in canonical group order, so they line up-with the key columns 'aggregate' gathers with @selectIndices@.--The kernel is strictly a FAST PATH: the matcher 'DataFrame.Internal.AggPlan.planAgg'-recognises a small set of expression shapes; anything it does not recognise-keeps the existing interpreter, so the general @aggregate@ API stays correct for-arbitrary expressions.--}-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)------------------------------------------------------------------------------------ Column extraction----------------------------------------------------------------------------------{- | 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------------------------------------------------------------------------------------ Scatter reductions----------------------------------------------------------------------------------{- | Run one fast-path reduction. Returns 'Nothing' when the value column is-not a non-null unboxed Int/Double column (then the caller falls back).--}-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))--{- | The monomorphic reduction body. @count@ always yields an Int column;-@sum@/@min@/@max@ preserve the element type; @mean@/@std@/@var@ produce Double;-@top2Sum@ produces Double.--}-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 #-}--{- | One pass filling running sum and count from value column @v@ over groups-@g@ into the supplied accumulator arrays.--}-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--{- | Sample variance (or its square root for sd) via a per-group Welford-recurrence, scattered into three unboxed arrays @(count, mean, m2)@. This is the-same numerically-stable update as the interpreter's @varianceStep@, so the-result is byte-identical to the existing CollectAgg path (and the db-benchmark-checksum is unchanged); the win is the unboxed scatter replacing the per-group-boxed @VarAcc@ fold over a materialized slice. Degenerate groups (n < 2) yield-0, matching @computeVariance@.--}-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 #-}--{- | Per-group sum of the two largest values: a 2-slot scatter holding the-running first and second maximum, then @m1 + m2@. Matches the @take 2 . sortBy-(flip compare)@ definition used by the benchmark's @top2Sum@.--}-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 #-}
− src/DataFrame/Internal/AggKernelDirect.hs
@@ -1,373 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--{- | Low-cardinality DIRECT-INDEXED accumulator fast path (research #4).--When the group-by key resolves to a small dense domain (@nGroups@ below-'directThreshold'), the dense @rowToGroup@ vector handed down by the grouping-layer already IS the group id per row, so we can bypass the @valueIndices@-gather entirely: scan @rowToGroup@ and the value column /in lockstep, in-original row order/, scattering each value straight into a per-group accumulator-array indexed by the dense id. Both arrays are read sequentially (no random-gather through @valueIndices@), and for a small domain the accumulator stays-cache-resident.--For parallelism we use the two-phase morsel pattern (#2): split the ROW range-into one contiguous chunk per capability, give each worker a private tiny-accumulator array, scan its chunk sequentially, then merge the @caps@ partials in-a cheap O(caps * nGroups) pass. This makes few-group questions scale (the work-per worker is a tight sequential pass) instead of being dominated by parallel-fan-out, which is exactly the regression the group-range gather suffered on Q4.--CRITICAL — byte-identical at any @-N@: the two-phase merge only changes the-fold/combine order, so it is admitted ONLY for reductions whose result is-independent of that order: integer sum (exact, associative), count, min, max,-and integer mean (an exact integer sum and count divided once at finalize). The-order-sensitive float reductions (Double sum/mean, variance, sd, top2) are NOT-routed here — the caller keeps the order-preserving group-range kernel for them,-so the db-benchmark checksums stay byte-identical between -N1 and -N8.--}-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. The db-benchmark low-cardinality questions (id1=100, id4=100, id6=1e5,-id2:id4=1e4) sit at or below it; wider domains keep the group-range kernel. The-admitted reductions are order-independent, so the per-worker accumulator merge is-exact regardless of size.--}-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. Returns-'Nothing' (so the caller falls back to the order-preserving kernel) unless BOTH:-(a) the reduction's result is order-independent at this element type — see the-module note — and (b) the value column is a clean unboxed Int/Double column.--@g@ is the dense @rowToGroup@ vector (group id per row, original row order);-@nGroups@ is the dense domain size, already verified @<= directThreshold@ by the-caller.--}-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 @[lo, hi)@ of-@[0, n)@, balanced evenly by row count; each worker @w@ runs @fill lo hi@-producing its OWN private accumulator (thread-local, no shared array, no sync).-Returns the partials in worker order for the caller's merge. Rethrows the first-worker 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
− src/DataFrame/Internal/AggKernelPar.hs
@@ -1,454 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--{- | Parallel scatter-accumulate aggregation kernel.--This builds on the sequential kernel ('DataFrame.Internal.AggKernel') and the-Round-5 grouping layout: 'groupBy' hands us @valueIndices@ (rows ordered by-group) and @offsets@ (per-group boundaries), so each group's rows are a-contiguous run @valueIndices[offsets[g] .. offsets[g+1])@ in their original row-order.--The parallel driver splits the dense group-id range @[0, nGroups)@ into one-contiguous chunk per capability, balanced by row count. Because group ranges are-disjoint and their @valueIndices@ runs are disjoint, every worker reads and-writes its own slice of the shared output array(s) — there is NO cross-worker-overlap and NO merge. And because each group's rows are visited in the same-original-row order as the sequential @rowToGroup@ scan, the per-group fold order-is unchanged, so the result is /byte-identical/ to the sequential kernel at any-@-N@ (no fold-order drift, even for the float sums).--Forks plain 'forkIO' workers (no sparks), one per capability; falls back to the-sequential 'DataFrame.Internal.AggKernel' path when @caps == 1@ or the row count-is below 'parThreshold'.--}-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,- )------------------------------------------------------------------------------------ Parallelisation policy----------------------------------------------------------------------------------{- | Below this many rows the fork overhead is not worth it; the sequential-kernel runs instead. Mirrors 'DataFrame.Internal.GroupingPar.parThreshold'.--}-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------------------------------------------------------------------------------------ Group-range partitioning by row count----------------------------------------------------------------------------------{- | Split @[0, nGroups)@ into @caps@ contiguous group ranges, each holding a-near-equal share of rows. Returns @caps + 1@ boundaries @b@ with @b[0] == 0@ and-@b[caps] == nGroups@; worker @w@ owns groups @[b[w], b[w+1])@. A row-balanced-split keeps skew low even when group sizes vary wildly.--}-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)-        -- First group whose start offset reaches @target@ (>= prev), or nGroups.-        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--{- | Fork @caps@ workers, worker @w@ running @act (b[w]) (b[w+1])@ over its-disjoint group range; join and rethrow the first failure. Sequential when-@caps == 1@.--}-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------------------------------------------------------------------------------------ Parallel single-column reductions----------------------------------------------------------------------------------{- | Parallel counterpart of 'scatterReduce'. Returns 'Nothing' on the same-columns the sequential kernel rejects (boxed/nullable/non-Int-Double); on the-sequential path or tiny inputs it delegates to 'scatterReduce'. The result is-byte-identical to 'scatterReduce' (same per-group fold order).--}-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 #-}--{- | Reconstruct @rowToGroup@ from the group layout for the sequential delegate.-Only used on the small/-N1 path, so the extra pass is negligible.--}-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))--{- | The monomorphic parallel reduction body. Each scatter allocates its full-@nGroups@ output array(s) once, then workers fill disjoint group ranges in-parallel. The result type follows the sequential kernel exactly.--}-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 #-}--{- | Per-group Welford variance/sd, parallel by group range. The recurrence is-applied in original-row order within each group, identical to the sequential-@varScatter@, so the result is byte-identical (no parallel-combine needed: each-group lives wholly inside one worker's range).--}-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 in original-row order inside one worker's-range. Byte-identical to 'momentScatter'; delegates to it on the sequential-path. Returns '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)
− src/DataFrame/Internal/AggPlan.hs
@@ -1,316 +0,0 @@-{-# 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' inspects a named output expression and, on a recognised shape over a-clean (non-null, unboxed Int/Double) column, returns the 'AggPlan' the caller-runs through the scatter kernel ('DataFrame.Internal.AggKernel'); anything it-does not recognise returns 'Nothing' so the caller keeps the existing-interpreter. Recognition is by the 'AggStrategy' name tag plus the shape of the-inner 'Expr' — no new constructor is needed, and the general @aggregate@ API is-unchanged.--'momentScatter' fuses the additive moment sums of two columns (count, Sx, Sy,-Sxx, Syy, Sxy) into one pass — the sufficient statistics for the Q9 regression-family. It is exposed for callers that want to collapse the six separate folds-into a single pass.--}-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. On a recognised shape over a present-clean column return @Just plan@; otherwise 'Nothing'. Nullable (bitmap) or-non-Int/Double value columns are rejected here so the scatter only ever 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 together form the sufficient statistics of two base columns @x@ and @y@.-The caller runs the fused 'momentScatter'/'momentScatterPar' once over-@(colX, colY)@ and binds each output name to the named 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 multi-aggregate 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 base columns @x@ and @y@ (after resolving derived product columns-through @derivingExpressions@). Returns 'Nothing' on any other set so the caller-falls back to the per-expression planner. Both base columns must be clean-unboxed Int/Double, the same gate the single-column scatter uses.--}-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 the five sums. Collapses the Q9 regression family's six independent folds-(and three derive passes) into a single fused pass. Returns '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
− src/DataFrame/Internal/Column.hs
@@ -1,1852 +0,0 @@-{-# 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--{- | Our representation of a column is a GADT that can store data based on the underlying data.--This allows us to pattern match on data kinds and limit some operations to only some-kinds of vectors. Nullability is represented via an optional bit-packed 'Bitmap':-@Nothing@ = no nulls; @Just bm@ = bit @i@ of @bm@ is 1 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: a shared UTF-8 byte buffer + (n+1) row offsets + an-    -- optional validity bitmap. No per-row Text header; Text is produced only-    -- on demand. Behaves as a column of Text (or Maybe Text when a bitmap is-    -- present). Only the CSV ingest path 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 =-        -- byte-aligned: simple slice; clamp so we never ask for more bytes than exist-        let startByte = start `shiftR` 3-            bytes = min ((len + 7) `shiftR` 3) (VU.length bm - startByte)-         in VU.slice startByte bytes bm-    | otherwise =-        -- non-aligned: unpack bit-by-bit and repack-        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)@.-When @a@ is unboxable, creates an 'UnboxedColumn' (more compact).-Otherwise creates a 'BoxedColumn'.-Always attaches a bitmap so the column is recognized as nullable even when-no 'Nothing' values are present (preserves the Maybe type marker).--}-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---- | Materialize an element from a column at index @i@, 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--{- | The universal cold fallback: decode a 'PackedText' into a-@BoxedColumn Text@ using the exact 'sliceTextVector' path the boxed-Text-builder used, so the result is 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)-    -- Gathered/selected payload: rows are non-contiguous, decode per row.-    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 TypedColumn is a wrapper around our type-erased column.-It is used to type check expressions on columns.--Note: there is no guarantee that the Phanton type is the-same as the underlying vector 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--{- | Checks if a column is of a given type values.-For nullable columns (@BoxedColumn (Just _)@ or @UnboxedColumn (Just _)@),-also returns @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-    -- Direct type match-    directMatch :: forall (b :: Type). TypeRep b -> Bool-    directMatch = isJust . testEquality (typeRep @a)-    -- For a nullable column (has bitmap), also accept a = Maybe b-    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-        -- user maps over Maybe a (nullable column as Maybe)-        Just Refl ->-            let !n = VB.length col-             in -- Build result directly without intermediate Maybe vector to avoid-                -- fusion forcing null slots via VU.convert.-                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 ->-                -- user maps over inner type a; preserve bitmap-                Right $ case sUnbox @c of-                    STrue -> UnboxedColumn bm (VU.generate (VB.length col) (f . VB.unsafeIndex col))-                    SFalse -> 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 -> 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) =-    -- Slice-preserving: share the byte buffer, permute via a selection vector-    -- instead of materializing boxed Text. Bitmap follows the same gather.-    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 ->-                -- Slice-preserving sentinel gather: share the buffer, permute-                -- offsets (negative sentinel -> empty slice), build the null-                -- bitmap so -1 rows read as null in left/outer joins.-                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 index for row i.-Avoids random column reads; random writes go to the accumulator array which is-small (nGroups entries) and typically cache-resident.-When @acc@ is unboxable, uses an unboxed mutable vector for the accumulator-array, eliminating pointer indirection on every read/write.--}-foldLinearGroups ::-    forall b acc.-    (Columnable b, Columnable acc) =>-    (acc -> b -> acc) ->-    acc ->-    Column ->-    VU.Vector Int -> -- rowToGroup (length n)-    Int -> -- nGroups-    Either DataFrameException Column-foldLinearGroups f seed col rowToGroup nGroups-    | nGroups == 0 = Right (fromVector @acc VB.empty)-    | otherwise = case col of-        UnboxedColumn _ (vec :: VU.Vector d) -> 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-                        }--    -- \| Allocate accumulators, run the traversal, return a frozen Column.-    -- When @acc@ is unboxable, uses an unboxed mutable vector (no pointer-    -- indirection per read/write) and returns UnboxedColumn directly —-    -- avoiding a round-trip through VB.Vector.-    runWith :: ((Int -> IO acc) -> (Int -> acc -> IO ()) -> IO ()) -> IO Column-    runWith body = case sUnbox @acc of-        STrue -> do-            accs <- VUM.replicate nGroups seed-            body (VUM.unsafeRead accs) (VUM.unsafeWrite accs)-            UnboxedColumn 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-            -- Fast path: both plain unboxed, no bitmaps involved in the output type-            | 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)-        -- Type mismatch or bitmap involvement: fall through to general toVector path-        _ -> 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))-            -- pad data with default (undefined slot, protected by bitmap)-            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-                        }-                    )--{- | Concatenates two columns.--Works similar to 'concatColumns', but unlike that function, it will also combine columns of different types-by wrapping the values in an Either.--E.g. combining Column containing [1,2] with Column containing ["a","b"]-will result in a Column containing [Left 1, Left 2, Right "a", Right "b"].--}--{- | O(n) Concatenate a list of same-type columns in a single allocation.-All columns must have the same constructor and element type (as they will-within a single Parquet column). Calls 'error' on mismatch.--}-concatManyColumns :: [Column] -> Column-concatManyColumns [] = fromList ([] :: [Maybe Int])-concatManyColumns [c] = c-concatManyColumns 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--{- | Converts a column to a vector of a specific type.--This is a type-safe conversion that requires the column's element type-to exactly match the requested type. You must specify the desired type-via type applications.--==== __Type Parameters__--[@a@] The element type to convert to-[@v@] The vector type (e.g., 'VU.Vector', 'VB.Vector')--==== __Examples__-->>> toVector @Int @VU.Vector column-Right (unboxed vector of Ints)-->>> toVector @Text @VB.Vector column-Right (boxed vector of Text)--==== __Returns__--* 'Right' - The converted vector if types match-* 'Left' 'TypeMismatchException' - If the column's type doesn't match the requested type--==== __See also__--For numeric conversions with automatic type coercion, see 'toDoubleVector',-'toFloatVector', and 'toIntVector'.--}-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.--{- | Converts a column to an unboxed vector of 'Double' values.--This function performs intelligent type coercion for numeric types:--* If the column is already 'Double', returns it directly-* If the column contains other floating-point types, converts via 'realToFrac'-* If the column contains integral types, converts via 'fromIntegral' (beware of overflow if the type is `Integer`).--==== __Optional column handling__--For 'OptionalColumn' types, 'Nothing' values are converted to @NaN@ (Not a Number).-This allows optional numeric data to be represented in the resulting vector.--==== __Returns__--* 'Right' - The converted 'Double' vector-* 'Left' 'TypeMismatchException' - If 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-                            }-                        )--{- | Converts a column to an unboxed vector of 'Float' values.--This function performs intelligent type coercion for numeric types:--* If the column is already 'Float', returns it directly-* If the column contains other floating-point types, converts via 'realToFrac'-* If the column contains integral types, converts via 'fromIntegral'-* If the column is boxed 'Integer', converts via 'fromIntegral' (beware of overflow for 64-bit integers and `Integer`)--==== __Optional column handling__--For 'OptionalColumn' types, 'Nothing' values are converted to @NaN@ (Not a Number).-This allows optional numeric data to be represented in the resulting vector.--==== __Returns__--* 'Right' - The converted 'Float' vector-* 'Left' 'TypeMismatchException' - If the column is not numeric--==== __Precision warning__--Converting from 'Double' to 'Float' may result in loss of 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-                            }-                        )--{- | Converts a column to an unboxed vector of 'Int' values.--This function performs intelligent type coercion for numeric types:--* If the column is already 'Int', returns it directly-* If the column contains floating-point types, rounds via 'round' and converts-* If the column contains other integral types, converts via 'fromIntegral'-* If the column is boxed 'Integer', converts via 'fromIntegral'--==== __Returns__--* 'Right' - The converted 'Int' vector-* 'Left' 'TypeMismatchException' - If the column is not numeric--==== __Note__--Unlike 'toDoubleVector' and 'toFloatVector', this function does NOT support-'OptionalColumn'. Optional columns must be handled separately.--==== __Rounding behavior__--Floating-point values are rounded to the nearest integer using 'round'.-For example: 2.5 rounds to 2, 3.5 rounds to 4 (banker's rounding).--}-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 @Int)-                            , expectedType = Right (typeRep @a)-                            , 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))
− src/DataFrame/Internal/ColumnBuilder.hs
@@ -1,299 +0,0 @@-{-# 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]
− src/DataFrame/Internal/ColumnMerge.hs
@@ -1,179 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--{- | Concatenation of per-chunk 'Column's (e.g. from parallel CSV chunks),-re-exported through 'DataFrame.Internal.ColumnBuilder'. Text columns can-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 time). Not yet called by any reader (a later ingest stage flips-'mergeTextChunks' to use it).--}-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 merged shared buffer + offsets-as 'PackedText' (no per-row 'Data.Text.Text' header, no eager UTF-8-validation pass — decode is deferred to the last mile).--}-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
− src/DataFrame/Internal/DataFrame.hs
@@ -1,378 +0,0 @@-{-# 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-    {- ^ Our main data structure stores a dataframe as-    a vector of columns. This improv-    -}-    , 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. A non-positive value on-any field means \"no limit\" on that axis.--* 'maxRows' — render at most this many rows from the top of the frame.-* 'maxColumns' — when the frame has more columns than this, the middle columns-  are collapsed into a single ellipsis column.-* 'maxCellWidth' — text in any individual cell (including headers and type-  rows) longer than this is truncated with a trailing ellipsis.--}-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--        -- Separate out cases dynamically so we don't end up making round trip-        -- string copies.-        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 (if anywhere)-to splice in the ellipsis column. The split puts the extra column on the-left for odd 'maxColumns'; the ellipsis is only inserted when it actually-saves space (i.e. the frame has more than 'maxColumns' + 1 columns).--}-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--{- | Retrieves a column by name from the dataframe, throwing an exception if not found.--This is an unsafe version of 'getColumn' that throws 'ColumnsNotFoundException'-if the column does not exist. Use this when you are certain the column exists.--==== __Throws__--* 'ColumnsNotFoundException' - if the column with the given name does not 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 "
− src/DataFrame/Internal/DictEncode.hs
@@ -1,159 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--{- | Dictionary-encode a text (or factor) group key to dense @Int@ codes-(research #4 / #5).--A text group key is hashed and bucketed exactly as the grouping hash table does,-but instead of carrying the full grouping output it only assigns each row a dense-first-appearance code @0 .. card-1@ (a NULL row gets its own reserved code) and-reports the cardinality. The intent was to feed those codes to the int fast path-(direct-indexed low-card, or a packed composite key) so a text group key reaches-it.--VERDICT (this round): routing through the codes profiled SLOWER than the existing-hash group-by on EVERY db-benchmark group-by question, so-'DataFrame.Internal.Grouping' does not take the dict path (see 'tryDictGroup'-there). The dict-build is its own hash pass over every row, and the hash-group-by already fuses hashing and grouping into one pass; substituting int codes-adds the encode pass without removing the dominant grouping work. Single low-card-(Q1 id1), single high-card (Q3/Q7 id3) and the composites (Q2 id1:id2, Q10 six-keys) were each measured and all lost.--This module remains a correct, unit-tested building block: it produces the codes-and the cardinality only; the routing decision lives in-'DataFrame.Internal.Grouping'.--}-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 dict-build pass is pure-overhead, so the caller keeps the plain hash group-by. Matches the direct-group-histogram budget.--}-dictMaxCardinality :: Int-dictMaxCardinality = 1048576--{- | Dictionary-encode a text-like column to dense first-appearance @Int@ codes.--Returns @Just (codes, cardinality)@ where @codes!i@ is the dense id of row @i@'s-value (a NULL row, when the column is nullable, is assigned its own reserved code-distinct from every present value) and @cardinality@ is the number of distinct-codes used. Returns 'Nothing' for any non-text column or when the cardinality-exceeds 'dictMaxCardinality' (so the caller falls back to the hash group-by).--Only 'PackedText' and boxed 'Data.Text.Text' columns are encoded; everything else-is 'Nothing'.--}-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@. The early bail lets a low-cardinality-PROBE (the single-key direct path) avoid a full high-cardinality pass when the-column turns out to be high-card.--}-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. Hashes each row's raw UTF-8 bytes (the same-'mixBytes' the grouping hash uses) and re-verifies byte equality on collisions,-assigning dense codes in first-appearance order. A null row hashes 'nullSalt' and-re-verifies as equal-to-null only.--}-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. Buckets every row through an-open-addressing table on its precomputed hash, re-verifying the real value with-@eqAt@ on a hash hit, assigning dense first-appearance codes. Bails to 'Nothing'-the moment the distinct count would exceed 'dictMaxCardinality'.--}-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))
− src/DataFrame/Internal/Expression.hs
@@ -1,472 +0,0 @@-{-# 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.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 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 @instance@s. The-generic 'UnUDF'/'BinUDF' carriers cover UDFs, dynamic-named, and arithmetic ops.-Method names match the carrier record fields ('NoFieldSelectors' frees them), so-existing construction and read sites are unchanged.--}-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 -- Swap to canonical order-                            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 column references using a map from column name to-replacement expression. Unlike folding 'replaceExpr' over the bindings, this is a-single parallel pass: every 'Col' reference is resolved against the original map,-so a swap such as @{a ↦ col b, b ↦ col a}@ is handled correctly (sequential-replacement would collapse both columns onto one).--Only 'Col' references are substituted; raw-text references inside 'CastWith' and-'Over' partition keys are left untouched (documented limitation — the fitted ML-transforms that use this only ever emit @Col@/@Lit@/@Unary@/@Binary@/@If@). A-binding whose replacement type does not match the referenced column's type is a-programmer error and raises an exception.--}-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--prettyPrint :: Expr a -> String-prettyPrint = go 0 0-  where-    indent :: Int -> String-    indent n = replicate (n * 2) ' '--    go :: Int -> Int -> Expr a -> String-    go depth prec expr = case expr of-        Col name -> T.unpack name-        CastWith name _ _ -> T.unpack name-        CastExprWith tag _ inner -> T.unpack tag ++ "(" ++ go depth 0 inner ++ ")"-        Lit value -> show value-        If cond t e ->-            let inner =-                    "if "-                        ++ go (depth + 1) 0 cond-                        ++ "\n"-                        ++ indent (depth + 1)-                        ++ "then "-                        ++ go (depth + 1) 0 t-                        ++ "\n"-                        ++ indent (depth + 1)-                        ++ "else "-                        ++ go (depth + 1) 0 e-             in if prec > 0 then "(" ++ inner ++ ")" else inner-        Unary op arg -> case unarySymbol op of-            Nothing -> T.unpack (unaryName op) ++ "(" ++ go depth 0 arg ++ ")"-            Just sym -> T.unpack sym ++ "(" ++ go depth 0 arg ++ ")"-        Binary op l r ->-            let p = binaryPrecedence op-                inner = case binarySymbol op of-                    Just name -> go depth p l ++ " " ++ T.unpack name ++ " " ++ go depth p r-                    Nothing ->-                        T.unpack (binaryName op) ++ "(" ++ go depth p l ++ ", " ++ go depth p r ++ ")"-             in if prec > p then "(" ++ inner ++ ")" else inner-        Agg (CollectAgg op _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"-        Agg (FoldAgg op _ _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"-        Agg (MergeAgg op _ _ _ _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"-        Over keys inner -> go depth 0 inner ++ ".over(" ++ show (map T.unpack keys) ++ ")"
− src/DataFrame/Internal/Grouping.hs
@@ -1,454 +0,0 @@-{-# 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) groups the dataframe by the given rows aggregating the remaining rows-into vector that should be reduced later.--Rows are bucketed with an unboxed open-addressing hash table-('DataFrame.Internal.HashTable') that maps each row's key-hash to a dense group-id, re-verifying the real key columns on every hash hit. This both cuts the-per-row boxed allocation of the previous 'Data.IntMap' bucketing (less GC) and-fixes a latent collision bug where two distinct keys sharing a hash were merged.-Groups are numbered in first-appearance order; 'valueIndices' / 'offsets' are-then derived by a stable counting sort on the group id.--}-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--{- | The low-cardinality direct-indexed grouping fast path-('DataFrame.Internal.GroupingDirect'). Fires only for a SINGLE clean small-range-@Int@ key column; one shared function feeds both -N1 and -N8 so the result is-identical at any capability count (parallel==sequential by construction). Returns-'Nothing' on any other key shape, falling through to the hash group-by.--}-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 (the codes ARE the-first-appearance group ids), then derive @valueIndices@/@offsets@ from those-codes with one counting sort.--PROFILED AS A LOSS, so this currently always falls back ('Nothing'). The-dict-build is its own hash pass over every row, and the hash group-by already-hashes and groups in one fused pass: at -N1 the dict path measured ~0.44s vs-~0.33s for the hash path on Q1 (id1, 100 groups) at 1e7 rows, and at -N8 the-parallel partitioned grouping is far faster than any sequential dict-build. The-high-card single keys (Q3/Q7 id3 ~1e5) and the multi-key composites (Q2 id1:id2,-Q10 six keys) were each tried and also lost — substituting int codes does not-remove the dominant grouping passes, it adds the encode passes on top. The-'DataFrame.Internal.DictEncode.dictEncodeColumnUpTo' step is kept and unit-tested-as a correct building block; the routing here is deliberately disabled. The-@_n@/@_df@/@_names@/@_col@ wiring is retained so re-enabling is a one-line change-should a parallel dict-encode ever change the verdict.--}-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 the single-key dict-encode probe would use: it bails to-'Nothing' once the distinct count passes this, so a high-card key never pays for-a full encode pass. 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; produces an-output bit-for-bit identical to 'groupBySeq'. Pure via 'unsafePerformIO' (the IO-is 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 via the open-addressing table, in-first-appearance order. Returns @(rowToGroup, repHash, repRow)@ where @repHash@ /-@repRow@ are the hash and representative row index of each group (indexed by the-first-appearance id). The table re-verifies the real key with 'eqKeyRow' on each-hash hit so colliding keys are kept 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-    -- At most n groups; trimmed to the actual count on freeze.-    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 are ordered-by ascending representative hash, tie-broken by representative row index. This-makes the emitted group order a deterministic function of the key set (not of-input row order), so set operations like @union a b@ and @union b a@ agree, and-reproduces the ascending-hash order of the previous 'Data.IntMap' grouping.-Returns @remap@ with @remap[firstAppearanceId] = canonicalId@.--The ordering is the stable hash-rank of 'DataFrame.Internal.RadixRank': @repRow@-is strictly ascending in first-appearance id order (a new group's representative-is the first row that reaches it, scanned in increasing row index), so the stable-sort's equal-hash tie-break reproduces the old @(hash, repRow)@ comparison. O(g)-with no boxed-tuple comparison sort — this keeps the @1e7@-distinct-group case-(Q10) off an @n log n@ list 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@ are equal across all key-columns, comparing validity bits first (a null equals only another null) then-the underlying value. Used by the hash table 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[g] = size of group g (g in [0, nGroups)). Slot nGroups stays 0 so-    -- the exclusive scan below lands n in offsets[nGroups].-    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-    -- Exclusive prefix scan of counts into offsets: offsets[k] is the start of-    -- group k and offsets[nGroups] == n.-    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-    -- 'counts' is repurposed as a per-group write cursor seeded at each group's-    -- start offset, giving a stable placement (rows keep original order).-    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)
− src/DataFrame/Internal/GroupingDirect.hs
@@ -1,260 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--{- | Low-cardinality DIRECT-INDEXED grouping fast path (research #4).--When the group-by key is a single clean (non-null) unboxed @Int@ column whose-value /range/ is small (@max - min + 1 <= directGroupThreshold@), the value-itself indexes a dense accumulator: there is no hashing, no key re-verification,-and no open-addressing probe. This is the X100 / ClickHouse FixedHashMap idea-applied to the grouping step — the dominant cost of the low-cardinality-db-benchmark questions (id4=100, id6=1e5) is the hash group-by, not the-aggregate scatter, so bypassing the hash here is the real lever.--The pipeline is three linear passes plus an O(range) compaction:--  1. min/max of the key (parallel range-reduce, order-independent).-  2. a per-value histogram (parallel per-thread histograms, exact-integer merge).-  3. compact non-empty values into dense ids in ASCENDING value order, then a-     stable placement pass building @valueIndices@.--The emitted group order is ascending key value — a deterministic function of the-key set (like the hash path's canonical order), so set operations stay-commutative and the db-benchmark checksums (order-independent sums) are-unchanged. Crucially this single function is shared by both the sequential and-parallel 'groupBy' entry points, so the parallel==sequential parity is automatic-(identical output by construction) at any @-N@.--}-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, an exclusive-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-    -- 1. Per-value histogram over the dense value index (parallel, exact).-    hist <- buildHistogram v mn range n-    -- 2. Compact non-empty values -> dense group ids (ascending value order),-    -- recording each value's group id and the group's row count.-    valToGroup <- VUM.replicate range (-1 :: Int)-    grpCount <- VUM.new range-    nGroups <- compact hist range valToGroup grpCount-    -- 3. Exclusive prefix scan of group counts -> offsets (length nGroups + 1).-    offsM <- VUM.new (nGroups + 1)-    cursor <- VUM.new nGroups-    scanOffsets grpCount nGroups offsM cursor-    -- 4. Stable placement: rowToGroup[i] and valueIndices in group order.-    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)
− src/DataFrame/Internal/GroupingPar.hs
@@ -1,355 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE Strict #-}--{- |-Parallel partitioned group-by assignment. Row indices are partitioned by the-high bits of their key-hash (a counting sort: per-range histogram, prefix-sum,-scatter into one index buffer laid out partition-by-partition). One task per-capability then groups its partitions with its OWN open-addressing hash table-('DataFrame.Internal.HashTable') — keys are disjoint across partitions, so the-per-partition group sets concatenate with NO merge.--The output @(rowToGroup, valueIndices, offsets)@ is /bit-for-bit identical/ to-the sequential 'DataFrame.Internal.Grouping.groupBy': groups are emitted in-ascending @(repHash, repRow)@ order (signed-Int order on the hash, tie-broken by-the representative row). Because the partition key is the top bits of a-sign-preserving unsigned remap of the same hash, partition order already agrees-with that global order; within a partition we sort the local groups by the same-key. This is the parallel==sequential correctness gate.--The driver forks plain 'forkIO' workers (no sparks) over a shared atomic-counter-work queue, so partition skew is balanced. A sequential fallback is used when-there is a single capability or the row count is below 'parThreshold' (decided in-'shouldParallelize', which 'groupBy' consults before calling here).--}-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-    -- Phase 1: counting sort of row indices by partition.-    (partStart, sortedRows) <- partitionRows n hashes p shift-    -- Phase 2: per-partition grouping + per-partition canonical ranking (both-    -- inside the parallel worker). localGid[pos] = local group id of the row at-    -- sorted position 'pos'; canonBoxes[part] = its rank vector.-    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-    -- Phase 3: global base ids (serial prefix sum; the ranking is already done).-    (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 that-pull partition indices off a shared counter. For partition @pp@ spanning-@[partStart[pp], partStart[pp+1])@ of @sortedRows@ a worker assigns dense local-group ids (first-appearance order) into @localGid@ at the same sorted positions,-recording each new group's representative hash and row into unboxed-first-appearance vectors. It then computes @canonBoxes[pp]@ — @canon[localGid] =-within-partition canonical rank — via a stable radix sort on the unsigned-representative hash (ties keep first-appearance order, which is ascending repRow,-so the @(key hash, repRow)@ order of the old comparison sort is reproduced with-no boxed tuples). @nLocalGroups[pp]@ holds the group count.--}-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-                -- repHash indexed by local gid (first-appearance order). The-                -- representative row is implicitly ascending in gid order (sorted-                -- positions are in original-row order, a stable counting sort),-                -- so the stable hash-rank below needs no explicit repRow.-                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-                -- Rank this partition's groups into canonical order (shared with-                -- the sequential path). repRow is ascending in gid order (stable-                -- counting sort keeps sorted positions in original-row order), so-                -- the stable hash-rank's tie-break reproduces (hash, repRow).-                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@-(length @p+1@, @globalBase[pp]@ is the first global id of partition @pp@,-@globalBase[p]@ the total). The per-partition canonical ranks were already-computed in 'runPartitions'; partitions are in ascending key order so prepending-@globalBase[pp]@ to each rank yields the sequential @canonicalRemap@ 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)@. For each sorted-position we know its partition, its local group id and the canonical maps, so the-global group id is @globalBase[pp] + canonOf[pp][localGid]@. @valueIndices@ is the-rows ordered by global group; @offsets@ the per-group boundaries; @rowToGroup@ the-inverse mapping 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)-    -- Global group id of each sorted position, plus per-group counts.-    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-    -- offsets = exclusive prefix sum of counts.-    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-    -- valueIndices: place each sorted position's row at its group's cursor.-    -- Iterating sorted positions in order keeps rows in original order within a-    -- group (the partition counting sort and grouping both preserve it).-    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
− src/DataFrame/Internal/Hash.hs
@@ -1,128 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE CPP #-}-{-# LANGUAGE MagicHash #-}--{- |-A poor-man's hash used by 'DataFrame.Internal.Grouping' to bucket rows-without depending on the @hashable@ package.--Each value is folded into an 'Int' accumulator with an FxHash-style step-(rotate, xor, multiply). It is intentionally small and not cryptographically-strong — it only needs to spread group-key tuples well enough that-'Data.IntMap' bucketing produces sensible groups.--}-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 of a nullable column, so that a-@Nothing@ does not hash to the same value as a present @Just x@ that happens to-store the same underlying bits (notably @Just 0@). A fixed distinctive constant-(the 64-bit golden-ratio mix constant) keeps null hashing deterministic; a real-value equal to it collides only as rarely as any other hash collision.--}-nullSalt :: Int-nullSalt = fromIntegral (0x9E3779B97F4A7C15 :: Word64)--{- | Mix an 'Int' into the accumulator.--An FxHash-style step (rotate the accumulator, xor the value, multiply by a large-odd constant). The rotate diffuses each value's bits across all positions before-the next is folded in, so small/adjacent integers — common as group keys — do-not produce the structured collisions that a plain @(acc `xor` x) * prime@ does-once several columns are combined. Grouping trusts hash equality, so this-robustness is what keeps distinct rows in distinct groups.--}-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 (rather than decoding and-mixing one codepoint at a time) cuts the multiply count ~8x on long keys while-staying collision-equivalent: UTF-8 is injective, so equal 'T.Text's mix to the-same value and distinct ones almost never collide. The trailing @len `mod` 8@-bytes are folded in individually.--}-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 #-}
− src/DataFrame/Internal/HashTable.hs
@@ -1,113 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}--{- |-A flat, unboxed, open-addressing (linear-probe) hash table that maps a row's-key-hash to a /dense group id/, verifying the real key on every hash hit.--The table is three parallel unboxed 'VUM.MVector's keyed by hash slot:--  * @htHash@   — the stored hash at each slot.-  * @htGroup@  — the dense group id stored at each slot (@-1@ marks an empty-    slot, since real group ids are @>= 0@).-  * @htRep@    — the representative row index of that group, used to re-verify-    the real key columns on a hash hit and so reject collisions.--It is 'PrimMonad'-polymorphic: it runs in 'Control.Monad.ST.ST' for the current-single-threaded 'DataFrame.Internal.Grouping.groupBy' and can run in 'IO' inside-a per-worker partition once grouping is parallelised. The lookup-or-insert loop-('htInsert') trusts the caller-supplied @eqRow@ predicate to compare the key-columns of two rows by index, fixing the hash-only bucketing that previously-merged colliding keys.--}-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@) in the table, returning its dense-group id. On an empty slot the row starts a new group: the caller's-@nextGroup@ thunk supplies the next dense id, and the row is recorded as that-group's representative. On a stored-hash match the real key is re-verified with-@eqRow rep row@ before the existing id is returned; a mismatch is a hash-collision and probing continues. The returned 'Bool' is 'True' when a new group-was created, letting the caller bump its group counter without a second read.--}-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 #-}
− src/DataFrame/Internal/Interpreter.hs
@@ -1,1103 +0,0 @@-{-# 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 a 'liftValue' — this is how the-old @Binary op (Lit l) right@ special cases are recovered without-explicit pattern matches in the evaluator.--}-liftValue2 ::-    (Columnable c, Columnable b, Columnable a) =>-    (c -> b -> a) ->-    Value c ->-    Value b ->-    Either DataFrameException (Value a)-liftValue2 f (Scalar l) (Scalar r) = Right (Scalar (f l r))-liftValue2 f (Scalar l) v = liftValue (f l) v-liftValue2 f v (Scalar r) = liftValue (`f` r) v-liftValue2 f (Flat l) (Flat r) = Flat <$> zipWithColumns f l r-liftValue2 f (Group ls) (Group rs)-    | V.length ls == V.length rs =-        Group <$> V.zipWithM (zipWithColumns f) ls rs--- Shape mismatches: aggregated vs. non-aggregated.-liftValue2 _ (Flat _) (Group _) =-    Left $ AggregatedAndNonAggregatedException "aggregated" "non-aggregated"-liftValue2 _ (Group _) (Flat _) =-    Left $ AggregatedAndNonAggregatedException "non-aggregated" "aggregated"-liftValue2 _ (Group _) (Group _) =-    Left $ InternalException "Group count mismatch in binary operation"-{-# 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-    -- Prefer the existing value; fall back to the new one.-    Nothing <|+> b = b-    a <|+> _ = a-enrichError _ e = e------------------------------------------------------------------------------------ Group slicing----------------------------------------------------------------------------------{- | Given a flat column and grouping metadata, produce one 'Column' per-group.  Each result column is an O(1) slice into a sorted copy of the-input — the sort happens once, not per-group.--}-sliceGroups :: Column -> VU.Vector Int -> VU.Vector Int -> V.Vector Column-sliceGroups col os indices = case col of-    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----------------------------------------------------------------------------------{- | Apply a result-handler @onResult@ to each element of a column after-coercing it to type @a@.  Covers three modes in one:--* @onResult = either (const Nothing) Just@  → like @cast@   (returns @Maybe a@)-* @onResult = either (const def) id@         → like @castWithDefault@ (returns @a@)-* @onResult = either (Left . T.pack) Right@  → like @castEither@       (returns @Either T.Text a@)--Numeric coercion handles Double, Float, and Int targets.  Text columns-(String / T.Text) are parsed via 'reads'.  Any other mismatch returns-'Left TypeMismatchException'.--}-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 the output type @b@ is @Maybe c@ (or @Maybe (Maybe c)@) and the-column stores plain @c@ values, wrap each element in 'Just'.-The @Maybe (Maybe c)@ case applies join semantics: instead of producing-a double-wrapped column, a @Maybe c@ column is returned, so-@castExpr \@(Maybe Double)@ on a @Double@ column yields @Maybe Double@-rather than @Maybe (Maybe Double)@.-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)--- Leaves -------------------------------------------------------------------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 ()-                        )--- CastWith -----------------------------------------------------------------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))--- CastExprWith -------------------------------------------------------------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--- Unary --------------------------------------------------------------------eval ctx expr@(Unary op (inner :: Expr b)) = addContext expr $ do-    v <- eval @b ctx inner-    liftValue (unaryFn op) v---- Binary -------------------------------------------------------------------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---- If -----------------------------------------------------------------------eval ctx expr@(If cond l r) = addContext expr $ do-    c <- eval @Bool ctx cond-    lv <- eval @a ctx l-    rv <- eval @a ctx r-    branchValue c lv rv---- Over (window function) ---------------------------------------------------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 ->-            -- Scalar agg (mean, sum, median): one value per group.-            -- Broadcast via rowToGroup: row i gets value at group rowToGroup[i].-            Right (Flat (atIndicesStable (rowToGroup gdf) groupCol))-        Group groupCols -> do-            -- Concatenate in sorted order, then unsort to original row order.-            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-                            )--- Aggregation: CollectAgg --------------------------------------------------eval ctx expr@(Agg (CollectAgg _ (f :: v b -> a)) inner) =-    addContext expr $ do-        v <- eval @b ctx inner-        case v of-            Scalar _ ->-                Left $-                    InternalException-                        "Cannot apply a collection aggregation to a scalar"-            Flat col ->-                Scalar <$> applyCollect @v @b @a f col-            Group gs ->-                Flat . fromVector-                    <$> V.mapM (applyCollect @v @b @a f) gs---- Aggregation: FoldAgg with seed -------------------------------------------eval ctx expr@(Agg (FoldAgg _ (Just seed) (f :: a -> b -> a)) inner) =-    addContext expr $ do-        v <- eval @b ctx inner-        case v of-            Scalar 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---- Aggregation: MergeAgg ----------------------------------------------------eval-    ctx-    expr@( Agg-                (MergeAgg _ seed (step :: acc -> b -> acc) _ (finalize :: acc -> a))-                (inner :: Expr b)-            ) =-        addContext expr $ do-            v <- eval @b ctx inner-            case v of-                Scalar 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---- Aggregation: FoldAgg without seed (fold1) --------------------------------eval ctx expr@(Agg (FoldAgg _ Nothing (f :: a -> b -> a)) inner) =-    addContext expr $-        case testEquality (typeRep @a) (typeRep @b) of-            Nothing ->-                Left $-                    InternalException-                        "Type mismatch in seedless fold: \-                        \accumulator and element types must match"-            Just Refl -> do-                v <- eval @b ctx inner-                case v of-                    Scalar _ ->-                        Left $-                            InternalException-                                "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.  This is the original top-level entry point; internally-it calls 'eval' and materialises the result.--NOTE: unlike the old implementation, 'Lit' values are no longer-eagerly broadcast.  The broadcast happens here, at the boundary,-via 'materialize'.--}-interpret ::-    forall a.-    (Columnable a) =>-    DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)-interpret df expr = do-    v <- eval (FlatCtx df) expr-    pure $ TColumn $ materialize @a (fst (dataframeDimensions df)) v--{- | Interpret an expression against a 'GroupedDataFrame',-distinguishing aggregated results from bare column references.-Internally calls 'eval'.--}-interpretAggregation ::-    forall a.-    (Columnable a) =>-    GroupedDataFrame ->-    Expr a ->-    Either DataFrameException (AggregationResult a)-interpretAggregation gdf expr = do-    v <- eval (GroupCtx gdf) expr-    case v of-        Scalar a ->-            Right $-                Aggregated $-                    TColumn $-                        broadcastScalar @a (numGroups gdf) a-        Flat col ->-            Right $ Aggregated $ TColumn col-        Group _ ->-            -- The Column payload is intentionally unused — the only-            -- call-site ('aggregate') immediately throws-            -- 'UnaggregatedException' on this constructor.-            Right $ UnAggregated $ BoxedColumn @T.Text Nothing V.empty
− src/DataFrame/Internal/Nullable.hs
@@ -1,500 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE FunctionalDependencies #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE TypeOperators #-}-{-# LANGUAGE UndecidableInstances #-}-{-# LANGUAGE UndecidableSuperClasses #-}--{- | Nullable-aware binary operations for expressions.--This module provides two type classes, 'NullableArithOp' and 'NullableCmpOp',-which enable operators like '.+', '.-', '.*', './', '.==' etc. to work-transparently across combinations of nullable (@Maybe a@) and non-nullable-(@a@) column types.--The partial functional dependencies uniquely determine the result type from-the operand types, so GHC infers it without annotations.--The four combinations covered for each class:--* @(a, a)@               — non-nullable × non-nullable-* @(Maybe a, a)@         — nullable × non-nullable-* @(a, Maybe a)@         — non-nullable × nullable-* @(Maybe a, Maybe a)@   — both nullable--== Usage--@--- Mixing nullable and non-nullable columns:-F.col \@Int \"x\" '.+' F.col \@(Maybe Int) \"y\"  -- :: Expr (Maybe Int)---- Both non-nullable (existing behaviour preserved):-F.col \@Int \"x\" '.+' F.col \@Int \"y\"           -- :: Expr Int---- Comparison with three-valued logic:-F.col \@(Maybe Int) \"x\" '.==' F.col \@Int \"y\"  -- :: Expr (Maybe Bool)-@--}-module DataFrame.Internal.Nullable (-    -- * Type family-    BaseType,--    -- * Arithmetic class-    NullableArithOp (..),--    -- * Comparison class-    NullableCmpOp (..),--    -- * Generalized nullable lift classes-    NullLift1Op (..),-    NullLift2Op (..),--    -- * Result-type type families (drive inference in nullLift / nullLift2)-    NullLift1Result,-    NullLift2Result,--    -- * Result-type type family for comparison operators-    NullCmpResult,--    -- * Numeric widening-    NumericWidenOp (..),-    widenArithOp,-    widenCmpOp,-    WidenResult,--    -- * Division widening (integral × integral → Double)-    DivWidenOp (..),-    divArithOp,-    WidenResultDiv,-) where--import Data.Int (Int32, Int64)-import DataFrame.Internal.Column (Columnable)-import DataFrame.Internal.Types (Promote, PromoteDiv)--{- | Strip one layer of 'Maybe'.--@-BaseType (Maybe a) = a-BaseType a         = a   -- for any non-Maybe type-@--}-type family BaseType a where-    BaseType (Maybe a) = a-    BaseType a = a--{- | Class for arithmetic binary operations that work transparently over-nullable and non-nullable column types.--The functional dependency @a b -> c@ ensures GHC can infer the result type @c@-from the operand types. The 'OVERLAPPABLE' pragma on the non-nullable instance-ensures the more specific @(Maybe a, Maybe a)@ instance wins when both operands-are nullable.--}-class-    ( Columnable a-    , Columnable b-    , Columnable c-    ) =>-    NullableArithOp a b c-        | a b -> c-    where-    {- | Lift an arithmetic function over the inner (non-Maybe) values.-    'Nothing' short-circuits: any 'Nothing' operand produces 'Nothing'.-    -}-    nullArithOp ::-        (BaseType a -> BaseType a -> BaseType a) ->-        a ->-        b ->-        c--{- | Compute the result type of a nullable comparison.--@-NullCmpResult (Maybe a) b = Maybe Bool-NullCmpResult a (Maybe b) = Maybe Bool   -- when a is apart from Maybe-NullCmpResult a b         = Bool-@--Used by the comparison operators ('.==', '.<', etc.) so GHC infers the-return type without an explicit annotation.--}-type family NullCmpResult a b where-    NullCmpResult (Maybe a) b = Maybe Bool-    NullCmpResult a (Maybe b) = Maybe Bool-    NullCmpResult a b = Bool--{- | Class for comparison binary operations that work transparently over-nullable and non-nullable column types.--No functional dependency on @e@: the 'OVERLAPPING'\/'OVERLAPPABLE' pragmas on-instances disambiguate at call sites without a FundDep (which would conflict-when both operands are @Maybe@). GHC selects the unique most-specific instance-from the concrete operand types.--}-class-    ( Columnable a-    , Columnable b-    , Columnable e-    ) =>-    NullableCmpOp a b e-    where-    {- | Lift a comparison function over the inner values (three-valued logic).-    Returns 'Nothing' when either operand is 'Nothing'.-    -}-    nullCmpOp ::-        (BaseType a -> BaseType a -> Bool) ->-        a ->-        b ->-        e--{- | Non-nullable × Non-nullable: apply directly, no wrapping.-Arithmetic result is @a@; comparison result is @Bool@.--}-instance-    {-# OVERLAPPABLE #-}-    (Columnable a, a ~ BaseType a) =>-    NullableArithOp a a a-    where-    nullArithOp f = f--instance-    {-# OVERLAPPABLE #-}-    (Columnable a, Columnable Bool, a ~ BaseType a) =>-    NullableCmpOp a a Bool-    where-    nullCmpOp f = f---- | Nullable × Non-nullable: 'Nothing' short-circuits.-instance-    (Columnable a, Columnable (Maybe a)) =>-    NullableArithOp (Maybe a) a (Maybe a)-    where-    nullArithOp _f Nothing _ = Nothing-    nullArithOp f (Just x) y = Just (f x y)--instance-    (Columnable a, Columnable (Maybe a), Columnable (Maybe Bool)) =>-    NullableCmpOp (Maybe a) a (Maybe Bool)-    where-    nullCmpOp _f Nothing _ = Nothing-    nullCmpOp f (Just x) y = Just (f x y)---- | Non-nullable × Nullable: 'Nothing' short-circuits.-instance-    ( Columnable a-    , Columnable (Maybe a)-    , a ~ BaseType a-    ) =>-    NullableArithOp a (Maybe a) (Maybe a)-    where-    nullArithOp _f _ Nothing = Nothing-    nullArithOp f x (Just y) = Just (f x y)--instance-    ( Columnable a-    , Columnable (Maybe a)-    , Columnable (Maybe Bool)-    , a ~ BaseType a-    ) =>-    NullableCmpOp a (Maybe a) (Maybe Bool)-    where-    nullCmpOp _f _ Nothing = Nothing-    nullCmpOp f x (Just y) = Just (f x y)---- | Nullable × Nullable: either 'Nothing' short-circuits.-instance-    {-# OVERLAPPING #-}-    (Columnable a, Columnable (Maybe a)) =>-    NullableArithOp (Maybe a) (Maybe a) (Maybe a)-    where-    nullArithOp _f Nothing _ = Nothing-    nullArithOp _f _ Nothing = Nothing-    nullArithOp f (Just x) (Just y) = Just (f x y)--instance-    {-# OVERLAPPING #-}-    (Columnable a, Columnable (Maybe a), Columnable (Maybe Bool)) =>-    NullableCmpOp (Maybe a) (Maybe a) (Maybe Bool)-    where-    nullCmpOp _f Nothing _ = Nothing-    nullCmpOp _f _ Nothing = Nothing-    nullCmpOp f (Just x) (Just y) = Just (f x y)---- ------------------------------------------------------------------------------ Generalized nullable lift (unary)--- -----------------------------------------------------------------------------{- | Lift a unary function over a column expression, propagating 'Nothing'.--When @a@ is non-nullable the function is applied directly; when @a = Maybe x@-the function is applied under the 'Just' and 'Nothing' short-circuits.--Use via 'DataFrame.Functions.nullLift'.--}--{- | Compute the result type of a nullable unary lift.--@-NullLift1Result (Maybe a) r = Maybe r-NullLift1Result a         r = r        -- for any non-Maybe a-@--Used by 'DataFrame.Functions.nullLift' so GHC can infer the return type-without an explicit annotation.--}-type family NullLift1Result a r where-    NullLift1Result (Maybe a) r = Maybe r-    NullLift1Result a r = r--class-    ( Columnable a-    , Columnable r-    , Columnable c-    ) =>-    NullLift1Op a r c-    where-    applyNull1 :: (BaseType a -> r) -> a -> c---- | Non-nullable: apply directly.-instance-    {-# OVERLAPPABLE #-}-    (Columnable a, Columnable r, a ~ BaseType a) =>-    NullLift1Op a r r-    where-    applyNull1 f = f---- | Nullable: propagate 'Nothing'.-instance-    {-# OVERLAPPING #-}-    (Columnable a, Columnable r, Columnable (Maybe r)) =>-    NullLift1Op (Maybe a) r (Maybe r)-    where-    applyNull1 _ Nothing = Nothing-    applyNull1 f (Just x) = Just (f x)---- ------------------------------------------------------------------------------ Generalized nullable lift (binary)--- -----------------------------------------------------------------------------{- | Lift a binary function over two column expressions, propagating 'Nothing'.--The four combinations:--* @(a, b)@               — both non-nullable: result is @r@-* @(Maybe a, b)@         — left nullable: result is @Maybe r@-* @(a, Maybe b)@         — right nullable: result is @Maybe r@-* @(Maybe a, Maybe b)@   — both nullable: result is @Maybe r@--Use via 'DataFrame.Functions.nullLift2'.--}--{- | Compute the result type of a nullable binary lift.--@-NullLift2Result (Maybe a) b         r = Maybe r-NullLift2Result a         (Maybe b) r = Maybe r   -- when a is apart from Maybe-NullLift2Result a         b         r = r-@--Used by 'DataFrame.Functions.nullLift2' so GHC can infer the return type.--}-type family NullLift2Result a b r where-    NullLift2Result (Maybe a) b r = Maybe r-    NullLift2Result a (Maybe b) r = Maybe r-    NullLift2Result a b r = r--class-    ( Columnable a-    , Columnable b-    , Columnable r-    , Columnable c-    ) =>-    NullLift2Op a b r c-    where-    applyNull2 :: (BaseType a -> BaseType b -> r) -> a -> b -> c---- | Both non-nullable: apply directly.-instance-    {-# OVERLAPPABLE #-}-    (Columnable a, Columnable b, Columnable r, a ~ BaseType a, b ~ BaseType b) =>-    NullLift2Op a b r r-    where-    applyNull2 f = f---- | Left nullable: 'Nothing' short-circuits.-instance-    {-# OVERLAPPABLE #-}-    (Columnable a, Columnable b, Columnable r, Columnable (Maybe r), b ~ BaseType b) =>-    NullLift2Op (Maybe a) b r (Maybe r)-    where-    applyNull2 _ Nothing _ = Nothing-    applyNull2 f (Just x) y = Just (f x y)---- | Right nullable: 'Nothing' short-circuits.-instance-    {-# OVERLAPPABLE #-}-    (Columnable a, Columnable b, Columnable r, Columnable (Maybe r), a ~ BaseType a) =>-    NullLift2Op a (Maybe b) r (Maybe r)-    where-    applyNull2 _ _ Nothing = Nothing-    applyNull2 f x (Just y) = Just (f x y)---- | Both nullable: either 'Nothing' short-circuits.-instance-    {-# OVERLAPPING #-}-    (Columnable a, Columnable b, Columnable r, Columnable (Maybe r)) =>-    NullLift2Op (Maybe a) (Maybe b) r (Maybe r)-    where-    applyNull2 _ Nothing _ = Nothing-    applyNull2 _ _ Nothing = Nothing-    applyNull2 f (Just x) (Just y) = Just (f x y)---- ------------------------------------------------------------------------------ Numeric widening--- -----------------------------------------------------------------------------{- | Widen two numeric base types to their promoted common type.--When @a ~ b@ the coercions are identity; otherwise one operand is widened-(e.g. 'Int' → 'Double').--}-class (Columnable (Promote a b)) => NumericWidenOp a b where-    widen1 :: a -> Promote a b-    widen2 :: b -> Promote a b---- | Same type: identity coercions.-instance {-# OVERLAPPING #-} (Columnable a) => NumericWidenOp a a where-    widen1 = id-    widen2 = id--instance NumericWidenOp Int Double where widen1 = fromIntegral; widen2 = id-instance NumericWidenOp Double Int where-    widen1 = id-    widen2 = fromIntegral-instance NumericWidenOp Float Double where widen1 = realToFrac; widen2 = id-instance NumericWidenOp Double Float where-    widen1 = id-    widen2 = realToFrac-instance NumericWidenOp Int32 Float where widen1 = fromIntegral; widen2 = id-instance NumericWidenOp Float Int32 where-    widen1 = id-    widen2 = fromIntegral-instance NumericWidenOp Int32 Double where widen1 = fromIntegral; widen2 = id-instance NumericWidenOp Double Int32 where-    widen1 = id-    widen2 = fromIntegral-instance NumericWidenOp Int64 Float where widen1 = fromIntegral; widen2 = id-instance NumericWidenOp Float Int64 where-    widen1 = id-    widen2 = fromIntegral-instance NumericWidenOp Int64 Double where widen1 = fromIntegral; widen2 = id-instance NumericWidenOp Double Int64 where-    widen1 = id-    widen2 = fromIntegral---- | Apply an arithmetic function after widening both operands to their common type.-widenArithOp ::-    forall a b.-    (NumericWidenOp a b) =>-    (Promote a b -> Promote a b -> Promote a b) ->-    a ->-    b ->-    Promote a b-widenArithOp f x y = f (widen1 @a @b x) (widen2 @a @b y)---- | Apply a comparison function after widening both operands to their common type.-widenCmpOp ::-    forall a b.-    (NumericWidenOp a b) =>-    (Promote a b -> Promote a b -> Bool) ->-    a ->-    b ->-    Bool-widenCmpOp f x y = f (widen1 @a @b x) (widen2 @a @b y)---- | Result type of a widening binary operator, accounting for nullable wrappers.-type WidenResult a b = NullLift2Result a b (Promote (BaseType a) (BaseType b))---- ------------------------------------------------------------------------------ Division widening (integral × integral → Double)--- -----------------------------------------------------------------------------{- | Like 'NumericWidenOp' but uses 'PromoteDiv': integral×integral → Double.-Floating types still dominate (Double > Float), and any two integral types-(same or mixed) are both widened to Double.--}-class (Columnable (PromoteDiv a b)) => DivWidenOp a b where-    divWiden1 :: a -> PromoteDiv a b-    divWiden2 :: b -> PromoteDiv a b---- Floating same-type (identity)-instance DivWidenOp Double Double where divWiden1 = id; divWiden2 = id-instance DivWidenOp Float Float where divWiden1 = id; divWiden2 = id---- Mixed Double/Float-instance DivWidenOp Double Float where divWiden1 = id; divWiden2 = realToFrac-instance DivWidenOp Float Double where divWiden1 = realToFrac; divWiden2 = id---- Double beats integral-instance DivWidenOp Double Int where divWiden1 = id; divWiden2 = fromIntegral-instance DivWidenOp Int Double where divWiden1 = fromIntegral; divWiden2 = id-instance DivWidenOp Double Int32 where divWiden1 = id; divWiden2 = fromIntegral-instance DivWidenOp Int32 Double where divWiden1 = fromIntegral; divWiden2 = id-instance DivWidenOp Double Int64 where divWiden1 = id; divWiden2 = fromIntegral-instance DivWidenOp Int64 Double where divWiden1 = fromIntegral; divWiden2 = id---- Float beats integral-instance DivWidenOp Float Int where divWiden1 = id; divWiden2 = fromIntegral-instance DivWidenOp Int Float where divWiden1 = fromIntegral; divWiden2 = id-instance DivWidenOp Float Int32 where divWiden1 = id; divWiden2 = fromIntegral-instance DivWidenOp Int32 Float where divWiden1 = fromIntegral; divWiden2 = id-instance DivWidenOp Float Int64 where divWiden1 = id; divWiden2 = fromIntegral-instance DivWidenOp Int64 Float where divWiden1 = fromIntegral; divWiden2 = id---- Integral × integral → Double-instance DivWidenOp Int Int where-    divWiden1 = fromIntegral-    divWiden2 = fromIntegral-instance DivWidenOp Int32 Int32 where-    divWiden1 = fromIntegral-    divWiden2 = fromIntegral-instance DivWidenOp Int64 Int64 where-    divWiden1 = fromIntegral-    divWiden2 = fromIntegral-instance DivWidenOp Int Int32 where-    divWiden1 = fromIntegral-    divWiden2 = fromIntegral-instance DivWidenOp Int32 Int where-    divWiden1 = fromIntegral-    divWiden2 = fromIntegral-instance DivWidenOp Int Int64 where-    divWiden1 = fromIntegral-    divWiden2 = fromIntegral-instance DivWidenOp Int64 Int where-    divWiden1 = fromIntegral-    divWiden2 = fromIntegral-instance DivWidenOp Int32 Int64 where-    divWiden1 = fromIntegral-    divWiden2 = fromIntegral-instance DivWidenOp Int64 Int32 where-    divWiden1 = fromIntegral-    divWiden2 = fromIntegral---- | Apply an arithmetic function after widening both operands via 'PromoteDiv'.-divArithOp ::-    forall a b.-    (DivWidenOp a b) =>-    (PromoteDiv a b -> PromoteDiv a b -> PromoteDiv a b) ->-    a ->-    b ->-    PromoteDiv a b-divArithOp f x y = f (divWiden1 @a @b x) (divWiden2 @a @b y)---- | Result type of a division-widening binary operator, accounting for nullable wrappers.-type WidenResultDiv a b =-    NullLift2Result a b (PromoteDiv (BaseType a) (BaseType b))
− src/DataFrame/Internal/PackedText.hs
@@ -1,169 +0,0 @@-{-# LANGUAGE BangPatterns #-}--{- | Packed-text payload + byte-slice primitives. A 'PackedTextData' shares a-single UTF-8 byte buffer across all rows of a string column, with @n+1@ row-offsets, so no per-row 'Data.Text.Text' header is materialized at freeze.-'Data.Text.Text' is produced only on demand (display, typed extraction) via-the same decode path that the boxed-Text builder used.--A gathered/joined/sorted result keeps sharing that buffer: instead of copying-bytes it carries a @ptSel@ selection vector that reindexes the base rows, so a-permuted or row-exploded column stays a 'PackedText' (shared buffer + permuted-indices) rather than materializing back to boxed 'Data.Text.Text'.--}-module DataFrame.Internal.PackedText (-    PackedTextData (..),-    mkPackedContiguous,-    packedGather,-    packedTake,-    packedRowOffsetVec,-    packedLength,-    packedSlice,-    packedIndexText,-    sliceEqBytes,-    sliceCmpBytes,-) where--import qualified Data.Text as T-import qualified Data.Text.Array as A-import qualified Data.Vector.Unboxed as VU--import Data.Ord (comparing)-import Data.Text.Internal (Text (Text))-import DataFrame.Internal.Utf8 (isValidUtf8Slice, lenientDecodeSlice)--{- | A shared UTF-8 byte buffer plus @n+1@ row offsets (base row @r@ spans bytes-@[offsets!r, offsets!(r+1))@). Validity lives in the enclosing column's-@Maybe Bitmap@, mirroring 'BoxedColumn'/'UnboxedColumn'.--@ptSel@ is an optional selection layer: when @Nothing@ the column is the-contiguous base (row @i@ == base row @i@). When @Just sel@, logical row @i@ is-base row @sel!i@; this is how a gather/join/sort result shares the buffer-without copying bytes. Out-of-range entries in @sel@ (e.g. a join @-1@-sentinel) decode to the empty slice and are masked by the column bitmap.--}-data PackedTextData = PackedTextData-    { ptBytes :: {-# UNPACK #-} !A.Array-    , ptOffsets :: {-# UNPACK #-} !(VU.Vector Int)-    , ptSel :: !(Maybe (VU.Vector Int))-    }---- | Build a contiguous packed payload (no selection): the freeze-path shape.-mkPackedContiguous :: A.Array -> VU.Vector Int -> PackedTextData-mkPackedContiguous arr offs = PackedTextData arr offs Nothing-{-# INLINE mkPackedContiguous #-}--{- | Reindex a packed payload by a selection vector, sharing the byte buffer-and base offsets. Logical row @i@ becomes base row @indices!i@. A negative or-out-of-range index decodes to the empty slice (callers mask it with a bitmap).-Composes with an existing selection so a gather of a gather still shares the-buffer.--}-packedGather :: VU.Vector Int -> PackedTextData -> PackedTextData-packedGather indices (PackedTextData arr offs msel) =-    let !base = VU.length offs - 1-        clamp r = if r >= 0 && r < base then r else -1-        sel' = case msel of-            Nothing -> VU.map clamp indices-            Just s ->-                VU.map-                    (\i -> if i >= 0 && i < VU.length s then clamp (VU.unsafeIndex s i) else -1)-                    indices-     in PackedTextData arr offs (Just sel')-{-# INLINE packedGather #-}--{- | Take the first @k@ logical rows, sharing the byte buffer. With a selection-layer the selection is sliced to @k@ entries; without one a base-row selection-@[0 .. k-1]@ is installed (slicing the contiguous offsets would still leave the-trailing bytes addressable, but a short selection caps 'packedLength' to @k@).-O(k), no byte copy or decode — the fix for cheap @take@/display on a 1e7-row-packed column.--}-packedTake :: Int -> PackedTextData -> PackedTextData-packedTake k (PackedTextData arr offs msel) =-    let !base = VU.length offs - 1-        !k' = max 0 k-     in case msel of-            Just s -> PackedTextData arr offs (Just (VU.take k' s))-            Nothing -> PackedTextData arr offs (Just (VU.enumFromN 0 (min k' base)))-{-# INLINE packedTake #-}---- | Map a logical row index to its base row, honoring any selection layer.-baseRow :: PackedTextData -> Int -> Int-baseRow (PackedTextData _ _ Nothing) i = i-baseRow (PackedTextData _ _ (Just sel)) i = VU.unsafeIndex sel i-{-# INLINE baseRow #-}---- | Row count: @length sel@ when selected, else @length offsets - 1@.-packedLength :: PackedTextData -> Int-packedLength (PackedTextData _ offs Nothing) = VU.length offs - 1-packedLength (PackedTextData _ _ (Just sel)) = VU.length sel-{-# INLINE packedLength #-}---- | Raw byte slice for logical row @i@: @(buffer, offset, length)@. The hot accessor.-packedSlice :: PackedTextData -> Int -> (A.Array, Int, Int)-packedSlice p@(PackedTextData arr offs _) i =-    let !r = baseRow p i-     in if r < 0-            then (arr, 0, 0)-            else-                let o = VU.unsafeIndex offs r in (arr, o, VU.unsafeIndex offs (r + 1) - o)-{-# INLINE packedSlice #-}--{- | The shared buffer + contiguous @n+1@ offsets when the payload is the-unselected base. A selected (gathered) payload has non-contiguous rows that a-single offset vector cannot express, so this returns @Nothing@ and the caller-decodes per-row via 'packedIndexText'. Lets the boxed-Text fallback take the-fast contiguous 'sliceTextVector' path when possible.--}-packedRowOffsetVec :: PackedTextData -> Maybe (A.Array, VU.Vector Int)-packedRowOffsetVec (PackedTextData arr offs Nothing) = Just (arr, offs)-packedRowOffsetVec _ = Nothing-{-# INLINE packedRowOffsetVec #-}--{- | On-demand single 'Data.Text.Text' for row @i@, using the same-validate-or-lenient decode as 'sliceTextVector' so output is bit-identical.--}-packedIndexText :: PackedTextData -> Int -> T.Text-packedIndexText p i =-    let (arr, o, l) = packedSlice p i-     in decodeField arr o l-{-# INLINE packedIndexText #-}---- Decode one field exactly as 'sliceTextVector' does per row.-decodeField :: A.Array -> Int -> Int -> T.Text-decodeField arr o l-    | l == 0 = T.empty-    | isValidUtf8Slice arr o l = Text arr o l-    | otherwise = lenientDecodeSlice arr o l-{-# INLINE decodeField #-}--{- | Byte-wise equality of two slices. UTF-8 is injective on valid scalar-sequences and lenient decode is deterministic, so this agrees with-@Text@'s '==' on the decoded values.--}-sliceEqBytes :: A.Array -> Int -> Int -> A.Array -> Int -> Int -> Bool-sliceEqBytes a ao al b bo bl-    | al /= bl = False-    | otherwise = go 0-  where-    go !k-        | k >= al = True-        | A.unsafeIndex a (ao + k) == A.unsafeIndex b (bo + k) = go (k + 1)-        | otherwise = False-{-# INLINE sliceEqBytes #-}--{- | Unsigned byte-lexicographic comparison (memcmp semantics). For-well-formed UTF-8 this matches 'Data.Text.compare' exactly, since UTF-8-byte order equals codepoint order for all valid scalars.--}-sliceCmpBytes :: A.Array -> Int -> Int -> A.Array -> Int -> Int -> Ordering-sliceCmpBytes a ao al b bo bl = go 0-  where-    !m = min al bl-    go !k-        | k >= m = compare al bl-        | otherwise = case comparing id (A.unsafeIndex a (ao + k)) (A.unsafeIndex b (bo + k)) of-            EQ -> go (k + 1)-            r -> r-{-# INLINE sliceCmpBytes #-}
− src/DataFrame/Internal/ParRadixSort.hs
@@ -1,294 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ScopedTypeVariables #-}--{- |-Parallel stable sort of row indices by the ascending unsigned order of a-per-row 'Int' hash. Shared by the join build-side 'CompactIndex' construction-(@DataFrame.Operations.Join@), which previously paid a single-threaded-comparison sort over the whole build side — the dominant serial cost of a large-inner join (the @1e7 x 1e7@ big-inner case).--@parSortByHash n hashes@ returns @(sortedHashes, sortedIndices)@ where-@sortedIndices@ lists @[0, n)@ in ascending 'sortKey' order of their hash, ties-broken by ascending original index (stable), and @sortedHashes[k] ==-hashes[sortedIndices[k]]@. Bit-for-bit identical to the old stable merge sort's-output ordering, so equal-hash rows stay contiguous (the run scan in-'buildCompactIndex' depends on this) and within a run keep original-row order.--Strategy (mirrors "DataFrame.Internal.GroupingPar"): a counting sort buckets-rows by the top @log2 p@ bits of their unsigned key into @p@ partitions laid out-in ascending key order; @caps@ 'forkIO' workers then LSD-radix-sort each-partition by the full 56 remaining low bits. Because partitions are already in-global key order and each per-partition sort is stable, concatenating them-reproduces the global stable order with no merge step. A sequential LSD radix-sort is used below 'parSortThreshold' or on a single capability.--}-module DataFrame.Internal.ParRadixSort (-    parSortByHash,-    parSortThreshold,-) 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.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM-import Data.Word (Word64)-import DataFrame.Internal.RadixRank (sortKey)-import System.IO.Unsafe (unsafePerformIO)--{- | Below this many rows the partition/fork overhead is not worth it; the-caller's sequential LSD radix path is used instead.--}-parSortThreshold :: Int-parSortThreshold = 500000--capabilities :: Int-capabilities = unsafePerformIO getNumCapabilities-{-# NOINLINE capabilities #-}--{- | Top-bits partition index of a hash: the high @64 - shift@ bits of its-unsigned 'sortKey'. Ascending partition order equals ascending key order.--}-partIx :: Int -> Int -> Int-partIx shift h = fromIntegral ((fromIntegral (sortKey h) :: Word64) `unsafeShiftR` shift)-{-# INLINE partIx #-}---- | Number of partitions: a power of two, at least @4 * caps@, 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 stable sort of @[0, n)@ by ascending unsigned hash order. See the-module header for the ordering contract.--}-parSortByHash :: Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)-parSortByHash n hashes-    | n <= 1 =-        (hashes, VU.enumFromN 0 n)-    | n < parSortThreshold || capabilities <= 1 =-        seqSortByHash n hashes-    | otherwise = unsafePerformIO (parSortByHashIO n hashes)-{-# NOINLINE parSortByHash #-}------------------------------------------------------------------------------------ Sequential LSD radix sort (also the per-partition worker kernel)----------------------------------------------------------------------------------{- | Stable LSD radix sort of @[0, n)@ by ascending 'sortKey' of their hash, 8-bits per pass over the full 64-bit key. Returns @(sortedHashes, sortedIndices)@.--}-seqSortByHash :: Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)-seqSortByHash n hashes = unsafePerformIO $ do-    keysA <- VUM.new n-    orderA <- VUM.new n-    let seed !i-            | i >= n = pure ()-            | otherwise = do-                VUM.unsafeWrite keysA i (sortKey (VU.unsafeIndex hashes i))-                VUM.unsafeWrite orderA i i-                seed (i + 1)-    seed 0-    keysB <- VUM.new n-    orderB <- VUM.new n-    radixPasses n keysA orderA keysB orderB-    order <- VU.unsafeFreeze orderA-    pure (VU.unsafeBackpermute hashes order, order)--{- | Run all eight stable 8-bit LSD passes, ping-ponging between the two-key/order buffer pairs so the sorted order lands back in @(keysA, orderA)@.-@keysA[i]@ must already hold @sortKey (hash of orderA[i])@ on entry.--}-radixPasses ::-    Int ->-    VUM.IOVector Int ->-    VUM.IOVector Int ->-    VUM.IOVector Int ->-    VUM.IOVector Int ->-    IO ()-radixPasses n keysA orderA keysB orderB = do-    counts <- VUM.new 256-    let pass ::-            Int ->-            VUM.IOVector Int ->-            VUM.IOVector Int ->-            VUM.IOVector Int ->-            VUM.IOVector Int ->-            IO ()-        pass !shiftBits !srcK !srcO !dstK !dstO = do-            VUM.set counts 0-            let count !i-                    | i >= n = pure ()-                    | otherwise = do-                        k <- VUM.unsafeRead srcK i-                        let !b = (k `unsafeShiftR` shiftBits) .&. 0xff-                        VUM.unsafeRead counts b >>= VUM.unsafeWrite counts b . (+ 1)-                        count (i + 1)-            count 0-            let scan !b !acc-                    | b >= 256 = pure ()-                    | otherwise = do-                        c <- VUM.unsafeRead counts b-                        VUM.unsafeWrite counts b acc-                        scan (b + 1) (acc + c)-            scan 0 0-            let place !i-                    | i >= n = pure ()-                    | otherwise = do-                        k <- VUM.unsafeRead srcK i-                        o <- VUM.unsafeRead srcO i-                        let !b = (k `unsafeShiftR` shiftBits) .&. 0xff-                        pos <- VUM.unsafeRead counts b-                        VUM.unsafeWrite counts b (pos + 1)-                        VUM.unsafeWrite dstK pos k-                        VUM.unsafeWrite dstO pos o-                        place (i + 1)-            place 0-    pass 0 keysA orderA keysB orderB-    pass 8 keysB orderB keysA orderA-    pass 16 keysA orderA keysB orderB-    pass 24 keysB orderB keysA orderA-    pass 32 keysA orderA keysB orderB-    pass 40 keysB orderB keysA orderA-    pass 48 keysA orderA keysB orderB-    pass 56 keysB orderB keysA orderA------------------------------------------------------------------------------------ Parallel path: counting-sort partition, then per-partition sort in parallel----------------------------------------------------------------------------------parSortByHashIO :: Int -> VU.Vector Int -> IO (VU.Vector Int, VU.Vector Int)-parSortByHashIO n hashes = do-    caps <- getNumCapabilities-    let !p = numPartitionsFor caps-        !shift = 64 - intLog2 p-    -- Phase 1: counting sort of row indices into ascending-key partitions.-    (partStart, partRows) <- partitionRows n hashes p shift-    -- Phase 2: stable-sort each partition by full key, in parallel. Each worker-    -- owns disjoint [partStart[pp], partStart[pp+1]) output ranges, so the-    -- single shared output buffers are written race-free.-    outOrder <- VUM.new n-    outKeys <- VUM.new n-    sortPartitions caps p partStart partRows hashes outOrder outKeys-    order <- VU.unsafeFreeze outOrder-    pure (VU.unsafeBackpermute hashes order, order)--{- | Bucket every row index into its top-bits 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 ascending key order.--}-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-    rowsM <- 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 rowsM pos i-                VUM.unsafeWrite cursor pp (pos + 1)-                place (i + 1)-    place 0-    partStart <- VU.unsafeFreeze partStartM-    partRows <- VU.unsafeFreeze rowsM-    pure (partStart, partRows)--{- | Stable-sort each partition by full key, writing sorted original indices-into @outOrder@ and their hashes into @outKeys@ at the partition's slot range.-Forks @caps@ workers that pull partition indices off a shared atomic counter.-Within a partition the counting sort already left rows in ascending original-order, so the LSD radix sort's stability reproduces the global @(key, row)@-order. Partitions below two elements are already sorted (counting sort kept-original order) and are copied directly.--}-sortPartitions ::-    Int ->-    Int ->-    VU.Vector Int ->-    VU.Vector Int ->-    VU.Vector Int ->-    VUM.IOVector Int ->-    VUM.IOVector Int ->-    IO ()-sortPartitions caps p partStart partRows hashes outOrder outKeys = do-    next <- newIORef 0-    let sortOne !pp = do-            let !s = VU.unsafeIndex partStart pp-                !e = VU.unsafeIndex partStart (pp + 1)-                !sz = e - s-            when (sz > 0) $-                if sz == 1-                    then do-                        let !r = VU.unsafeIndex partRows s-                        VUM.unsafeWrite outOrder s r-                        VUM.unsafeWrite outKeys s (VU.unsafeIndex hashes r)-                    else do-                        keysA <- VUM.new sz-                        orderA <- VUM.new sz-                        let seed !i-                                | i >= sz = pure ()-                                | otherwise = do-                                    let !r = VU.unsafeIndex partRows (s + i)-                                    VUM.unsafeWrite keysA i (sortKey (VU.unsafeIndex hashes r))-                                    VUM.unsafeWrite orderA i r-                                    seed (i + 1)-                        seed 0-                        keysB <- VUM.new sz-                        orderB <- VUM.new sz-                        radixPasses sz keysA orderA keysB orderB-                        let emit !i-                                | i >= sz = pure ()-                                | otherwise = do-                                    o <- VUM.unsafeRead orderA i-                                    VUM.unsafeWrite outOrder (s + i) o-                                    VUM.unsafeWrite outKeys (s + i) (VU.unsafeIndex hashes o)-                                    emit (i + 1)-                        emit 0-        worker = do-            i <- atomicModifyIORef' next (\j -> (j + 1, j))-            when (i < p) $ sortOne i >> worker-    forkJoin_ (replicate caps worker)---- | 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
− src/DataFrame/Internal/RadixRank.hs
@@ -1,114 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ScopedTypeVariables #-}--{- |-Stable rank of a set of group representatives by the ascending unsigned order of-their hash. Shared by the sequential ('DataFrame.Internal.Grouping') and parallel-('DataFrame.Internal.GroupingPar') group-by canonical-ordering steps so they-stay bit-for-bit identical.--@rankByHash readHash ng@ returns @rank@ with @rank[gid] = position@ of group-@gid@ when groups are ordered by ascending unsigned 'sortKey' of @readHash gid@.-A stable LSD radix sort (8 bits per pass, 8 passes) keeps groups with equal-hash in their original @gid@ order; callers number @gid@s so that this matches-the @repRow@ tie-break of the old comparison sort. @O(ng)@, no boxed tuples or-comparison closures — the lever for the @1e7@-distinct-group case (Q10).--}-module DataFrame.Internal.RadixRank (-    rankByHash,-    sortKey,-) where--import Control.Monad (when)-import Control.Monad.Primitive (PrimMonad)-import Data.Bits (unsafeShiftR, (.&.))-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM-import Data.Word (Word64)--{- | Unsigned sort key of a hash: ascending 'Word64' order of @sortKey h@ equals-ascending signed-'Int' order of @h@. Reinterpreted back to 'Int' for the-byte-wise radix passes (the @.&. 0xff@ byte mask makes the arithmetic shift's-sign extension irrelevant).--}-sortKey :: Int -> Int-sortKey h = fromIntegral (fromIntegral h + 0x8000000000000000 :: Word64)-{-# INLINE sortKey #-}---- | See the module header. @readHash@ supplies the hash of local group @gid@.-rankByHash ::-    forall m. (PrimMonad m) => (Int -> m Int) -> Int -> m (VU.Vector Int)-rankByHash readHash ng = do-    rankM <- VUM.new (max 1 ng)-    if ng <= 1-        then when (ng == 1) (VUM.unsafeWrite rankM 0 0)-        else do-            keysA <- VUM.new ng-            orderA <- VUM.new ng-            let seed !i-                    | i >= ng = pure ()-                    | otherwise = do-                        h <- readHash i-                        VUM.unsafeWrite keysA i (sortKey h)-                        VUM.unsafeWrite orderA i i-                        seed (i + 1)-            seed 0-            keysB <- VUM.new ng-            orderB <- VUM.new ng-            counts <- VUM.new 256-            let pass ::-                    Int ->-                    VUM.MVector (VUM.PrimState m) Int ->-                    VUM.MVector (VUM.PrimState m) Int ->-                    VUM.MVector (VUM.PrimState m) Int ->-                    VUM.MVector (VUM.PrimState m) Int ->-                    m ()-                pass !shiftBits !srcK !srcO !dstK !dstO = do-                    VUM.set counts 0-                    let count !i-                            | i >= ng = pure ()-                            | otherwise = do-                                k <- VUM.unsafeRead srcK i-                                let !b = (k `unsafeShiftR` shiftBits) .&. 0xff-                                VUM.unsafeRead counts b >>= VUM.unsafeWrite counts b . (+ 1)-                                count (i + 1)-                    count 0-                    let scan !b !acc-                            | b >= 256 = pure ()-                            | otherwise = do-                                c <- VUM.unsafeRead counts b-                                VUM.unsafeWrite counts b acc-                                scan (b + 1) (acc + c)-                    scan 0 0-                    let place !i-                            | i >= ng = pure ()-                            | otherwise = do-                                k <- VUM.unsafeRead srcK i-                                o <- VUM.unsafeRead srcO i-                                let !b = (k `unsafeShiftR` shiftBits) .&. 0xff-                                pos <- VUM.unsafeRead counts b-                                VUM.unsafeWrite counts b (pos + 1)-                                VUM.unsafeWrite dstK pos k-                                VUM.unsafeWrite dstO pos o-                                place (i + 1)-                    place 0-            -- 8 stable passes over the 64-bit key; ping-pong so the final-            -- sorted order lands back in (keysA, orderA).-            pass 0 keysA orderA keysB orderB-            pass 8 keysB orderB keysA orderA-            pass 16 keysA orderA keysB orderB-            pass 24 keysB orderB keysA orderA-            pass 32 keysA orderA keysB orderB-            pass 40 keysB orderB keysA orderA-            pass 48 keysA orderA keysB orderB-            pass 56 keysB orderB keysA orderA-            -- orderA[rank] = gid; invert to rank[gid] = rank.-            let inv !r-                    | r >= ng = pure ()-                    | otherwise = do-                        g <- VUM.unsafeRead orderA r-                        VUM.unsafeWrite rankM g r-                        inv (r + 1)-            inv 0-    VU.unsafeFreeze rankM-{-# INLINEABLE rankByHash #-}
− src/DataFrame/Internal/Row.hs
@@ -1,209 +0,0 @@-{-# LANGUAGE ExistentialQuantification #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE InstanceSigs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Internal.Row 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 Control.Exception (throw)-import Data.Function (on)-import Data.Maybe (catMaybes, fromMaybe, isNothing, mapMaybe)-import Data.Type.Equality (TestEquality (..))-import Data.Typeable (type (:~:) (..))-import DataFrame.Errors (DataFrameException (..))-import DataFrame.Internal.Column-import DataFrame.Internal.DataFrame-import DataFrame.Internal.Expression (Expr (..))-import DataFrame.Internal.PackedText (packedIndexText, packedLength)-import Type.Reflection (typeOf, typeRep)--data Any where-    Value :: (Columnable a) => a -> Any-    -- Saves us the extra indirection we get from making Value (Maybe a)-    -- and having to unpack it again to check for nulls.-    -- Instead, we just have Null as a separate constructor.-    Null :: Any--instance Eq Any where-    (==) :: Any -> Any -> Bool-    (Value a) == (Value b) = fromMaybe False $ do-        Refl <- testEquality (typeOf a) (typeOf b)-        return $ a == b-    Null == Null = True-    _ == _ = False--instance Show Any where-    show :: Any -> String-    show (Value a) = T.unpack (showValue a)-    show Null = "null"--showValue :: forall a. (Columnable a) => a -> T.Text-showValue v = case testEquality (typeRep @a) (typeRep @T.Text) of-    Just Refl -> v-    Nothing -> case testEquality (typeRep @a) (typeRep @String) of-        Just Refl -> T.pack v-        Nothing -> (T.pack . show) v---- | Wraps a value into an \Any\ type. This helps up represent rows as heterogenous lists.-toAny :: forall a. (Columnable a) => a -> Any-toAny = Value---- | Unwraps a value from an \Any\ type. A 'Null' cell yields 'Nothing'.-fromAny :: forall a. (Columnable a) => Any -> Maybe a-fromAny Null = Nothing-fromAny (Value (v :: b)) = do-    Refl <- testEquality (typeRep @a) (typeRep @b)-    pure v--{- | Wrap a column cell into an 'Any', honouring the column's null bitmap: a slot-marked invalid becomes 'Null', any other slot becomes a 'Value'. Only needs-@Columnable a@ (the stored element type), not @Columnable (Maybe a)@.--}-cellAny :: (Columnable a) => Maybe Bitmap -> Int -> a -> Any-cellAny Nothing _ x = Value x-cellAny (Just bm) i x = if bitmapTestBit bm i then Value x else Null--type Row = V.Vector Any--(!?) :: [a] -> Int -> Maybe a-(!?) [] _ = Nothing-(!?) (x : _) 0 = Just x-(!?) (_x : xs) n = (!?) xs (n - 1)--{- | Reconstruct column @i@ from a list of rows. The element type is taken from-the first non-'Null' cell; cells of a different type are skipped. If any cell is-'Null' the result is a nullable column (built with 'fromMaybeVec', which needs-only @Columnable a@), so a round-trip through 'toRowList'/'fromRows' preserves-nulls.--}-mkColumnFromRow :: Int -> [[Any]] -> Column-mkColumnFromRow i rows =-    let cells = mapMaybe (!? i) rows-     in case L.find isValue cells of-            Nothing -> fromList ([] :: [T.Text])-            Just (Value (_ :: a)) ->-                let collect Null = Just (Nothing :: Maybe a)-                    collect (Value (v' :: b)) =-                        case testEquality (typeRep @a) (typeRep @b) of-                            Just Refl -> Just (Just v')-                            Nothing -> Nothing-                    maybes = mapMaybe collect cells-                 in if any isNothing maybes-                        then fromMaybeVec (V.fromList maybes)-                        else fromList (catMaybes maybes)-            Just Null -> fromList ([] :: [T.Text]) -- unreachable: find isValue-  where-    isValue (Value _) = True-    isValue Null = False--{- | Converts the entire dataframe to a list of rows.--Each row contains all columns in the dataframe, ordered by their column indices.-The rows are returned in their natural order (from index 0 to n-1).--==== __Examples__-->>> toRowList df-[[("name", "Alice"), ("age", 25), ...], [("name", "Bob"), ("age", 30), ...], ...]--==== __Performance note__--This function materializes all rows into a list, which may be memory-intensive-for large dataframes. Consider using 'toRowVector' if you need random access-or streaming operations.--}-toRowList :: DataFrame -> [[(T.Text, Any)]]-toRowList df =-    let-        names = map fst (L.sortBy (compare `on` snd) $ M.toList (columnIndices df))-     in-        map-            (zip names . V.toList . mkRowRep df names)-            [0 .. (fst (dataframeDimensions df) - 1)]--{- | Converts the dataframe to a vector of rows with only the specified columns.--Each row will contain only the columns named in the @names@ parameter.-This is useful when you only need a subset of columns or want to control-the column order in the resulting rows.--==== __Parameters__--[@names@] List of column names to include in each row. The order of names-          determines the order of fields in the resulting rows.--[@df@] The dataframe to convert.--==== __Examples__-->>> toRowVector ["name", "age"] df-Vector of rows with only name and age fields-->>> toRowVector [] df  -- Empty column list-Vector of empty rows (one per dataframe row)--}-toRowVector :: [T.Text] -> DataFrame -> V.Vector Row-toRowVector names df = V.generate (fst (dataframeDimensions df)) (mkRowRep df names)--{- | Given a row gets the value associated with a field.--==== __Examples__-->>> map (rowValue (F.col @Int "age")) (toRowList df)-[25,30, ...]--}-rowValue :: forall a. Expr a -> [(T.Text, Any)] -> Maybe a-rowValue (Col name) row = lookup name row >>= fromAny @a-rowValue _ _ = error "Can only get rowValue of column reference"--mkRowFromArgs :: [T.Text] -> DataFrame -> Int -> Row-mkRowFromArgs names df i = V.map get (V.fromList names)-  where-    get name = case getColumn name df of-        Nothing ->-            throw $-                ColumnsNotFoundException-                    [name]-                    "[INTERNAL] mkRowFromArgs"-                    (M.keys $ columnIndices df)-        Just (BoxedColumn bm column) -> cellAny bm i (column V.! i)-        Just (UnboxedColumn bm column) -> cellAny bm i (column VU.! i)-        Just (PackedText bm p) -> cellAny bm i (packedIndexText p i)---- This function will return the items in the order that is specified--- by the user. For example, if the dataframe consists of the columns--- "Age", "Pclass", "Name", and the user asks for ["Name", "Age"],--- this will order the values in the order ["Mr Smith", 50]-mkRowRep :: DataFrame -> [T.Text] -> Int -> Row-mkRowRep df names i = V.generate (L.length names) (\index -> get (names' V.! index))-  where-    names' = V.fromList names-    throwError name =-        error $-            "Column "-                ++ T.unpack name-                ++ " has less items than "-                ++ "the other columns at index "-                ++ show i-    get name = case getColumn name df of-        Just (BoxedColumn bm c) -> case c V.!? i of-            Just e -> cellAny bm i e-            Nothing -> throwError name-        Just (UnboxedColumn bm c) -> case c VU.!? i of-            Just e -> cellAny bm i e-            Nothing -> throwError name-        Just (PackedText bm p)-            | i < packedLength p -> cellAny bm i (packedIndexText p i)-            | otherwise -> throwError name-        Nothing ->-            throw $ ColumnsNotFoundException [name] "mkRowRep" (M.keys $ columnIndices df)
− src/DataFrame/Internal/RowHash.hs
@@ -1,232 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE LambdaCase #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--{- | Row-hash kernels with a parallel driver.--The per-row key hash is the sole input to grouping and the join build/probe.-For a single wide pass over many rows (notably a 1e7-row text/factor join key)-the hashing is the dominant cost and is embarrassingly parallel: each row's hash-depends only on that row's own bytes, so hashing disjoint row ranges into-disjoint slots of one shared vector is race-free and produces a result-/bit-for-bit identical/ to the sequential single-pass hash.--'hashRowRange' is the shared per-range kernel (used sequentially and by every-worker); 'computeRowHashesIO' forks one worker per capability over contiguous-row ranges above 'parRowHashThreshold', else runs the range once. The mixing per-column type mirrors the grouping hash exactly so grouping and joins agree.--}-module DataFrame.Internal.RowHash (-    computeRowHashesIO,-    hashRowRange,-    parRowHashThreshold,-) where--import Control.Concurrent (forkIO, getNumCapabilities)-import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)-import Control.Exception (SomeException, throwIO, try)-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 System.IO.Unsafe (unsafePerformIO)-import Type.Reflection (typeRep)--import DataFrame.Internal.Column (Bitmap, Column (..), bitmapTestBit)-import DataFrame.Internal.Hash (-    fnvOffset,-    mixBytes,-    mixDouble,-    mixInt,-    mixShow,-    mixText,-    nullSalt,- )-import DataFrame.Internal.PackedText (-    PackedTextData (..),-    packedSlice,- )-import DataFrame.Internal.Types (-    SBool (..),-    sFloating,-    sIntegral,- )--{- | At least this many rows make the fork/coordination overhead of the parallel-hash worth it. Below it the sequential single range is used. Matches the-grouping/join parallel thresholds so the whole pipeline switches together.--}-parRowHashThreshold :: Int-parRowHashThreshold = 200000--capabilities :: Int-capabilities = unsafePerformIO getNumCapabilities-{-# NOINLINE capabilities #-}--{- | Compute the per-row key hash over the (already selected) key columns of an-@n@-row frame. Forks one worker per capability over contiguous row ranges when-the row count justifies it (>= 'parRowHashThreshold' and more than one-capability); otherwise hashes the single full range. The output is identical for-any capability count: each row's hash is a pure function of its own bytes and-workers own disjoint row ranges.--}-computeRowHashesIO :: Int -> [Column] -> IO (VU.Vector Int)-computeRowHashesIO n selected = do-    mv <- VUM.unsafeNew (max 1 n)-    let runRange lo hi = hashRowRange mv lo hi selected-    if n >= parRowHashThreshold && capabilities > 1-        then 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 (runRange lo hi) >>= putMVar var)-                    pure var-            vars <- mapM spawn [0 .. caps - 1]-            rs <- mapM takeMVar vars-            mapM_ (either (throwIO @SomeException) pure) rs-        else runRange 0 n-    VU.unsafeFreeze (VUM.slice 0 n mv)--{- | Mix every selected column over the row range @[lo, hi)@ into @mv@, seeding-each slot with 'fnvOffset' first. The seeding and per-column mixing must match-'DataFrame.Operations.Aggregation.computeRowHashes' byte-for-byte so grouping-and joins bucket identically.--}-hashRowRange :: VUM.IOVector Int -> Int -> Int -> [Column] -> IO ()-hashRowRange mv lo hi cols = do-    seedRange mv lo hi-    mapM_ (mixColumnRange mv lo hi) cols--seedRange :: VUM.IOVector Int -> Int -> Int -> IO ()-seedRange mv lo hi = go lo-  where-    go !i-        | i >= hi = pure ()-        | otherwise = VUM.unsafeWrite mv i fnvOffset >> go (i + 1)--{- | Fold one column's values over @[lo, hi)@ into the running hashes. The branch-structure mirrors the sequential grouping hash: typed unboxed fast paths, then a-'mixShow' fallback, with the null bitmap mixing 'nullSalt'.--}-mixColumnRange :: VUM.IOVector Int -> Int -> Int -> Column -> IO ()-mixColumnRange mv lo hi = \case-    UnboxedColumn ubm (v :: VU.Vector a) ->-        case testEquality (typeRep @a) (typeRep @Int) of-            Just Refl -> unboxedRange mv lo hi ubm mixInt v-            Nothing ->-                case testEquality (typeRep @a) (typeRep @Double) of-                    Just Refl -> unboxedRange mv lo hi ubm mixDouble v-                    Nothing ->-                        case sIntegral @a of-                            STrue ->-                                unboxedRange mv lo hi ubm (\h d -> mixInt h (fromIntegral @a @Int d)) v-                            SFalse ->-                                case sFloating @a of-                                    STrue ->-                                        unboxedRange mv lo hi ubm (\h d -> mixDouble h (realToFrac d :: Double)) v-                                    SFalse ->-                                        unboxedRange mv lo hi ubm mixShow v-    BoxedColumn bm (v :: V.Vector a) ->-        case testEquality (typeRep @a) (typeRep @T.Text) of-            Just Refl -> boxedRange mv lo hi bm mixText v-            Nothing -> boxedRange mv lo hi bm mixShow v-    PackedText bm p -> packedRange mv lo hi bm p--{- | Mix an unboxed column's range, mixing 'nullSalt' at null slots. @INLINE@d to-specialise on the element type and mixing function per call site.--}-unboxedRange ::-    (VU.Unbox a) =>-    VUM.IOVector Int ->-    Int ->-    Int ->-    Maybe Bitmap ->-    (Int -> a -> Int) ->-    VU.Vector a ->-    IO ()-unboxedRange mv lo hi ubm mix v = go lo-  where-    go !i-        | i >= hi = pure ()-        | otherwise = do-            h <- VUM.unsafeRead mv i-            let !h' = case ubm of-                    Just bm | not (bitmapTestBit bm i) -> mixInt h nullSalt-                    _ -> mix h (VU.unsafeIndex v i)-            VUM.unsafeWrite mv i h'-            go (i + 1)-{-# INLINE unboxedRange #-}--boxedRange ::-    VUM.IOVector Int ->-    Int ->-    Int ->-    Maybe Bitmap ->-    (Int -> a -> Int) ->-    V.Vector a ->-    IO ()-boxedRange mv lo hi bm mix v = go lo-  where-    go !i-        | i >= hi = pure ()-        | otherwise = do-            h <- VUM.unsafeRead mv i-            let !h' = case bm of-                    Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt-                    _ -> mix h (V.unsafeIndex v i)-            VUM.unsafeWrite mv i h'-            go (i + 1)-{-# INLINE boxedRange #-}--{- | Mix a packed-text column's range over its raw UTF-8 byte slices. The-contiguous (unselected) payload is the hot path: hoist the byte buffer and the-@n+1@ offset vector out of the loop and index them directly, so each row mixes-@[offs!i, offs!(i+1))@ with no per-row selection 'Maybe' test or 'packedSlice'-tuple. A selected payload (a gather/join result) falls back to 'packedSlice'.--}-packedRange ::-    VUM.IOVector Int ->-    Int ->-    Int ->-    Maybe Bitmap ->-    PackedTextData ->-    IO ()-packedRange mv lo hi bm p =-    case ptSel p of-        Nothing -> contiguous (ptBytes p) (ptOffsets p)-        Just _ -> selected-  where-    valid i = case bm of-        Just bm' -> bitmapTestBit bm' i-        Nothing -> True-    contiguous !arr !offs = go lo-      where-        go !i-            | i >= hi = pure ()-            | otherwise = do-                h <- VUM.unsafeRead mv i-                let !o = VU.unsafeIndex offs i-                    !l = VU.unsafeIndex offs (i + 1) - o-                    !h' = if valid i then mixBytes h arr o l else mixInt h nullSalt-                VUM.unsafeWrite mv i h'-                go (i + 1)-    selected = go lo-      where-        go !i-            | i >= hi = pure ()-            | otherwise = do-                h <- VUM.unsafeRead mv i-                let !h' =-                        if valid i-                            then let (arr, o, l) = packedSlice p i in mixBytes h arr o l-                            else mixInt h nullSalt-                VUM.unsafeWrite mv i h'-                go (i + 1)-{-# INLINE packedRange #-}
− src/DataFrame/Internal/Simplify.hs
@@ -1,422 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE MultiWayIf #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE PatternSynonyms #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.Internal.Simplify (-    simplify,-    simplifyPredicatePair,--    -- * Path-condition entailment (for fitted-tree pruning)-    PredFact,-    factTrue,-    factFalse,-    entails,-) where--import Control.Monad (guard)-import Data.Maybe (fromMaybe)-import Data.Type.Equality (testEquality, (:~:) (Refl))-import Type.Reflection (eqTypeRep, typeRep, (:~~:) (HRefl), pattern App)--import DataFrame.Internal.Column (Columnable)-import DataFrame.Internal.Expression (-    BinaryOp,-    Expr (..),-    UnaryOp (unaryName),-    eqExpr,-    normalize,- )-import DataFrame.Operators (-    NullAnd,-    NullEq,-    NullGeq,-    NullGt,-    NullLeq,-    NullLt,-    NullNeq,-    NullOr,-    (.==.),- )--simplify :: forall a. (Columnable a) => Expr a -> Expr a-simplify e-    | isBoolish @a = fixpoint (10 :: Int) e-    | otherwise = e-  where-    fixpoint 0 x = x-    fixpoint n x = let x' = simplifyB x in if eqExpr x x' then x else fixpoint (n - 1) x'--isBoolish :: forall a. (Columnable a) => Bool-isBoolish =-    case ( testEquality (typeRep @a) (typeRep @Bool)-         , testEquality (typeRep @a) (typeRep @(Maybe Bool))-         ) of-        (Just Refl, _) -> True-        (_, Just Refl) -> True-        _ -> False--data Conn = ConnAnd | ConnOr--connOf :: forall op c b r. (BinaryOp op) => op c b r -> Maybe Conn-connOf _-    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullAnd) = Just ConnAnd-    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullOr) = Just ConnOr-    | otherwise = Nothing--simplifyB :: forall a. (Columnable a) => Expr a -> Expr a-simplifyB expr = case expr of-    Binary (op :: op c b a) l r-        | Just conn <- connOf op-        , Just Refl <- testEquality (typeRep @c) (typeRep @a)-        , Just Refl <- testEquality (typeRep @b) (typeRep @a) ->-            let l' = simplifyB l; r' = simplifyB r-             in fromMaybe (Binary op l' r') (combine conn l' r')-        | otherwise -> expr-    Unary (op :: op b a) inner-        | Just Refl <- testEquality (typeRep @a) (typeRep @Bool)-        , Just Refl <- testEquality (typeRep @b) (typeRep @Bool)-        , unaryName op == "not" ->-            simplifyNot op (simplifyB inner)-        | otherwise -> expr-    If c t f ->-        let c' = simplify c-            t' = simplifyB t-            f' = simplifyB f-         in case asBoolLit c' of-                Just True -> t'-                Just False -> f'-                Nothing-                    | eqExpr t' f' -> t'-                    | Just Refl <- testEquality (typeRep @a) (typeRep @Bool)-                    , asBoolLit t' == Just True-                    , asBoolLit f' == Just False ->-                        c'-                    | otherwise -> If c' t' f'-    _ -> expr--simplifyNot :: (UnaryOp op) => op Bool Bool -> Expr Bool -> Expr Bool-simplifyNot op inner = case asBoolLit inner of-    Just b -> Lit (not b)-    Nothing -> case inner of-        Unary (op2 :: op2 b2 Bool) inner2-            | unaryName op2 == "not"-            , Just Refl <- testEquality (typeRep @b2) (typeRep @Bool) ->-                inner2-        _ -> Unary op inner--combine :: (Columnable a) => Conn -> Expr a -> Expr a -> Maybe (Expr a)-combine ConnAnd = combineAnd-combine ConnOr = combineOr--asBoolLit :: forall a. (Columnable a) => Expr a -> Maybe Bool-asBoolLit (Lit v) =-    case testEquality (typeRep @a) (typeRep @Bool) of-        Just Refl -> Just v-        Nothing -> case testEquality (typeRep @a) (typeRep @(Maybe Bool)) of-            Just Refl -> v-            Nothing -> Nothing-asBoolLit _ = Nothing--{- | Polymorphic boolean literal: @Lit b@ for @Expr Bool@, @Lit (Just b)@ for-@Expr (Maybe Bool)@.--}-litBoolish :: forall a. (Columnable a) => Bool -> Maybe (Expr a)-litBoolish v =-    case testEquality (typeRep @a) (typeRep @Bool) of-        Just Refl -> Just (Lit v)-        Nothing -> case testEquality (typeRep @a) (typeRep @(Maybe Bool)) of-            Just Refl -> Just (Lit (Just v))-            Nothing -> Nothing--combineAnd :: (Columnable a) => Expr a -> Expr a -> Maybe (Expr a)-combineAnd l r-    | eqExpr l r = Just l-    | asBoolLit l == Just False = litBoolish False-    | asBoolLit r == Just False = litBoolish False-    | asBoolLit l == Just True = Just r-    | asBoolLit r == Just True = Just l-    | absorbs ConnOr l r = Just l-    | absorbs ConnOr r l = Just r-    | otherwise = simplifyPredicatePair True l r--combineOr :: (Columnable a) => Expr a -> Expr a -> Maybe (Expr a)-combineOr l r-    | eqExpr l r = Just l-    | asBoolLit l == Just True = litBoolish True-    | asBoolLit r == Just True = litBoolish True-    | asBoolLit l == Just False = Just r-    | asBoolLit r == Just False = Just l-    | absorbs ConnAnd l r = Just l-    | absorbs ConnAnd r l = Just r-    | otherwise = simplifyPredicatePair False l r--absorbs :: (Columnable a) => Conn -> Expr a -> Expr a -> Bool-absorbs conn x (Binary (op :: op c b a) ya yb)-    | Just c' <- connOf op-    , sameConn conn c'-    , Just Refl <- testEquality (typeRep @c) (typeRep @a)-    , Just Refl <- testEquality (typeRep @b) (typeRep @a) =-        eqExpr x ya || eqExpr x yb-absorbs _ _ _ = False--sameConn :: Conn -> Conn -> Bool-sameConn ConnAnd ConnAnd = True-sameConn ConnOr ConnOr = True-sameConn _ _ = False--data Cmp = CLt | CLeq | CGt | CGeq | CEq | CNeq deriving (Eq)--data NullK = Total | FalseOnNull | UnknownOnNull deriving (Eq)--data Atom = Atom-    { aCmp :: Cmp-    , aThr :: !Double-    , aKey :: String-    , aNull :: NullK-    , aIntegral :: Bool-    }--cmpOf :: forall op c b r. (BinaryOp op) => op c b r -> Maybe Cmp-cmpOf _-    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullLt) = Just CLt-    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullLeq) = Just CLeq-    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullGt) = Just CGt-    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullGeq) = Just CGeq-    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullEq) = Just CEq-    | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullNeq) = Just CNeq-    | otherwise = Nothing--isLower, isUpper :: Cmp -> Bool-isLower c = c == CGt || c == CGeq-isUpper c = c == CLt || c == CLeq---- | True if @x@ is a @Maybe _@ type.-isMaybeTy :: forall x. (Columnable x) => Bool-isMaybeTy = case typeRep @x of-    App con _ -> case eqTypeRep con (typeRep @Maybe) of Just HRefl -> True; _ -> False-    _ -> False--litDouble :: forall b. (Columnable b) => Expr b -> Maybe Double-litDouble (Lit v) =-    case testEquality (typeRep @b) (typeRep @Double) of-        Just Refl -> Just v-        Nothing -> case testEquality (typeRep @b) (typeRep @Int) of-            Just Refl -> Just (fromIntegral v)-            Nothing -> case testEquality (typeRep @b) (typeRep @(Maybe Double)) of-                Just Refl -> v-                Nothing -> case testEquality (typeRep @b) (typeRep @(Maybe Int)) of-                    Just Refl -> fromIntegral <$> v-                    Nothing -> Nothing-litDouble _ = Nothing--{- | True for a column lifted from an integral type (never NaN): @toDouble (col …)@-or a column whose type is itself integral.--}-integralColE :: forall c. (Columnable c) => Expr c -> Bool-integralColE (Unary op _) = unaryName op == "toDouble"-integralColE _ =-    or-        [ matches @Int-        , matches @(Maybe Int)-        ]-  where-    matches :: forall t. (Columnable t) => Bool-    matches = case testEquality (typeRep @c) (typeRep @t) of Just Refl -> True; _ -> False--atomOf :: forall a. (Columnable a) => Expr a -> Maybe Atom-atomOf (Unary fm (Binary (op :: op c b r) (colE :: Expr c) litE))-    | unaryName fm == "fromMaybe"-    , Just cmp <- cmpOf op-    , Just t <- litDouble litE =-        Just (Atom cmp t (show (normalize colE)) FalseOnNull (integralColE colE))-atomOf (Binary (op :: op c b a) (colE :: Expr c) litE)-    | Just cmp <- cmpOf op-    , Just t <- litDouble litE =-        let nk = if isMaybeTy @c then UnknownOnNull else Total-         in Just (Atom cmp t (show (normalize colE)) nk (integralColE colE))-atomOf _ = Nothing--simplifyPredicatePair ::-    forall a. (Columnable a) => Bool -> Expr a -> Expr a -> Maybe (Expr a)-simplifyPredicatePair isAnd a b = do-    atomA <- atomOf a-    atomB <- atomOf b-    guard (aKey atomA == aKey atomB)-    let nk = aNull atomA-        integral = aIntegral atomA-    if isAnd-        then andAtoms a atomA b atomB nk integral-        else orAtoms a atomA b atomB nk integral---- | Contradiction folds to a literal False unless null-rows make it unknown.-litFalseGated :: (Columnable a) => NullK -> Maybe (Expr a)-litFalseGated UnknownOnNull = Nothing-litFalseGated _ = litBoolish False--{- | Tautology to literal True is sound only for total (never-null) atoms; the-exhaustive-cover form additionally needs a non-NaN (integral) column.--}-litTrueTotal :: (Columnable a) => NullK -> Maybe (Expr a)-litTrueTotal Total = litBoolish True-litTrueTotal _ = Nothing--andAtoms ::-    (Columnable a) =>-    Expr a -> Atom -> Expr a -> Atom -> NullK -> Bool -> Maybe (Expr a)-andAtoms a atomA b atomB nk _ =-    let cA = aCmp atomA; tA = aThr atomA; cB = aCmp atomB; tB = aThr atomB-     in if-            | isLower cA, isLower cB, cA == cB -> Just (if tA >= tB then a else b)-            | isUpper cA, isUpper cB, cA == cB -> Just (if tA <= tB then a else b)-            | isLower cA, isUpper cB -> lu cA tA cB tB-            | isUpper cA, isLower cB -> lu cB tB cA tA-            | cA == CEq, cB == CEq -> if tA == tB then Just a else litFalseGated nk-            | cA == CEq, cB == CNeq -> if tA == tB then litFalseGated nk else Just a-            | cA == CNeq, cB == CEq -> if tA == tB then litFalseGated nk else Just b-            | cA == CEq -> if satisfies tA cB tB then Just a else litFalseGated nk-            | cB == CEq -> if satisfies tB cA tA then Just b else litFalseGated nk-            | cA == CNeq, cB == CNeq -> Nothing-            | cA == CNeq -> if outside tA cB tB then Just b else Nothing-            | cB == CNeq -> if outside tB cA tA then Just a else Nothing-            | otherwise -> Nothing-  where-    lu lc lo uc hi-        | lo > hi = litFalseGated nk-        | lo == hi, lc == CGeq, uc == CLeq = pointEq a lo-        | lo == hi = litFalseGated nk-        | otherwise = Nothing--orAtoms ::-    (Columnable a) =>-    Expr a -> Atom -> Expr a -> Atom -> NullK -> Bool -> Maybe (Expr a)-orAtoms a atomA b atomB nk integral =-    let cA = aCmp atomA; tA = aThr atomA; cB = aCmp atomB; tB = aThr atomB-     in if-            | isLower cA, isLower cB, cA == cB -> Just (if tA <= tB then a else b)-            | isUpper cA, isUpper cB, cA == cB -> Just (if tA >= tB then a else b)-            | isUpper cA-            , isLower cB-            , nk == Total-            , integral-            , covers cB tB cA tA ->-                litTrueTotal nk-            | isLower cA-            , isUpper cB-            , nk == Total-            , integral-            , covers cA tA cB tB ->-                litTrueTotal nk-            | cA == CNeq, cB == CNeq -> if tA == tB then Just a else litTrueTotal nk-            | cA == CEq, cB == CNeq -> if tA == tB then litTrueTotal nk else Just b-            | cA == CNeq, cB == CEq -> if tA == tB then litTrueTotal nk else Just a-            | cA == CEq, cB == CEq -> if tA == tB then Just a else Nothing-            | otherwise -> Nothing--{- | Build @col == t@ for the point-collapse rule; only strict @Expr Bool@ over a-@Double@ column (otherwise bail).--}-pointEq :: forall a. (Columnable a) => Expr a -> Double -> Maybe (Expr a)-pointEq atom lo = case testEquality (typeRep @a) (typeRep @Bool) of-    Just Refl -> (\colE -> colE .==. Lit lo) <$> recoverColD atom-    Nothing -> Nothing--recoverColD :: Expr x -> Maybe (Expr Double)-recoverColD (Binary _ (colE :: Expr c) _) =-    case testEquality (typeRep @c) (typeRep @Double) of-        Just Refl -> Just colE-        _ -> Nothing-recoverColD (Unary _ inner) = recoverColD inner-recoverColD _ = Nothing--covers :: Cmp -> Double -> Cmp -> Double -> Bool-covers lowerCmp lo upperCmp hi =-    lo < hi || (lo == hi && (lowerCmp == CGeq || upperCmp == CLeq))--satisfies :: Double -> Cmp -> Double -> Bool-satisfies t CGt tb = t > tb-satisfies t CGeq tb = t >= tb-satisfies t CLt tb = t < tb-satisfies t CLeq tb = t <= tb-satisfies _ _ _ = False--outside :: Double -> Cmp -> Double -> Bool-outside t CGt tb = t <= tb-outside t CGeq tb = t < tb-outside t CLt tb = t >= tb-outside t CLeq tb = t > tb-outside _ _ _ = False---- ------------------------------------------------------------------------------ Path-condition entailment for fitted-tree pruning.--- ------------------------------------------------------------------------------- | A known same-column threshold fact accumulated along a tree path.-data PredFact = PredFact !String !Cmp !Double---- | The fact a branch's true edge establishes (the condition holds).-factTrue :: Expr Bool -> Maybe PredFact-factTrue e = (\a -> PredFact (aKey a) (aCmp a) (aThr a)) <$> atomOf e--{- | The fact a branch's false edge establishes (the negated condition). Only-sound for non-NaN (integral) columns — a NaN row takes the false edge too,-so @¬(x>t)@ is not a clean @x<=t@ bound for floats.--}-factFalse :: Expr Bool -> Maybe PredFact-factFalse e = do-    a <- atomOf e-    guard (aIntegral a && aNull a == Total)-    nc <- negCmp (aCmp a)-    pure (PredFact (aKey a) nc (aThr a))--negCmp :: Cmp -> Maybe Cmp-negCmp CLt = Just CGeq-negCmp CLeq = Just CGt-negCmp CGt = Just CLeq-negCmp CGeq = Just CLt-negCmp _ = Nothing--{- | @entails facts cond@: 'Just' 'True' when the path facts force @cond@ true,-'Just' 'False' when they force it false, 'Nothing' when undecided.--}-entails :: [PredFact] -> Expr Bool -> Maybe Bool-entails facts cond = do-    a <- atomOf cond-    let decisions =-            [ d-            | PredFact fk fc ft <- facts-            , fk == aKey a-            , Just d <- [factImplies (fc, ft) (aCmp a, aThr a)]-            ]-    case decisions of-        (d : _) -> Just d-        [] -> Nothing--{- | Does the fact's solution set sit inside @cond@ ('Just' 'True'), disjoint-from it ('Just' 'False'), or neither ('Nothing')? Boundary strictness is-honoured: e.g. @x<=t@ does NOT entail @x<t@, and @x>=t ∧ x<=t@ is not empty.--}-factImplies :: (Cmp, Double) -> (Cmp, Double) -> Maybe Bool-factImplies (fc, ft) (cc, tc)-    | isLower fc, isLower cc, subset = Just True-    | isUpper fc, isUpper cc, subset = Just True-    | isLower fc, isUpper cc, disjointAtEq = Just False-    | isUpper fc, isLower cc, disjointBelow = Just False-    | otherwise = Nothing-  where-    fIncl = fc == CGeq || fc == CLeq-    cIncl = cc == CGeq || cc == CLeq-    -- same-direction containment: strictly tighter, or equal threshold where the-    -- fact's boundary inclusivity is no stronger than the condition's.-    subset =-        (if isLower fc then ft > tc else ft < tc)-            || (ft == tc && (not fIncl || cIncl))-    -- lower fact ∩ upper cond empty: fact starts above cond's top, or they meet-    -- at a point that is not in both.-    disjointAtEq = ft > tc || (ft == tc && not (fIncl && cIncl))-    -- upper fact ∩ lower cond empty (mirror).-    disjointBelow = ft < tc || (ft == tc && not (fIncl && cIncl))
− src/DataFrame/Internal/Types.hs
@@ -1,161 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE ConstraintKinds #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE DeriveTraversable #-}-{-# LANGUAGE ExistentialQuantification #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE PolyKinds #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE UndecidableInstances #-}--module DataFrame.Internal.Types where--import Data.Int (Int16, Int32, Int64, Int8)-import Data.Kind (Constraint, Type)-import Data.Typeable (Typeable)-import qualified Data.Vector.Unboxed as VU-import Data.Word (Word16, Word32, Word64, Word8)--type Columnable' a = (Typeable a, Show a, Eq a)--{- | Inline replacement for @Data.These.These@ to keep @dataframe-core@ free-of the @these@ package dependency. Only the three constructors and the-derived classes are used internally.--}-data These a b = This a | That b | These a b-    deriving (Eq, Ord, Show, Read, Functor, Foldable, Traversable)--{- | A type with column representations used to select the-"right" representation when specializing the `toColumn` function.--}-data Rep-    = RBoxed-    | RUnboxed-    | RNullableBoxed---- | Type-level if statement.-type family If (cond :: Bool) (yes :: k) (no :: k) :: k where-    If 'True yes _ = yes-    If 'False _ no = no---- | All unboxable types (according to the `vector` package).-type family Unboxable (a :: Type) :: Bool where-    Unboxable Int = 'True-    Unboxable Int8 = 'True-    Unboxable Int16 = 'True-    Unboxable Int32 = 'True-    Unboxable Int64 = 'True-    Unboxable Word = 'True-    Unboxable Word8 = 'True-    Unboxable Word16 = 'True-    Unboxable Word32 = 'True-    Unboxable Word64 = 'True-    Unboxable Char = 'True-    Unboxable Bool = 'True-    Unboxable Double = 'True-    Unboxable Float = 'True-    Unboxable _ = 'False--type family Numeric (a :: Type) :: Bool where-    Numeric Integer = 'True-    Numeric Int = 'True-    Numeric Int8 = 'True-    Numeric Int16 = 'True-    Numeric Int32 = 'True-    Numeric Int64 = 'True-    Numeric Word = 'True-    Numeric Word8 = 'True-    Numeric Word16 = 'True-    Numeric Word32 = 'True-    Numeric Word64 = 'True-    Numeric Double = 'True-    Numeric Float = 'True-    Numeric _ = 'False---- | Compute the column representation tag for any 'a'.-type family KindOf a :: Rep where-    KindOf (Maybe a) = 'RNullableBoxed-    KindOf a = If (Unboxable a) 'RUnboxed 'RBoxed---- | Type-level boolean for constraint/type comparison.-data SBool (b :: Bool) where-    STrue :: SBool 'True-    SFalse :: SBool 'False---- | The runtime witness for our type-level branching.-class SBoolI (b :: Bool) where-    sbool :: SBool b--instance SBoolI 'True where sbool = STrue-instance SBoolI 'False where sbool = SFalse---- | Type-level function to determine whether or not a type is unboxa-sUnbox :: forall a. (SBoolI (Unboxable a)) => SBool (Unboxable a)-sUnbox = sbool @(Unboxable a)--sNumeric :: forall a. (SBoolI (Numeric a)) => SBool (Numeric a)-sNumeric = sbool @(Numeric a)--type family When (flag :: Bool) (c :: Constraint) :: Constraint where-    When 'True c = c-    When 'False c = () -- empty constraint--type UnboxIf a = When (Unboxable a) (VU.Unbox a)--type family IntegralTypes (a :: Type) :: Bool where-    IntegralTypes Integer = 'True-    IntegralTypes Int = 'True-    IntegralTypes Int8 = 'True-    IntegralTypes Int16 = 'True-    IntegralTypes Int32 = 'True-    IntegralTypes Int64 = 'True-    IntegralTypes Word = 'True-    IntegralTypes Word8 = 'True-    IntegralTypes Word16 = 'True-    IntegralTypes Word32 = 'True-    IntegralTypes Word64 = 'True-    IntegralTypes _ = 'False--sIntegral :: forall a. (SBoolI (IntegralTypes a)) => SBool (IntegralTypes a)-sIntegral = sbool @(IntegralTypes a)--type IntegralIf a = When (IntegralTypes a) (Integral a)--type family FloatingTypes (a :: Type) :: Bool where-    FloatingTypes Float = 'True-    FloatingTypes Double = 'True-    FloatingTypes _ = 'False--sFloating :: forall a. (SBoolI (FloatingTypes a)) => SBool (FloatingTypes a)-sFloating = sbool @(FloatingTypes a)--type FloatingIf a = When (FloatingTypes a) (Real a, Fractional a)--{- | Numeric type promotion: resolves the common type for mixed arithmetic.-Double dominates over Float/Int; Float dominates over Int; same types stay unchanged.--}-type family Promote (a :: Type) (b :: Type) :: Type where-    Promote a a = a-    Promote Double _ = Double-    Promote _ Double = Double-    Promote Float _ = Float-    Promote _ Float = Float-    Promote Int64 _ = Int64-    Promote _ Int64 = Int64-    Promote Int32 _ = Int32-    Promote _ Int32 = Int32-    Promote a _ = a--{- | Like 'Promote', but integral × integral → Double for use with './' .-Double\/Float still dominate; any two integral types (same or mixed) become Double.--}-type family PromoteDiv (a :: Type) (b :: Type) :: Type where-    PromoteDiv Double _ = Double-    PromoteDiv _ Double = Double-    PromoteDiv Float _ = Float-    PromoteDiv _ Float = Float-    PromoteDiv _ _ = Double -- Int/Int32/Int64 in any combination
− src/DataFrame/Internal/Utf8.hs
@@ -1,98 +0,0 @@-{-# LANGUAGE BangPatterns #-}--{- | UTF-8 validation and @decodeUtf8Lenient@-parity slice decoding used by-'DataFrame.Internal.ColumnBuilder' to turn shared byte buffers into 'Text'.--}-module DataFrame.Internal.Utf8 (-    isValidUtf8Slice,-    isUtf8Boundary,-    lenientDecodeSlice,-    sliceTextVector,-) where--import qualified Data.Text as T-import qualified Data.Text.Array as A-import qualified Data.Vector as VB-import qualified Data.Vector.Mutable as VBM-import qualified Data.Vector.Unboxed as VU--import Data.Text.Internal (Text (..))-import Data.Text.Internal.Encoding.Utf8 (-    DecoderResult (..),-    utf8DecodeContinue,-    utf8DecodeStart,- )-import Data.Text.Internal.Validate (isValidUtf8ByteArray)-import Data.Word (Word8)---- | Whether @len@ bytes starting at @off@ are well-formed UTF-8.-isValidUtf8Slice :: A.Array -> Int -> Int -> Bool-isValidUtf8Slice = isValidUtf8ByteArray-{-# INLINE isValidUtf8Slice #-}--{- | Whether a byte may start a code point (i.e. is not a continuation-byte). Field slices of a valid buffer are themselves valid iff every-field starts on a boundary.--}-isUtf8Boundary :: Word8 -> Bool-isUtf8Boundary w = w < 0x80 || w >= 0xC0-{-# INLINE isUtf8Boundary #-}--{- | Decode a byte slice exactly like @decodeUtf8Lenient@: greedy decode at-each position; any byte that cannot begin a complete, valid sequence within-the slice becomes one U+FFFD and decoding resumes at the next byte.--}-lenientDecodeSlice :: A.Array -> Int -> Int -> T.Text-lenientDecodeSlice arr off len = T.pack (go off)-  where-    !end = off + len-    go !i-        | i >= end = []-        | otherwise = case tryDecode i of-            Just (c, i') -> c : go i'-            Nothing -> '\xFFFD' : go (i + 1)-    tryDecode !i = loop (utf8DecodeStart (A.unsafeIndex arr i)) (i + 1)-      where-        loop (Accept c) !j = Just (c, j)-        loop Reject _ = Nothing-        loop (Incomplete st cp) !j-            | j >= end = Nothing-            | otherwise = loop (utf8DecodeContinue (A.unsafeIndex arr j) st cp) (j + 1)--{- | Slice forced 'Text' values off a shared array; row @i@ spans bytes-@[offs!i, offs!(i+1))@. The offsets need not start at byte 0, so a row-sub-range of a larger offset vector slices independently (parallel text-merging uses this). Fast path: validate the spanned bytes once and check-every field starts on a code-point boundary. Slow path: per-field-validation with lenient decoding of invalid fields.--}-sliceTextVector :: A.Array -> VU.Vector Int -> VB.Vector T.Text-sliceTextVector arr offs = VB.create $ do-    mv <- VBM.unsafeNew n-    let fill dec = go 0-          where-            go !i-                | i >= n = pure ()-                | otherwise = do-                    let o = VU.unsafeIndex offs i-                        !t = dec o (VU.unsafeIndex offs (i + 1) - o)-                    VBM.unsafeWrite mv i t-                    go (i + 1)-    if fast then fill mkSlice else fill decodeField-    pure mv-  where-    n = VU.length offs - 1-    base = VU.unsafeIndex offs 0-    used = VU.unsafeIndex offs n-    boundariesOk !i-        | i >= n = True-        | otherwise =-            let o = VU.unsafeIndex offs i-             in (o >= used || isUtf8Boundary (A.unsafeIndex arr o))-                    && boundariesOk (i + 1)-    fast = isValidUtf8Slice arr base (used - base) && boundariesOk 0-    mkSlice o l = if l == 0 then T.empty else Text arr o l-    decodeField o l-        | l == 0 = T.empty-        | isValidUtf8Slice arr o l = Text arr o l-        | otherwise = lenientDecodeSlice arr o l
− src/DataFrame/Operators.hs
@@ -1,425 +0,0 @@-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE TypeOperators #-}--module DataFrame.Operators where--import Data.Function ((&))-import qualified Data.Text as T-import DataFrame.Internal.Column (Columnable)-import DataFrame.Internal.Expression (-    BinUDF (MkBinaryOp),-    BinaryOp (-        binaryCommutative,-        binaryFn,-        binaryName,-        binaryPrecedence,-        binarySymbol-    ),-    Expr (Binary, Col, If, Lit, Unary),-    NamedExpr,-    UExpr (UExpr),-    UnUDF (MkUnaryOp),- )-import DataFrame.Internal.Nullable (-    BaseType,-    DivWidenOp,-    NullCmpResult,-    NullLift2Op (applyNull2),-    NullableCmpOp (nullCmpOp),-    NumericWidenOp,-    WidenResult,-    WidenResultDiv,-    divArithOp,-    widenArithOp,-    widenCmpOp,- )-import DataFrame.Internal.Types (Promote, PromoteDiv)--infixr 8 .^^, .^^., .^, .^.-infixl 7 .*, ./, .*., ./.-infixl 6 .+, .-, .+., .-.-infix 4 .==, .==., .<, .<., .<=, .<=., .>=, .>=., .>, .>., ./=, ./=.-infixr 3 .&&, .&&.-infixr 2 .||, .||.-infixr 0 .=--(|>) :: a -> (a -> b) -> b-(|>) = (&)--as :: (Columnable a) => Expr a -> T.Text -> NamedExpr-as expr colName = (colName, UExpr expr)--name :: (Show a) => Expr a -> T.Text-name (Col n) = n-name other =-    error $-        "You must call `name` on a column reference. Not the expression: " ++ show other--col :: (Columnable a) => T.Text -> Expr a-col = Col--ifThenElse :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a-ifThenElse = If--lit :: (Columnable a) => a -> Expr a-lit = Lit--(.=) :: (Columnable a) => T.Text -> Expr a -> NamedExpr-(.=) = flip as--liftDecorated ::-    (Columnable a, Columnable b) =>-    (a -> b) -> T.Text -> Maybe T.Text -> Expr a -> Expr b-liftDecorated f opName rep = Unary (MkUnaryOp f opName rep)--lift2Decorated ::-    (Columnable c, Columnable b, Columnable a) =>-    (c -> b -> a) ->-    T.Text ->-    Maybe T.Text ->-    Bool ->-    Int ->-    Expr c ->-    Expr b ->-    Expr a-lift2Decorated f opName rep comm prec =-    Binary (MkBinaryOp f opName rep comm prec)--data NullEq a b c where-    NullEq ::-        ( NumericWidenOp (BaseType a) (BaseType b)-        , NullLift2Op a b Bool (NullCmpResult a b)-        , Eq (Promote (BaseType a) (BaseType b))-        ) =>-        NullEq a b (NullCmpResult a b)--data NullNeq a b c where-    NullNeq ::-        ( NumericWidenOp (BaseType a) (BaseType b)-        , NullLift2Op a b Bool (NullCmpResult a b)-        , Eq (Promote (BaseType a) (BaseType b))-        ) =>-        NullNeq a b (NullCmpResult a b)--data NullLt a b c where-    NullLt ::-        ( NumericWidenOp (BaseType a) (BaseType b)-        , NullLift2Op a b Bool (NullCmpResult a b)-        , Ord (Promote (BaseType a) (BaseType b))-        ) =>-        NullLt a b (NullCmpResult a b)--data NullGt a b c where-    NullGt ::-        ( NumericWidenOp (BaseType a) (BaseType b)-        , NullLift2Op a b Bool (NullCmpResult a b)-        , Ord (Promote (BaseType a) (BaseType b))-        ) =>-        NullGt a b (NullCmpResult a b)--data NullLeq a b c where-    NullLeq ::-        ( NumericWidenOp (BaseType a) (BaseType b)-        , NullLift2Op a b Bool (NullCmpResult a b)-        , Ord (Promote (BaseType a) (BaseType b))-        ) =>-        NullLeq a b (NullCmpResult a b)--data NullGeq a b c where-    NullGeq ::-        ( NumericWidenOp (BaseType a) (BaseType b)-        , NullLift2Op a b Bool (NullCmpResult a b)-        , Ord (Promote (BaseType a) (BaseType b))-        ) =>-        NullGeq a b (NullCmpResult a b)--data NullAnd a b c where-    NullAnd ::-        (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>-        NullAnd a b (NullCmpResult a b)--data NullOr a b c where-    NullOr ::-        (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>-        NullOr a b (NullCmpResult a b)--instance BinaryOp NullEq where-    binaryFn NullEq = applyNull2 (widenCmpOp (==))-    binaryName NullEq = "eq"-    binarySymbol NullEq = Just ".=="-    binaryCommutative NullEq = True-    binaryPrecedence NullEq = 4-instance BinaryOp NullNeq where-    binaryFn NullNeq = applyNull2 (widenCmpOp (/=))-    binaryName NullNeq = "neq"-    binarySymbol NullNeq = Just "./="-    binaryCommutative NullNeq = True-    binaryPrecedence NullNeq = 4-instance BinaryOp NullLt where-    binaryFn NullLt = applyNull2 (widenCmpOp (<))-    binaryName NullLt = "lt"-    binarySymbol NullLt = Just ".<"-    binaryPrecedence NullLt = 4-instance BinaryOp NullGt where-    binaryFn NullGt = applyNull2 (widenCmpOp (>))-    binaryName NullGt = "gt"-    binarySymbol NullGt = Just ".>"-    binaryPrecedence NullGt = 4-instance BinaryOp NullLeq where-    binaryFn NullLeq = applyNull2 (widenCmpOp (<=))-    binaryName NullLeq = "leq"-    binarySymbol NullLeq = Just ".<="-    binaryPrecedence NullLeq = 4-instance BinaryOp NullGeq where-    binaryFn NullGeq = applyNull2 (widenCmpOp (>=))-    binaryName NullGeq = "geq"-    binarySymbol NullGeq = Just ".>="-    binaryPrecedence NullGeq = 4-instance BinaryOp NullAnd where-    binaryFn NullAnd = nullCmpOp (&&)-    binaryName NullAnd = "nulland"-    binarySymbol NullAnd = Just ".&&"-    binaryCommutative NullAnd = True-    binaryPrecedence NullAnd = 3-instance BinaryOp NullOr where-    binaryFn NullOr = nullCmpOp (||)-    binaryName NullOr = "nullor"-    binarySymbol NullOr = Just ".||"-    binaryCommutative NullOr = True-    binaryPrecedence NullOr = 2--(.==.) ::-    (Columnable a, Eq a) =>-    Expr a ->-    Expr a ->-    Expr Bool-(.==.) = lift2Decorated (==) "eq" (Just ".==.") True 4--(./=.) ::-    (Columnable a, Eq a) =>-    Expr a ->-    Expr a ->-    Expr Bool-(./=.) = lift2Decorated (/=) "neq" (Just "./=.") True 4--(.<.) ::-    (Columnable a, Ord a) =>-    Expr a ->-    Expr a ->-    Expr Bool-(.<.) = lift2Decorated (<) "lt" (Just ".<.") False 4--(.>.) ::-    (Columnable a, Ord a) =>-    Expr a ->-    Expr a ->-    Expr Bool-(.>.) = lift2Decorated (>) "gt" (Just ".>.") False 4--(.<=.) ::-    (Columnable a, Ord a) =>-    Expr a ->-    Expr a ->-    Expr Bool-(.<=.) = lift2Decorated (<=) "leq" (Just ".<=.") False 4--(.>=.) ::-    (Columnable a, Ord a) =>-    Expr a ->-    Expr a ->-    Expr Bool-(.>=.) = lift2Decorated (>=) "geq" (Just ".>=.") False 4--(.+.) :: (Columnable a, Num a) => Expr a -> Expr a -> Expr a-(.+.) = (+)--(.-.) :: (Columnable a, Num a) => Expr a -> Expr a -> Expr a-(.-.) = (-)--(.*.) :: (Columnable a, Num a) => Expr a -> Expr a -> Expr a-(.*.) = (*)--(./.) :: (Columnable a, Fractional a) => Expr a -> Expr a -> Expr a-(./.) = (/)---- Nullable-aware arithmetic operators--{- | Nullable-aware addition. Works for all combinations of nullable\/non-nullable operands.-@col \@Int "x" .+ col \@(Maybe Int) "y"  -- :: Expr (Maybe Int)@--}-(.+) ::-    ( NumericWidenOp (BaseType a) (BaseType b)-    , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b)-    , Num (Promote (BaseType a) (BaseType b))-    ) =>-    Expr a ->-    Expr b ->-    Expr (WidenResult a b)-(.+) = lift2Decorated (applyNull2 (widenArithOp (+))) "nulladd" (Just ".+") True 6---- | Nullable-aware subtraction.-(.-) ::-    ( NumericWidenOp (BaseType a) (BaseType b)-    , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b)-    , Num (Promote (BaseType a) (BaseType b))-    ) =>-    Expr a ->-    Expr b ->-    Expr (WidenResult a b)-(.-) = lift2Decorated (applyNull2 (widenArithOp (-))) "nullsub" (Just ".-") False 6---- | Nullable-aware multiplication.-(.*) ::-    ( NumericWidenOp (BaseType a) (BaseType b)-    , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b)-    , Num (Promote (BaseType a) (BaseType b))-    ) =>-    Expr a ->-    Expr b ->-    Expr (WidenResult a b)-(.*) = lift2Decorated (applyNull2 (widenArithOp (*))) "nullmul" (Just ".*") True 7---- | Nullable-aware division. Integral operands are promoted to Double.-(./) ::-    ( DivWidenOp (BaseType a) (BaseType b)-    , NullLift2Op a b (PromoteDiv (BaseType a) (BaseType b)) (WidenResultDiv a b)-    , Fractional (PromoteDiv (BaseType a) (BaseType b))-    ) =>-    Expr a ->-    Expr b ->-    Expr (WidenResultDiv a b)-(./) = lift2Decorated (applyNull2 (divArithOp (/))) "nulldiv" (Just "./") False 7---- Nullable-aware comparison operators (three-valued logic: Nothing if either operand is Nothing)--{- | Nullable-aware equality. Widens numeric operands to their common type,-so @Expr Double .== Expr Int@ typechecks. Returns @Maybe Bool@ when either-operand is nullable.--}-(.==) ::-    ( NumericWidenOp (BaseType a) (BaseType b)-    , NullLift2Op a b Bool (NullCmpResult a b)-    , Eq (Promote (BaseType a) (BaseType b))-    ) =>-    Expr a ->-    Expr b ->-    Expr (NullCmpResult a b)-(.==) = Binary NullEq---- | Nullable-aware inequality. Widens numeric operands to their common type.-(./=) ::-    ( NumericWidenOp (BaseType a) (BaseType b)-    , NullLift2Op a b Bool (NullCmpResult a b)-    , Eq (Promote (BaseType a) (BaseType b))-    ) =>-    Expr a ->-    Expr b ->-    Expr (NullCmpResult a b)-(./=) = Binary NullNeq---- | Nullable-aware less-than. Widens numeric operands to their common type.-(.<) ::-    ( NumericWidenOp (BaseType a) (BaseType b)-    , NullLift2Op a b Bool (NullCmpResult a b)-    , Ord (Promote (BaseType a) (BaseType b))-    ) =>-    Expr a ->-    Expr b ->-    Expr (NullCmpResult a b)-(.<) = Binary NullLt---- | Nullable-aware greater-than. Widens numeric operands to their common type.-(.>) ::-    ( NumericWidenOp (BaseType a) (BaseType b)-    , NullLift2Op a b Bool (NullCmpResult a b)-    , Ord (Promote (BaseType a) (BaseType b))-    ) =>-    Expr a ->-    Expr b ->-    Expr (NullCmpResult a b)-(.>) = Binary NullGt--{- | Nullable-aware less-than-or-equal. Widens numeric operands to their-common type, so @Expr Double .<= Expr Int@ typechecks.--}-(.<=) ::-    ( NumericWidenOp (BaseType a) (BaseType b)-    , NullLift2Op a b Bool (NullCmpResult a b)-    , Ord (Promote (BaseType a) (BaseType b))-    ) =>-    Expr a ->-    Expr b ->-    Expr (NullCmpResult a b)-(.<=) = Binary NullLeq---- | Nullable-aware greater-than-or-equal. Widens numeric operands to their common type.-(.>=) ::-    ( NumericWidenOp (BaseType a) (BaseType b)-    , NullLift2Op a b Bool (NullCmpResult a b)-    , Ord (Promote (BaseType a) (BaseType b))-    ) =>-    Expr a ->-    Expr b ->-    Expr (NullCmpResult a b)-(.>=) = Binary NullGeq--(.&&.) :: Expr Bool -> Expr Bool -> Expr Bool-(.&&.) = lift2Decorated (&&) "and" (Just ".&&.") True 3--(.||.) :: Expr Bool -> Expr Bool -> Expr Bool-(.||.) = lift2Decorated (||) "or" (Just ".||.") True 2---- | Nullable-aware logical AND. Returns @Maybe Bool@ when either operand is nullable.-(.&&) ::-    (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>-    Expr a ->-    Expr b ->-    Expr (NullCmpResult a b)-(.&&) = Binary NullAnd---- | Nullable-aware logical OR. Returns @Maybe Bool@ when either operand is nullable.-(.||) ::-    (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>-    Expr a ->-    Expr b ->-    Expr (NullCmpResult a b)-(.||) = Binary NullOr--(.^^) ::-    ( Columnable (BaseType a)-    , Columnable (BaseType b)-    , Fractional (BaseType a)-    , Integral (BaseType b)-    , NumericWidenOp (BaseType a) (BaseType b)-    , NullLift2Op a b (BaseType a) a-    , Num (Promote (BaseType a) (BaseType b))-    ) =>-    Expr a -> Expr b -> Expr a-(.^^) = lift2Decorated (applyNull2 (^^)) "pow" (Just ".^^") False 8--(.^) ::-    ( Columnable (BaseType a)-    , Columnable (BaseType b)-    , Num (BaseType a)-    , Integral (BaseType b)-    , NumericWidenOp (BaseType a) (BaseType b)-    , NullLift2Op a b (BaseType a) a-    , Num (Promote (BaseType a) (BaseType b))-    ) =>-    Expr a -> Expr b -> Expr a-(.^) = lift2Decorated (applyNull2 (^)) "pow" (Just ".^") False 8---- Same-type (non-nullable) exponentiation operators--(.^^.) ::-    (Columnable a, Columnable b, Fractional a, Integral b) =>-    Expr a -> Expr b -> Expr a-(.^^.) = lift2Decorated (^^) "pow" (Just ".^^.") False 8--(.^.) ::-    (Columnable a, Columnable b, Num a, Integral b) =>-    Expr a -> Expr b -> Expr a-(.^.) = lift2Decorated (^) "pow" (Just ".^.") False 8
src/DataFrame/Typed/Freeze.hs view
@@ -9,16 +9,22 @@     -- * Safe boundary     freeze,     freezeWithError,+    freezeOrThrow,      -- * Escape hatches     thaw,     unsafeFreeze,++    -- * Frame coercion+    ToDataFrame (..), ) where +import Control.Exception (throwIO) import qualified Data.Text as T import Type.Reflection (SomeTypeRep)  import Data.List (stripPrefix)+import DataFrame.Errors (DataFrameException (InternalException)) import qualified DataFrame.Internal.Column as C import DataFrame.Internal.DataFrame (columnNames) import qualified DataFrame.Internal.DataFrame as D@@ -43,11 +49,28 @@     Left err -> Left err     Right _ -> Right (TDF df) +{- | Validate and wrap like 'freezeWithError', but throw a 'DataFrameException'+in 'IO' on mismatch. The throwing boundary used by the typed readers+(@readCsv@ \/ @readParquet@).+-}+freezeOrThrow ::+    forall cols. (KnownSchema cols) => D.DataFrame -> IO (TypedDataFrame cols)+freezeOrThrow = either (throwIO . InternalException) pure . freezeWithError @cols+ {- | Unwrap a typed DataFrame back to the untyped representation. Always safe; discards type information. -} thaw :: TypedDataFrame cols -> D.DataFrame thaw (TDF df) = df++class ToDataFrame f where+    toDataFrame :: f -> D.DataFrame++instance ToDataFrame D.DataFrame where+    toDataFrame = id++instance ToDataFrame (TypedDataFrame cols) where+    toDataFrame = thaw  {- | Wrap an untyped DataFrame without any validation. Used internally after delegation where the library guarantees schema correctness.
src/DataFrame/Typed/Generic.hs view
@@ -86,7 +86,6 @@ import qualified DataFrame.Internal.DataFrame as D import DataFrame.Typed.Record (requireColumn) import DataFrame.Typed.Schema (Append)-import DataFrame.Typed.Types (Column) import DataFrame.Typed.Util (camelToSnake)  {- | Field-name policy applied to record selectors when computing@@ -97,15 +96,15 @@ -} data NameCase = SnakeCase | IdentityCase -{- | The schema type @[Column name ty, ...]@ derived from the 'Rep' of a+{- | The schema type @'[ '(name, ty), ...]@ derived from the 'Rep' of a record type, with the given 'NameCase' applied to each field name. -}-type family RepToSchema (nc :: NameCase) (r :: Type -> Type) :: [Type] where+type family RepToSchema (nc :: NameCase) (r :: Type -> Type) :: [(Symbol, Type)] where     RepToSchema nc (M1 D _ f) = RepToSchema nc f     RepToSchema nc (M1 C _ f) = RepToSchema nc f     RepToSchema nc (a :*: b) = Append (RepToSchema nc a) (RepToSchema nc b)     RepToSchema nc (M1 S ('MetaSel ('Just name) _ _ _) (K1 _ a)) =-        '[Column (TransformName nc name) a]+        '(TransformName nc name, a) ': '[]  type family TransformName (nc :: NameCase) (name :: Symbol) :: Symbol where     TransformName 'SnakeCase s = CamelToSnake s
src/DataFrame/Typed/Record.hs view
@@ -38,6 +38,7 @@ import DataFrame.Internal.DataFrame (fromNamedColumns) import qualified DataFrame.Internal.DataFrame as D import DataFrame.Typed.Types (TypedDataFrame (..))+import GHC.TypeLits  {- | Bridge a Haskell record type @a@ to a typed-dataframe schema. @@ -51,7 +52,7 @@ @Left err@ if a column is missing or has the wrong type. -} class HasSchema a where-    type Schema a :: [Type]+    type Schema a :: [(Symbol, Type)]     toColumns :: [a] -> [(T.Text, C.Column)]     fromColumns :: D.DataFrame -> Either T.Text [a] 
src/DataFrame/Typed/Schema.hs view
@@ -20,6 +20,7 @@     HasName,     RemoveColumn,     Impute,+    SetColumnType,     SubsetSchema,     ExcludeSchema,     RenameInSchema,@@ -32,6 +33,12 @@     AssertPresent,     AssertAllPresent,     AssertKeyTypesMatch,+    AssertDisjoint,+    AssertAllColumnsHaveType,+    AssertRealColumn,+    AllColumnsReal,+    AllDouble,+    IsRealType,     IsElem,      -- * Maybe-stripping families@@ -55,25 +62,28 @@      -- * KnownSchema class     KnownSchema (..),+    schemaColumnNames,      -- * Helpers     AllKnownSymbol (..), ) where +import Data.Int (Int16, Int32, Int64, Int8) import Data.Kind (Constraint, Type) import Data.Proxy (Proxy (..)) import qualified Data.Text as T+import qualified Data.Vector.Unboxed as VU+import Data.Word (Word16, Word32, Word64, Word8) import GHC.TypeLits import Type.Reflection (SomeTypeRep, Typeable, someTypeRep)  import DataFrame.Internal.Column (Columnable)-import DataFrame.Internal.Types (These)-import DataFrame.Typed.Types (Column)+import DataFrame.Internal.Column.Types (These)  -- | Look up the element type of a column by name.-type family Lookup (name :: Symbol) (cols :: [Type]) :: Type where-    Lookup name (Column name a ': _) = a-    Lookup name (Column _ _ ': rest) = Lookup name rest+type family Lookup (name :: Symbol) (cols :: [(Symbol, Type)]) :: Type where+    Lookup name ('(name, a) ': _) = a+    Lookup name (_ ': rest) = Lookup name rest     Lookup name '[] =         TypeError             ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")@@ -82,46 +92,56 @@ 'TypeError' when the column is not found.  Use together with 'AssertPresent' so the error fires exactly once. -}-type family SafeLookup (name :: Symbol) (cols :: [Type]) :: Type where-    SafeLookup name (Column name a ': _) = a-    SafeLookup name (Column _ _ ': rest) = SafeLookup name rest+type family SafeLookup (name :: Symbol) (cols :: [(Symbol, Type)]) :: Type where+    SafeLookup name ('(name, a) ': _) = a+    SafeLookup name (_ ': rest) = SafeLookup name rest     SafeLookup name '[] = Int  -- | Unwrap a Maybe from a type after we impute values.-type family Impute (name :: Symbol) (cols :: [Type]) :: [Type] where-    Impute name (Column name (Maybe a) ': rest) = Column name a ': rest-    Impute name (Column name _ ': rest) =+type family Impute (name :: Symbol) (cols :: [(Symbol, Type)]) :: [(Symbol, Type)] where+    Impute name ('(name, Maybe a) ': rest) = '(name, a) ': rest+    Impute name ('(name, _) ': rest) =         TypeError             ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' is not of kind Maybe *")     Impute name (col ': rest) = col ': Impute name rest     Impute name '[] = '[] +type family+    SetColumnType (name :: Symbol) (b :: Type) (cols :: [(Symbol, Type)]) ::+        [(Symbol, Type)]+    where+    SetColumnType name b ('(name, _) ': rest) = '(name, b) ': rest+    SetColumnType name b (col ': rest) = col ': SetColumnType name b rest+    SetColumnType name b '[] =+        TypeError+            ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")+ -- | Add type to the end of a list. type family Snoc (xs :: [k]) (x :: k) :: [k] where     Snoc '[] x = '[x]     Snoc (y ': ys) x = y ': Snoc ys x  -- | Check whether a column name exists in a schema (type-level Bool).-type family HasName (name :: Symbol) (cols :: [Type]) :: Bool where-    HasName name (Column name _ ': _) = 'True-    HasName name (Column _ _ ': rest) = HasName name rest+type family HasName (name :: Symbol) (cols :: [(Symbol, Type)]) :: Bool where+    HasName name ('(name, _) ': _) = 'True+    HasName name (_ ': rest) = HasName name rest     HasName name '[] = 'False  -- | Remove a column by name from a schema.-type family RemoveColumn (name :: Symbol) (cols :: [Type]) :: [Type] where-    RemoveColumn name (Column name _ ': rest) = rest+type family RemoveColumn (name :: Symbol) (cols :: [(Symbol, Type)]) :: [(Symbol, Type)] where+    RemoveColumn name ('(name, _) ': rest) = rest     RemoveColumn name (col ': rest) = col ': RemoveColumn name rest     RemoveColumn name '[] = '[]  -- | Select a subset of columns by a list of names.-type family SubsetSchema (names :: [Symbol]) (cols :: [Type]) :: [Type] where+type family SubsetSchema (names :: [Symbol]) (cols :: [(Symbol, Type)]) :: [(Symbol, Type)] where     SubsetSchema '[] cols = '[]-    SubsetSchema (n ': ns) cols = Column n (Lookup n cols) ': SubsetSchema ns cols+    SubsetSchema (n ': ns) cols = '(n, Lookup n cols) ': SubsetSchema ns cols  -- | Exclude columns by a list of names.-type family ExcludeSchema (names :: [Symbol]) (cols :: [Type]) :: [Type] where+type family ExcludeSchema (names :: [Symbol]) (cols :: [(Symbol, Type)]) :: [(Symbol, Type)] where     ExcludeSchema names '[] = '[]-    ExcludeSchema names (Column n a ': rest) =+    ExcludeSchema names ('(n, a) ': rest) =         ExcludeSchemaHelper (IsElem n names) n a names rest  type family@@ -130,12 +150,12 @@         (n :: Symbol)         (a :: Type)         (names :: [Symbol])-        (rest :: [Type]) ::-        [Type]+        (rest :: [(Symbol, Type)]) ::+        [(Symbol, Type)]     where     ExcludeSchemaHelper 'True n a names rest = ExcludeSchema names rest     ExcludeSchemaHelper 'False n a names rest =-        Column n a ': ExcludeSchema names rest+        '(n, a) ': ExcludeSchema names rest  -- | Type-level elem for Symbols type family IsElem (x :: Symbol) (xs :: [Symbol]) :: Bool where@@ -144,15 +164,21 @@     IsElem x (_ ': xs) = IsElem x xs  -- | Rename a column in the schema.-type family RenameInSchema (old :: Symbol) (new :: Symbol) (cols :: [Type]) :: [Type] where-    RenameInSchema old new (Column old a ': rest) = Column new a ': rest+type family+    RenameInSchema (old :: Symbol) (new :: Symbol) (cols :: [(Symbol, Type)]) ::+        [(Symbol, Type)]+    where+    RenameInSchema old new ('(old, a) ': rest) = '(new, a) ': rest     RenameInSchema old new (col ': rest) = col ': RenameInSchema old new rest     RenameInSchema old new '[] =         TypeError             ('Text "Cannot rename: column '" ':<>: 'Text old ':<>: 'Text "' not found")  -- | Rename multiple columns.-type family RenameManyInSchema (pairs :: [(Symbol, Symbol)]) (cols :: [Type]) :: [Type] where+type family+    RenameManyInSchema (pairs :: [(Symbol, Symbol)]) (cols :: [(Symbol, Type)]) ::+        [(Symbol, Type)]+    where     RenameManyInSchema '[] cols = cols     RenameManyInSchema ('(old, new) ': rest) cols =         RenameManyInSchema rest (RenameInSchema old new cols)@@ -163,24 +189,27 @@     Append (x ': xs) ys = x ': Append xs ys  -- | Reverse a type-level list.-type family Reverse (xs :: [Type]) :: [Type] where+type family Reverse (xs :: [(Symbol, Type)]) :: [(Symbol, Type)] where     Reverse xs = ReverseAcc xs '[] -type family ReverseAcc (xs :: [Type]) (acc :: [Type]) :: [Type] where+type family+    ReverseAcc (xs :: [(Symbol, Type)]) (acc :: [(Symbol, Type)]) ::+        [(Symbol, Type)]+    where     ReverseAcc '[] acc = acc     ReverseAcc (x ': xs) acc = ReverseAcc xs (x ': acc)  -- | Extract column names as a type-level list of Symbols.-type family ColumnNames (cols :: [Type]) :: [Symbol] where+type family ColumnNames (cols :: [(Symbol, Type)]) :: [Symbol] where     ColumnNames '[] = '[]-    ColumnNames (Column n _ ': rest) = n ': ColumnNames rest+    ColumnNames ('(n, _) ': rest) = n ': ColumnNames rest  -- | Assert that a column name is absent from the schema (for derive/insert).-type family AssertAbsent (name :: Symbol) (cols :: [Type]) :: Constraint where+type family AssertAbsent (name :: Symbol) (cols :: [(Symbol, Type)]) :: Constraint where     AssertAbsent name cols = AssertAbsentHelper name (HasName name cols) cols  type family-    AssertAbsentHelper (name :: Symbol) (found :: Bool) (cols :: [Type]) ::+    AssertAbsentHelper (name :: Symbol) (found :: Bool) (cols :: [(Symbol, Type)]) ::         Constraint     where     AssertAbsentHelper name 'False cols = ()@@ -193,11 +222,11 @@             )  -- | Assert that a column name is present in the schema.-type family AssertPresent (name :: Symbol) (cols :: [Type]) :: Constraint where+type family AssertPresent (name :: Symbol) (cols :: [(Symbol, Type)]) :: Constraint where     AssertPresent name cols = AssertPresentHelper name (HasName name cols) cols  type family-    AssertPresentHelper (name :: Symbol) (found :: Bool) (cols :: [Type]) ::+    AssertPresentHelper (name :: Symbol) (found :: Bool) (cols :: [(Symbol, Type)]) ::         Constraint     where     AssertPresentHelper name 'True cols = ()@@ -206,7 +235,7 @@             ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")  -- | Assert that a column name is present in the schema.-type family AssertAllPresent (name :: [Symbol]) (cols :: [Type]) :: Constraint where+type family AssertAllPresent (name :: [Symbol]) (cols :: [(Symbol, Type)]) :: Constraint where     AssertAllPresent (name ': rest) cols =         AssertAllPresentHelper (HasName name cols) name rest cols     AssertAllPresent '[] cols = ()@@ -216,7 +245,7 @@         (found :: Bool)         (name :: Symbol)         (rest :: [Symbol])-        (cols :: [Type]) ::+        (cols :: [(Symbol, Type)]) ::         Constraint     where     AssertAllPresentHelper 'True name rest cols = AssertAllPresent rest cols@@ -231,7 +260,10 @@ error fires exactly once. -} type family-    AssertKeyTypesMatch (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) ::+    AssertKeyTypesMatch+        (keys :: [Symbol])+        (left :: [(Symbol, Type)])+        (right :: [(Symbol, Type)]) ::         Constraint     where     AssertKeyTypesMatch '[] left right = ()@@ -258,50 +290,150 @@                 ':<>: 'Text " in the right table"             ) +type family+    AssertDisjoint (left :: [(Symbol, Type)]) (right :: [(Symbol, Type)]) ::+        Constraint+    where+    AssertDisjoint left right =+        AssertDisjointHelper (SharedNames left right) left right++type family+    AssertDisjointHelper+        (shared :: [Symbol])+        (left :: [(Symbol, Type)])+        (right :: [(Symbol, Type)]) ::+        Constraint+    where+    AssertDisjointHelper '[] left right = ()+    AssertDisjointHelper (n ': ns) left right =+        TypeError+            ( 'Text "Cannot horizontally merge: column '"+                ':<>: 'Text n+                ':<>: 'Text "' appears in both schemas"+            )++type family+    AssertAllColumnsHaveType+        (names :: [Symbol])+        (a :: Type)+        (cols :: [(Symbol, Type)]) ::+        Constraint+    where+    AssertAllColumnsHaveType '[] a cols = ()+    AssertAllColumnsHaveType (n ': ns) a cols =+        ( SafeLookup n cols ~ a+        , AssertPresent n cols+        , AssertAllColumnsHaveType ns a cols+        )++-- | Is @a@ a real, unboxed numeric type — i.e. a valid numeric-column element?+type family IsRealType (a :: Type) :: Bool where+    IsRealType Int = 'True+    IsRealType Int8 = 'True+    IsRealType Int16 = 'True+    IsRealType Int32 = 'True+    IsRealType Int64 = 'True+    IsRealType Word = 'True+    IsRealType Word8 = 'True+    IsRealType Word16 = 'True+    IsRealType Word32 = 'True+    IsRealType Word64 = 'True+    IsRealType Double = 'True+    IsRealType Float = 'True+    IsRealType _ = 'False++{- | Emit a readable compile error when the column named @name@ is not a+real-number type, naming the calling function @fn@, the column, and the type it+actually has. Used by the numeric extractors so a wrong column type reads as a+repairable message rather than a bare @No instance for Real …@.+-}+type family AssertRealColumn (fn :: Symbol) (name :: Symbol) (a :: Type) :: Constraint where+    AssertRealColumn fn name a = AssertRealColumnGo fn name a (IsRealType a)++type family+    AssertRealColumnGo (fn :: Symbol) (name :: Symbol) (a :: Type) (isReal :: Bool) ::+        Constraint+    where+    AssertRealColumnGo fn name a 'True = ()+    AssertRealColumnGo fn name a 'False =+        TypeError+            ( 'Text fn+                ':<>: 'Text ": expected a real number column for '"+                ':<>: 'Text name+                ':<>: 'Text "' but instead you gave "+                ':<>: 'ShowType a+            )++{- | Constraint that every column in the schema is a real (numeric), unboxed+type. Lets the whole-frame matrix extractors ('toDoubleMatrix' and friends) be+total — a non-numeric or nullable column is a compile error (with the offending+column named, via 'AssertRealColumn'), not a runtime 'Left'.+-}+type family AllColumnsReal (fn :: Symbol) (cols :: [(Symbol, Type)]) :: Constraint where+    AllColumnsReal fn '[] = ()+    AllColumnsReal fn ('(n, a) ': rest) =+        (AssertRealColumn fn n a, Real a, VU.Unbox a, AllColumnsReal fn rest)++-- TODO: mchavinda - we can generalist to AllX+type family AllDouble (cols :: [(Symbol, Type)]) :: Constraint where+    AllDouble '[] = ()+    AllDouble ('(n, Double) ': rest) = AllDouble rest+    AllDouble ('(n, a) ': rest) =+        TypeError+            ( 'Text "Column '"+                ':<>: 'Text n+                ':<>: 'Text "' must be Double for this model, but is "+                ':<>: 'ShowType a+                ':$$: 'Text "Convert it (toDouble) or drop it before fitting."+            )+ {- | Strip 'Maybe' from all columns. Used by 'filterAllJust'. -@Column "x" (Maybe Double)@ becomes @Column "x" Double@.-@Column "y" Int@ stays @Column "y" Int@.+@'("x", (Maybe Double)@ becomes @'("x", Double@.))+@'("y", Int@ stays @'("y", Int@.)) -}-type family StripAllMaybe (cols :: [Type]) :: [Type] where+type family StripAllMaybe (cols :: [(Symbol, Type)]) :: [(Symbol, Type)] where     StripAllMaybe '[] = '[]-    StripAllMaybe (Column n (Maybe a) ': rest) = Column n a ': StripAllMaybe rest-    StripAllMaybe (Column n a ': rest) = Column n a ': StripAllMaybe rest+    StripAllMaybe ('(n, Maybe a) ': rest) = '(n, a) ': StripAllMaybe rest+    StripAllMaybe ('(n, a) ': rest) = '(n, a) ': StripAllMaybe rest  {- | Strip 'Maybe' from a single named column. Used by 'filterJust'. -@StripMaybeAt "x" '[Column "x" (Maybe Double), Column "y" Int]@-  = @'[Column "x" Double, Column "y" Int]@+@StripMaybeAt "x" '[ '("x", Maybe Double), '("y", Int)]@+  = @'[ '("x", Double), '("y", Int)]@ -}-type family StripMaybeAt (name :: Symbol) (cols :: [Type]) :: [Type] where-    StripMaybeAt name (Column name (Maybe a) ': rest) = Column name a ': rest-    StripMaybeAt name (Column name a ': rest) = Column name a ': rest+type family StripMaybeAt (name :: Symbol) (cols :: [(Symbol, Type)]) :: [(Symbol, Type)] where+    StripMaybeAt name ('(name, Maybe a) ': rest) = '(name, a) ': rest+    StripMaybeAt name ('(name, a) ': rest) = '(name, a) ': rest     StripMaybeAt name (col ': rest) = col ': StripMaybeAt name rest     StripMaybeAt name '[] =         TypeError             ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")  -- | Extract column names that appear in both schemas.-type family SharedNames (left :: [Type]) (right :: [Type]) :: [Symbol] where+type family SharedNames (left :: [(Symbol, Type)]) (right :: [(Symbol, Type)]) :: [Symbol] where     SharedNames '[] right = '[]-    SharedNames (Column n _ ': rest) right =+    SharedNames ('(n, _) ': rest) right =         SharedNamesHelper (HasName n right) n rest right  type family     SharedNamesHelper         (found :: Bool)         (n :: Symbol)-        (rest :: [Type])-        (right :: [Type]) ::+        (rest :: [(Symbol, Type)])+        (right :: [(Symbol, Type)]) ::         [Symbol]     where     SharedNamesHelper 'True n rest right = n ': SharedNames rest right     SharedNamesHelper 'False n rest right = SharedNames rest right  -- | Columns from @left@ whose names do NOT appear in @right@.-type family UniqueLeft (left :: [Type]) (rightNames :: [Symbol]) :: [Type] where+type family+    UniqueLeft (left :: [(Symbol, Type)]) (rightNames :: [Symbol]) ::+        [(Symbol, Type)]+    where     UniqueLeft '[] _ = '[]-    UniqueLeft (Column n a ': rest) rn =+    UniqueLeft ('(n, a) ': rest) rn =         UniqueLeftHelper (IsElem n rn) n a rest rn  type family@@ -309,26 +441,29 @@         (found :: Bool)         (n :: Symbol)         (a :: Type)-        (rest :: [Type])+        (rest :: [(Symbol, Type)])         (rn :: [Symbol]) ::-        [Type]+        [(Symbol, Type)]     where     UniqueLeftHelper 'True n a rest rn = UniqueLeft rest rn-    UniqueLeftHelper 'False n a rest rn = Column n a ': UniqueLeft rest rn+    UniqueLeftHelper 'False n a rest rn = '(n, a) ': UniqueLeft rest rn  type family ToMaybe (a :: Type) :: Type where     ToMaybe (Maybe a) = Maybe a     ToMaybe a = Maybe a  -- | Wrap column types in Maybe; idempotent on already-optional columns.-type family WrapMaybe (cols :: [Type]) :: [Type] where+type family WrapMaybe (cols :: [(Symbol, Type)]) :: [(Symbol, Type)] where     WrapMaybe '[] = '[]-    WrapMaybe (Column n a ': rest) = Column n (ToMaybe a) ': WrapMaybe rest+    WrapMaybe ('(n, a) ': rest) = '(n, ToMaybe a) ': WrapMaybe rest  -- | Wrap selected columns in Maybe by name list.-type family WrapMaybeColumns (names :: [Symbol]) (cols :: [Type]) :: [Type] where+type family+    WrapMaybeColumns (names :: [Symbol]) (cols :: [(Symbol, Type)]) ::+        [(Symbol, Type)]+    where     WrapMaybeColumns names '[] = '[]-    WrapMaybeColumns names (Column n a ': rest) =+    WrapMaybeColumns names ('(n, a) ': rest) =         WrapMaybeColumnsHelper (IsElem n names) n a names rest  type family@@ -337,18 +472,24 @@         (n :: Symbol)         (a :: Type)         (names :: [Symbol])-        (rest :: [Type]) ::-        [Type]+        (rest :: [(Symbol, Type)]) ::+        [(Symbol, Type)]     where     WrapMaybeColumnsHelper 'True n a names rest =-        Column n (ToMaybe a) ': WrapMaybeColumns names rest+        '(n, ToMaybe a) ': WrapMaybeColumns names rest     WrapMaybeColumnsHelper 'False n a names rest =-        Column n a ': WrapMaybeColumns names rest+        '(n, a) ': WrapMaybeColumns names rest  -- | Columns in left whose names collide with right (excluding keys).-type family CollidingColumns (left :: [Type]) (right :: [Type]) (keys :: [Symbol]) :: [Type] where+type family+    CollidingColumns+        (left :: [(Symbol, Type)])+        (right :: [(Symbol, Type)])+        (keys :: [Symbol]) ::+        [(Symbol, Type)]+    where     CollidingColumns '[] _ _ = '[]-    CollidingColumns (Column n a ': rest) right keys =+    CollidingColumns ('(n, a) ': rest) right keys =         CollidingColumnsHelper1 (IsElem n keys) n a rest right keys  type family@@ -356,10 +497,10 @@         (isKey :: Bool)         (n :: Symbol)         (a :: Type)-        (rest :: [Type])-        (right :: [Type])+        (rest :: [(Symbol, Type)])+        (right :: [(Symbol, Type)])         (keys :: [Symbol]) ::-        [Type]+        [(Symbol, Type)]     where     CollidingColumnsHelper1 'True n a rest right keys =         CollidingColumns rest right keys@@ -371,18 +512,24 @@         (inRight :: Bool)         (n :: Symbol)         (a :: Type)-        (rest :: [Type])-        (right :: [Type])+        (rest :: [(Symbol, Type)])+        (right :: [(Symbol, Type)])         (keys :: [Symbol]) ::-        [Type]+        [(Symbol, Type)]     where     CollidingColumnsHelper2 'True n a rest right keys =-        Column n (These a (Lookup n right)) ': CollidingColumns rest right keys+        '(n, These a (Lookup n right)) ': CollidingColumns rest right keys     CollidingColumnsHelper2 'False n a rest right keys =         CollidingColumns rest right keys  -- | Inner join result schema.-type family InnerJoinSchema (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) :: [Type] where+type family+    InnerJoinSchema+        (keys :: [Symbol])+        (left :: [(Symbol, Type)])+        (right :: [(Symbol, Type)]) ::+        [(Symbol, Type)]+    where     InnerJoinSchema keys left right =         Append             (SubsetSchema keys left)@@ -395,7 +542,13 @@             )  -- | Left join result schema.-type family LeftJoinSchema (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) :: [Type] where+type family+    LeftJoinSchema+        (keys :: [Symbol])+        (left :: [(Symbol, Type)])+        (right :: [(Symbol, Type)]) ::+        [(Symbol, Type)]+    where     LeftJoinSchema keys left right =         Append             (SubsetSchema keys left)@@ -408,7 +561,13 @@             )  -- | Right join result schema.-type family RightJoinSchema (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) :: [Type] where+type family+    RightJoinSchema+        (keys :: [Symbol])+        (left :: [(Symbol, Type)])+        (right :: [(Symbol, Type)]) ::+        [(Symbol, Type)]+    where     RightJoinSchema keys left right =         Append             (SubsetSchema keys right)@@ -422,8 +581,11 @@  -- | Full outer join result schema. type family-    FullOuterJoinSchema (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) ::-        [Type]+    FullOuterJoinSchema+        (keys :: [Symbol])+        (left :: [(Symbol, Type)])+        (right :: [(Symbol, Type)]) ::+        [(Symbol, Type)]     where     FullOuterJoinSchema keys left right =         Append@@ -437,9 +599,12 @@             )  -- | Extract Column entries from a schema whose names appear in @keys@.-type family GroupKeyColumns (keys :: [Symbol]) (cols :: [Type]) :: [Type] where+type family+    GroupKeyColumns (keys :: [Symbol]) (cols :: [(Symbol, Type)]) ::+        [(Symbol, Type)]+    where     GroupKeyColumns keys '[] = '[]-    GroupKeyColumns keys (Column n a ': rest) =+    GroupKeyColumns keys ('(n, a) ': rest) =         GroupKeyColumnsHelper (IsElem n keys) n a keys rest  type family@@ -448,15 +613,15 @@         (n :: Symbol)         (a :: Type)         (keys :: [Symbol])-        (rest :: [Type]) ::-        [Type]+        (rest :: [(Symbol, Type)]) ::+        [(Symbol, Type)]     where     GroupKeyColumnsHelper 'True n a keys rest =-        Column n a ': GroupKeyColumns keys rest+        '(n, a) ': GroupKeyColumns keys rest     GroupKeyColumnsHelper 'False n a keys rest = GroupKeyColumns keys rest  -- | Provides runtime evidence of a schema: a list of (name, TypeRep) pairs.-class KnownSchema (cols :: [Type]) where+class KnownSchema (cols :: [(Symbol, Type)]) where     schemaEvidence :: [(T.Text, SomeTypeRep)]  instance KnownSchema '[] where@@ -464,11 +629,23 @@  instance     (KnownSymbol name, Typeable a, Columnable a, KnownSchema rest) =>-    KnownSchema (Column name a ': rest)+    KnownSchema ('(name, a) ': rest)     where     schemaEvidence =         (T.pack (symbolVal (Proxy @name)), someTypeRep (Proxy @a))             : schemaEvidence @rest++{- | The column names a schema declares, in schema order. Pass it to a reader's+options to fetch only those columns:++@+D.readCsvWithOpts+    D.defaultReadOptions{D.readColumns = Just (schemaColumnNames \@(Schema Customer))}+    "customers.csv"+@+-}+schemaColumnNames :: forall cols. (KnownSchema cols) => [T.Text]+schemaColumnNames = map fst (schemaEvidence @cols)  -- | A class that provides a list of 'Text' values for a type-level list of Symbols. class AllKnownSymbol (names :: [Symbol]) where
src/DataFrame/Typed/Types.hs view
@@ -2,22 +2,25 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE ExistentialQuantification #-} {-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-}-{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-}  module DataFrame.Typed.Types (     -- * Core phantom-typed wrapper     TypedDataFrame (..), -    -- * Column phantom type (no constructors)-    Column,-     -- * Typed expressions (schema-validated)     TExpr (..), +    -- * Lifting untyped expressions into the typed world+    AsTExpr,+    ToTExpr (..),+     -- * Typed sort orders     TSortOrder (..), @@ -37,17 +40,17 @@  import qualified Data.Text as T import DataFrame.Internal.Column (Columnable)+import DataFrame.Internal.Column.Types (These (..)) import qualified DataFrame.Internal.DataFrame as D import DataFrame.Internal.Expression (Expr, NamedExpr, UExpr (..))-import DataFrame.Internal.Types (These (..))  {- | A phantom-typed wrapper over the untyped 'DataFrame'. -The type parameter @cols@ is a type-level list of @Column name ty@ entries+The type parameter @cols@ is a type-level list of @'(name, ty)@ pairs that tracks the schema at compile time. All operations delegate to the untyped core at runtime and update the phantom type at compile time. -}-newtype TypedDataFrame (cols :: [Type]) = TDF {unTDF :: D.DataFrame}+newtype TypedDataFrame (cols :: [(Symbol, Type)]) = TDF {unTDF :: D.DataFrame}  instance Show (TypedDataFrame cols) where     show (TDF df) = show df@@ -55,12 +58,6 @@ instance Eq (TypedDataFrame cols) where     (TDF a) == (TDF b) = a == b -{- | A phantom type that pairs a type-level column name ('Symbol')-with its element type. Has no value-level constructors — used-purely at the type level to describe schemas.--}-data Column (name :: Symbol) (a :: Type)- {- | A typed expression validated against schema @cols@, producing values of type @a@.  Unlike the untyped 'Expr a', a 'TExpr' can only be constructed through@@ -69,15 +66,33 @@  Use 'unTExpr' to extract the underlying 'Expr' for delegation to the untyped API. -}-newtype TExpr (cols :: [Type]) a = TExpr {unTExpr :: Expr a}+newtype TExpr (cols :: [(Symbol, Type)]) a = TExpr {unTExpr :: Expr a} +-- | Shows the underlying expression; the schema phantom is type-level only.+instance (Show a) => Show (TExpr cols a) where+    showsPrec d (TExpr e) = showsPrec d e++{- | The typed counterpart of an untyped expression type for schema @cols@:+@AsTExpr cols (Expr r) = TExpr cols r@. Lets a result type follow the frame —+an @Expr@ over a plain frame becomes a @TExpr@ over a typed one.+-}+type family AsTExpr (cols :: [(Symbol, Type)]) (e :: Type) :: Type where+    AsTExpr cols (Expr r) = TExpr cols r++-- | Lift an untyped expression into its 'TExpr' for schema @cols@.+class ToTExpr (cols :: [(Symbol, Type)]) e where+    toTExpr :: e -> AsTExpr cols e++instance ToTExpr cols (Expr r) where+    toTExpr = TExpr+ -- | A typed sort order validated against schema @cols@.-data TSortOrder (cols :: [Type]) where+data TSortOrder (cols :: [(Symbol, Type)]) where     Asc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols     Desc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols  -- | A phantom-typed wrapper over 'GroupedDataFrame'.-newtype TypedGrouped (keys :: [Symbol]) (cols :: [Type])+newtype TypedGrouped (keys :: [Symbol]) (cols :: [(Symbol, Type)])     = TGD {unTGD :: D.GroupedDataFrame}  {- | Internal aggregation chain. Each cons prepends a 'Column' to the@@ -91,7 +106,7 @@   . as \@\"avg_age\" (F.mean age) @ -}-data TAgg (keys :: [Symbol]) (cols :: [Type]) (aggs :: [Type]) where+data TAgg (keys :: [Symbol]) (cols :: [(Symbol, Type)]) (aggs :: [(Symbol, Type)]) where     TAggNil :: TAgg keys cols '[]     TAggCons ::         (Columnable a) =>@@ -101,7 +116,7 @@         TExpr cols a ->         -- | rest         TAgg keys cols aggs ->-        TAgg keys cols (Column name a ': aggs)+        TAgg keys cols ('(name, a) ': aggs)  {- | Extract the runtime 'NamedExpr' list from a 'TAgg', in declaration order (reversed from the cons-built order).