packages feed

dataframe-core 1.1.0.4 → 2.6.0.0

raw patch · 79 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.6.0.0 synopsis:           Core data structures for the dataframe library. description:     Minimal interchange-format types for the @dataframe@ ecosystem:@@ -31,46 +31,66 @@ 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.Distinct+                        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.Reflect+                        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.Interpreter.Kernels                         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,104 @@+{-# 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' or 'Markdown'.+data RenderFormat = Plain | Markdown+    deriving (Show, Eq)++-- Utility functions to show a DataFrame as a Markdown-ish table.++-- Fill functions, adapted from StackOverflow question 5929377.+type Filler = Int -> T.Text -> T.Text++-- A description of one table column.+data ColDesc t = ColDesc+    { colTitleFill :: Filler+    , colTitle :: T.Text+    , colValueFill :: Filler+    }++-- Pad text to a given width with a character, aligned 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.+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.+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,572 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | The low-cardinality dense reduction kernel (direct-indexed by group).+module DataFrame.Internal.Aggregation.Kernel.Dense (+    denseReduce,+    denseMaxMinusMin,+) 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 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 float reductions run one sequential pass.+seqFloatGroups :: Int+seqFloatGroups = 65536++-- | Run a recognised reduction through the direct-indexed path, if admitted.+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.+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))+    _ -> Nothing++-- | The reductions admitted over a Double column.+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))+    _ -> Nothing++-- | The fused @max a - min b@ direct pass over both extrema at once.+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 for the 'testEquality' dispatch above.+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 workers over disjoint row ranges, each with a private accumulator.+runPartialsOver ::+    Int -> Int -> (Int -> Int -> IO acc) -> IO [acc]+runPartialsOver n _caps = parallelChunks parThreshold n++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++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 group sums, sequential at small group counts.+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 (sum, count) per group, divided once at finalize.+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++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++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++-- | One streaming pass accumulating @max a@ and @min b@ together.+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.+{-# 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++-- | Integer mean in one fused pass: running sum and count per group.+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++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 per-worker Double partials into the first worker's, in worker 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 by exact integer addition.+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 mean.+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/Distinct.hs view
@@ -0,0 +1,609 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE UnboxedTuples #-}++module DataFrame.Internal.Aggregation.Kernel.Distinct (+    distinctSupported,+    countDistinctColumn,+    countDistinctGroups,+) where++import Control.Monad (forM_, replicateM, void, when)+import Data.Bits (+    countLeadingZeros,+    finiteBitSize,+    unsafeShiftL,+    unsafeShiftR,+    xor,+    (.&.),+ )+import Data.Char (ord)+import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef, writeIORef)+import Data.Int (Int16, Int32, Int64, Int8)+import Data.Primitive.ByteArray (MutableByteArray (..))+import qualified Data.Text.Array as A+import Data.Text.Internal (Text (Text))+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))+import qualified Data.Vector as V+import qualified Data.Vector.Primitive as VP+import qualified Data.Vector.Primitive.Mutable as VPM+import qualified Data.Vector.Unboxed as VU+import Data.Vector.Unboxed.Base (+    MVector (MV_Int),+    Vector (V_Int, V_Int64, V_Word64),+ )+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.Word (Word16, Word32, Word64, Word8)+import GHC.Exts (Int (I#), fetchAddIntArray#, (+#))+import GHC.Float (castDoubleToWord64)+import GHC.IO (IO (IO))+import System.IO.Unsafe (unsafePerformIO)+import Type.Reflection (typeRep)++import DataFrame.Internal.Algorithms.Hash (fnvOffset, mixBytes)+import DataFrame.Internal.Column (Column (..))+import DataFrame.Internal.Control.Concurrent (+    capabilities,+    forkJoin_,+    parThreshold,+    shouldParallelize,+    splitChunkRange,+ )+import DataFrame.Internal.Data.PackedText (+    PackedOffsets (..),+    PackedTextData,+    packedLength,+    packedRowOffsets,+    packedSlice,+ )++-- | How a column's rows become set keys.+data Keys+    = -- | One 64-bit word per row; equal words iff equal values.+      WordKeys !(VU.Vector Int)+    | -- | A UTF-8 byte slice per row.+      ByteKeys !Int !TextRows++-- | Where a text row's bytes live, with kernel instances for each layout.+data TextRows+    = Contig32 !A.Array !(VU.Vector Int32)+    | Contig64 !A.Array !(VU.Vector Int)+    | AnyRows (Int -> (A.Array, Int, Int))++-- | Whether 'countDistinctColumn' / 'countDistinctGroups' accept this column.+distinctSupported :: Column -> Bool+distinctSupported c = case columnKeys c of+    Just _ -> True+    Nothing -> False++columnKeys :: Column -> Maybe Keys+columnKeys col = case col of+    UnboxedColumn Nothing (v :: VU.Vector a)+        | Just Refl <- testEquality (typeRep @a) (typeRep @Int) -> Just (WordKeys v)+        | Just Refl <- testEquality (typeRep @a) (typeRep @Int64) ->+            Just (WordKeys (int64Words v))+        | Just Refl <- testEquality (typeRep @a) (typeRep @Word64) ->+            Just (WordKeys (word64Words v))+        | Just Refl <- testEquality (typeRep @a) (typeRep @Int32) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Int16) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Int8) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Word) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Word32) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Word16) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Word8) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Double) ->+            Just (WordKeys (VU.map doubleWord v))+        | Just Refl <- testEquality (typeRep @a) (typeRep @Char) ->+            Just (WordKeys (VU.map ord v))+        | Just Refl <- testEquality (typeRep @a) (typeRep @Bool) ->+            Just (WordKeys (VU.map fromEnum v))+    BoxedColumn Nothing (v :: V.Vector a)+        | Just Refl <- testEquality (typeRep @a) (typeRep @Text) ->+            Just+                ( ByteKeys+                    (V.length v)+                    (AnyRows (\i -> case V.unsafeIndex v i of Text arr o l -> (arr, o, l)))+                )+    PackedText Nothing p -> Just (ByteKeys (packedLength p) (packedRows p))+    _ -> Nothing+  where+    widen :: (Integral b, VU.Unbox b) => VU.Vector b -> Maybe Keys+    widen = Just . WordKeys . VU.map fromIntegral++packedRows :: PackedTextData -> TextRows+packedRows p = case packedRowOffsets p of+    Just (arr, Offs32 offs) -> Contig32 arr offs+    Just (arr, Offs64 offs) -> Contig64 arr offs+    Nothing -> AnyRows (packedSlice p)++-- | Reinterpret 64-bit words as 'Int' without copying (same width and layout).+int64Words :: VU.Vector Int64 -> VU.Vector Int+int64Words v@(V_Int64 (VP.Vector o l ba))+    | finiteBitSize (0 :: Int) == 64 = V_Int (VP.Vector o l ba)+    | otherwise = VU.map fromIntegral v++word64Words :: VU.Vector Word64 -> VU.Vector Int+word64Words v@(V_Word64 (VP.Vector o l ba))+    | finiteBitSize (0 :: Int) == 64 = V_Int (VP.Vector o l ba)+    | otherwise = VU.map fromIntegral v++-- | Equal doubles share a word: @-0.0@ joins @0.0@ and every NaN is one value.+doubleWord :: Double -> Int+doubleWord d+    | isNaN d = 0x7ff8000000000000+    | d == 0 = 0+    | otherwise = fromIntegral (castDoubleToWord64 d)++-- | Number of distinct values in a column, or 'Nothing' if not admitted.+countDistinctColumn :: Column -> Maybe Int+countDistinctColumn col = case columnKeys col of+    Nothing -> Nothing+    Just (WordKeys v) -> Just $! unsafePerformIO (wordGlobal v)+    Just (ByteKeys n rows) -> Just $! unsafePerformIO $ case rows of+        Contig32 arr offs -> bytesGlobal n (contigSlice arr offs)+        Contig64 arr offs -> bytesGlobal n (contigSlice arr offs)+        AnyRows slice -> bytesGlobal n slice+{-# NOINLINE countDistinctColumn #-}++-- | Distinct values per group, or 'Nothing' if the column is not admitted.+countDistinctGroups :: VU.Vector Int -> Int -> Column -> Maybe (VU.Vector Int)+countDistinctGroups rtg nGroups col = case columnKeys col of+    Nothing -> Nothing+    Just (WordKeys v) -> Just $! unsafePerformIO (wordGrouped rtg nGroups v)+    Just (ByteKeys _ rows) -> Just $! unsafePerformIO $ case rows of+        Contig32 arr offs -> bytesGrouped rtg nGroups (contigSlice arr offs)+        Contig64 arr offs -> bytesGrouped rtg nGroups (contigSlice arr offs)+        AnyRows slice -> bytesGrouped rtg nGroups slice+{-# NOINLINE countDistinctGroups #-}++contigSlice ::+    (Integral o, VU.Unbox o) => A.Array -> VU.Vector o -> Int -> (A.Array, Int, Int)+contigSlice arr offs i =+    let !o = fromIntegral (VU.unsafeIndex offs i)+     in (arr, o, fromIntegral (VU.unsafeIndex offs (i + 1)) - o)+{-# INLINE contigSlice #-}++wordGlobal :: VU.Vector Int -> IO Int+wordGlobal v = do+    let !n = VU.length v+        !pb = partitionBits n+        key = VU.unsafeIndex v+    Scattered w0 _ starts <-+        scatterRows+            n+            pb+            False+            False+            (fmix . key)+            (\r r' -> key r == key r')+            (\i _ -> key i)+            (const 0)+    counts <- forPartitions (VU.length starts - 1) (newIORef (0 :: Int)) $ \acc p -> do+        c <- dedupWords pb w0 (VU.unsafeIndex starts p) (VU.unsafeIndex starts (p + 1))+        modifyStrict acc (+ c)+    sum <$> mapM readIORef counts++wordGrouped :: VU.Vector Int -> Int -> VU.Vector Int -> IO (VU.Vector Int)+wordGrouped rtg nGroups v = do+    let !n = VU.length v+        !pb = partitionBits n+        key = VU.unsafeIndex v+        gid = VU.unsafeIndex rtg+    Scattered w0 w1 starts <-+        scatterRows+            n+            pb+            False+            True+            (\i -> pairHash (gid i) (key i))+            (\r r' -> key r == key r' && gid r == gid r')+            (\i _ -> key i)+            gid+    groupCounts nGroups (VU.length starts - 1) $ \counts p ->+        void $+            dedupPairs+                pb+                w0+                w1+                (VU.unsafeIndex starts p)+                (VU.unsafeIndex starts (p + 1))+                pairHash'+                (\a0 a1 b0 b1 -> a0 == b0 && a1 == b1)+                (\_ g -> bumpGroup counts g)+  where+    pairHash' k g = pairHash g k++bytesGlobal :: Int -> (Int -> (A.Array, Int, Int)) -> IO Int+bytesGlobal n slice = do+    let !pb = partitionBits n+        hashRow i = case slice i of (arr, o, l) -> fmix (mixBytes fnvOffset arr o l)+    Scattered w0 w1 starts <-+        scatterRows n pb True True hashRow (rowsEq slice) (\_ h -> h) id+    counts <- forPartitions (VU.length starts - 1) (newIORef (0 :: Int)) $ \acc p -> do+        c <-+            dedupPairs+                pb+                w0+                w1+                (VU.unsafeIndex starts p)+                (VU.unsafeIndex starts (p + 1))+                const+                (\h r h' r' -> h == h' && rowsEq slice r r')+                (\_ _ -> pure ())+        modifyStrict acc (+ c)+    sum <$> mapM readIORef counts++bytesGrouped ::+    VU.Vector Int -> Int -> (Int -> (A.Array, Int, Int)) -> IO (VU.Vector Int)+bytesGrouped rtg nGroups slice = do+    let !n = VU.length rtg+        !pb = partitionBits n+        gid = VU.unsafeIndex rtg+        hashRow i = case slice i of+            (arr, o, l) -> pairHash (gid i) (mixBytes fnvOffset arr o l)+    Scattered w0 w1 starts <-+        scatterRows+            n+            pb+            True+            True+            hashRow+            (\r r' -> gid r == gid r' && rowsEq slice r r')+            (\_ h -> h)+            id+    groupCounts nGroups (VU.length starts - 1) $ \counts p ->+        void $+            dedupPairs+                pb+                w0+                w1+                (VU.unsafeIndex starts p)+                (VU.unsafeIndex starts (p + 1))+                const+                (\h r h' r' -> h == h' && gid r == gid r' && rowsEq slice r r')+                (\_ r -> bumpGroup counts (gid r))++rowsEq :: (Int -> (A.Array, Int, Int)) -> Int -> Int -> Bool+rowsEq slice r r' = case slice r of+    (a, o, l) -> case slice r' of+        (a', o', l') -> l == l' && A.compare a o a' o' l == EQ+{-# INLINE rowsEq #-}++{-# INLINE bytesGlobal #-}+{-# INLINE bytesGrouped #-}++-- | Murmur3's 64-bit finaliser: every input bit reaches every output bit.+fmix :: Int -> Int+fmix k =+    let !h0 = fromIntegral k :: Word64+        !h1 = (h0 `xor` (h0 `unsafeShiftR` 33)) * 0xff51afd7ed558ccd+        !h2 = (h1 `xor` (h1 `unsafeShiftR` 33)) * 0xc4ceb9fe1a85ec53+     in fromIntegral (h2 `xor` (h2 `unsafeShiftR` 33))+{-# INLINE fmix #-}++pairHash :: Int -> Int -> Int+pairHash g k = fmix (k `xor` (g * 0x9E3779B97F4A7C1))+{-# INLINE pairHash #-}++-- | Partition of a hash: its top @pb@ bits (@pb >= 1@).+partOf :: Int -> Int -> Int+partOf pb h = fromIntegral ((fromIntegral h :: Word64) `unsafeShiftR` (64 - pb))+{-# INLINE partOf #-}++-- | Slot of a hash in a @2^bits@ table: the bits just below the partition bits.+slotOf :: Int -> Int -> Int -> Int+slotOf pb bits h =+    fromIntegral+        (((fromIntegral h :: Word64) `unsafeShiftL` pb) `unsafeShiftR` (64 - bits))+{-# INLINE slotOf #-}++-- | Partition count as a bit width, sized for cache-resident sets.+partitionBits :: Int -> Int+partitionBits n+    | not (shouldParallelize parThreshold n) = 1+    | otherwise =+        max+            1+            ( min+                maxPartitionBits+                (max (ceilLog2 (n `div` rowsPerPartition)) (ceilLog2 (4 * capabilities)))+            )++maxPartitionBits :: Int+maxPartitionBits = 8++rowsPerPartition :: Int+rowsPerPartition = 65536++ceilLog2 :: Int -> Int+ceilLog2 x+    | x <= 1 = 0+    | otherwise = finiteBitSize x - countLeadingZeros (x - 1)++-- | Entries in each chunk's cache of recently emitted keys.+cacheSize :: Int+cacheSize = 4096++-- | Partition mark of a row dropped as a cached duplicate.+droppedRow :: Word16+droppedRow = maxBound++-- | Rows scattered by partition: key words, optional second words, starts.+data Scattered = Scattered !(VU.Vector Int) !(VU.Vector Int) !(VU.Vector Int)++-- | Hash every row to a partition, then write each row's words into it.+scatterRows ::+    Int ->+    Int ->+    -- | keep the hash between passes+    Bool ->+    -- | write a second word per row+    Bool ->+    -- | row hash+    (Int -> Int) ->+    -- | do two rows hold the same key?+    (Int -> Int -> Bool) ->+    -- | first word, from row and hash+    (Int -> Int -> Int) ->+    -- | second word, from row+    (Int -> Int) ->+    IO Scattered+scatterRows n pb keepHash twoWords rowHash rowEq rowW0 rowW1 = do+    let !np = 1 `unsafeShiftL` pb+        chunks+            | shouldParallelize parThreshold n = splitChunkRange capabilities n+            | otherwise = [(0, n)]+        !nc = length chunks+    hashes <- if keepHash then VUM.unsafeNew n else VUM.unsafeNew 0+    parts <- VUM.unsafeNew n :: IO (VUM.IOVector Word16)+    hist <- VUM.replicate (nc * np) (0 :: Int)+    forkJoin_+        [ do+            cacheHash <- VUM.unsafeNew cacheSize+            cacheRow <- VUM.replicate cacheSize (-1 :: Int)+            let !base = ci * np+                go !i+                    | i >= hi = pure ()+                    | otherwise = do+                        let !h = rowHash i+                            !c = h .&. (cacheSize - 1)+                        when keepHash (VUM.unsafeWrite hashes i h)+                        r <- VUM.unsafeRead cacheRow c+                        dup <-+                            if r < 0+                                then pure False+                                else do+                                    h' <- VUM.unsafeRead cacheHash c+                                    pure (h' == h && rowEq r i)+                        if dup+                            then VUM.unsafeWrite parts i droppedRow+                            else do+                                VUM.unsafeWrite cacheHash c h+                                VUM.unsafeWrite cacheRow c i+                                let !p = partOf pb h+                                VUM.unsafeWrite parts i (fromIntegral p)+                                VUM.unsafeModify hist (+ 1) (base + p)+                        go (i + 1)+            go lo+        | (ci, (lo, hi)) <- zip [0 ..] chunks+        ]+    -- Exclusive prefix sum, partition-major, turns counts into write cursors.+    starts <- VUM.unsafeNew (np + 1)+    let prefix !p !run+            | p >= np = VUM.unsafeWrite starts np run+            | otherwise = do+                VUM.unsafeWrite starts p run+                let walk !ci !r+                        | ci >= nc = pure r+                        | otherwise = do+                            c <- VUM.unsafeRead hist (ci * np + p)+                            VUM.unsafeWrite hist (ci * np + p) r+                            walk (ci + 1) (r + c)+                run' <- walk 0 run+                prefix (p + 1) run'+    prefix 0 0+    kept <- VUM.unsafeRead starts np+    w0 <- VUM.unsafeNew kept+    w1 <- if twoWords then VUM.unsafeNew kept else VUM.unsafeNew 0+    forkJoin_+        [ let !base = ci * np+              go !i+                | i >= hi = pure ()+                | otherwise = do+                    p <- VUM.unsafeRead parts i+                    when (p /= droppedRow) $ do+                        h <- if keepHash then VUM.unsafeRead hashes i else pure 0+                        let !slot = base + fromIntegral p+                        pos <- VUM.unsafeRead hist slot+                        VUM.unsafeWrite hist slot (pos + 1)+                        VUM.unsafeWrite w0 pos (rowW0 i h)+                        when twoWords (VUM.unsafeWrite w1 pos (rowW1 i))+                    go (i + 1)+           in go lo+        | (ci, (lo, hi)) <- zip [0 ..] chunks+        ]+    Scattered+        <$> VU.unsafeFreeze w0+        <*> VU.unsafeFreeze w1+        <*> VU.unsafeFreeze starts+{-# INLINE scatterRows #-}++-- | Run @body@ over partitions on a worker pool pulling from a shared counter.+forPartitions :: Int -> IO s -> (s -> Int -> IO ()) -> IO [s]+forPartitions np newState body = do+    next <- newIORef 0+    let !width = if np > 1 && capabilities > 1 then min capabilities np else 1+    states <- replicateM width newState+    let worker st = do+            p <- atomicModifyIORef' next (\j -> (j + 1, j))+            when (p < np) (body st p >> worker st)+    forkJoin_ (map worker states)+    pure states+{-# INLINE forPartitions #-}++-- | Per-group counts from a body that bumps a counter once per new pair.+groupCounts :: Int -> Int -> (Counts -> Int -> IO ()) -> IO (VU.Vector Int)+groupCounts nGroups np body+    | nGroups * capabilities <= 1 `unsafeShiftL` 22 = do+        arrs <- forPartitions np (VUM.replicate nGroups 0) $ \arr p ->+            body (Counts False arr) p+        total <- VUM.replicate nGroups 0+        forM_ arrs $ \arr ->+            forM_ [0 .. nGroups - 1] $ \g -> do+                c <- VUM.unsafeRead arr g+                VUM.unsafeModify total (+ c) g+        VU.unsafeFreeze total+    | otherwise = do+        shared <- VUM.replicate nGroups 0+        _ <- forPartitions np (pure ()) $ \() p -> body (Counts True shared) p+        VU.unsafeFreeze shared+{-# INLINE groupCounts #-}++-- | A group-count array and whether other workers share it.+data Counts = Counts !Bool !(VUM.IOVector Int)++bumpGroup :: Counts -> Int -> IO ()+bumpGroup (Counts shared arr) g+    | shared = atomicIncr arr g+    | otherwise = VUM.unsafeModify arr (+ 1) g+{-# INLINE bumpGroup #-}++atomicIncr :: VUM.IOVector Int -> Int -> IO ()+atomicIncr (MV_Int (VPM.MVector (I# off) _ (MutableByteArray mba))) (I# i) =+    IO+        (\s -> case fetchAddIntArray# mba (off +# i) 1# s of (# s', _ #) -> (# s', () #))+{-# INLINE atomicIncr #-}++modifyStrict :: IORef Int -> (Int -> Int) -> IO ()+modifyStrict ref f = readIORef ref >>= \x -> writeIORef ref $! f x+{-# INLINE modifyStrict #-}++-- | Starting table width for a partition of @rows@ entries.+initialBits :: Int -> Int+initialBits rows = max 4 (min 12 (ceilLog2 (2 * rows)))++-- | Distinct words in @w0[s, e)@, with 'minBound' tracked by a flag.+dedupWords :: Int -> VU.Vector Int -> Int -> Int -> IO Int+dedupWords pb w0 s e = do+    let !bits0 = initialBits (e - s)+    t0 <- VUM.replicate (1 `unsafeShiftL` bits0) emptyWord+    let go !t !bits !cnt !sawEmpty !i+            | i >= e = pure (if sawEmpty then cnt + 1 else cnt)+            | otherwise = do+                let !k = VU.unsafeIndex w0 i+                if k == emptyWord+                    then go t bits cnt True (i + 1)+                    else do+                        let !mask = (1 `unsafeShiftL` bits) - 1+                            probe !sl = do+                                x <- VUM.unsafeRead t sl+                                if x == emptyWord+                                    then VUM.unsafeWrite t sl k >> pure True+                                    else+                                        if x == k+                                            then pure False+                                            else probe ((sl + 1) .&. mask)+                        new <- probe (slotOf pb bits (fmix k))+                        if not new+                            then go t bits cnt sawEmpty (i + 1)+                            else+                                if 2 * (cnt + 1) > mask + 1+                                    then do+                                        t' <- growWords pb t bits+                                        go t' (bits + 1) (cnt + 1) sawEmpty (i + 1)+                                    else go t bits (cnt + 1) sawEmpty (i + 1)+    go t0 bits0 0 False s++emptyWord :: Int+emptyWord = minBound++growWords :: Int -> VUM.IOVector Int -> Int -> IO (VUM.IOVector Int)+growWords pb t bits = do+    let !bits' = bits + 1+        !mask' = (1 `unsafeShiftL` bits') - 1+    t' <- VUM.replicate (mask' + 1) emptyWord+    forM_ [0 .. VUM.length t - 1] $ \j -> do+        k <- VUM.unsafeRead t j+        when (k /= emptyWord) $ do+            let place !sl = do+                    x <- VUM.unsafeRead t' sl+                    if x == emptyWord+                        then VUM.unsafeWrite t' sl k+                        else place ((sl + 1) .&. mask')+            place (slotOf pb bits' (fmix k))+    pure t'++-- | Distinct @(w0, w1)@ entries in @[s, e)@ under @eq@, slotted by @hashOf@.+dedupPairs ::+    Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    Int ->+    (Int -> Int -> Int) ->+    (Int -> Int -> Int -> Int -> Bool) ->+    (Int -> Int -> IO ()) ->+    IO Int+dedupPairs pb w0 w1 s e hashOf eq onNew = do+    let !bits0 = initialBits (e - s)+    a0 <- VUM.unsafeNew (1 `unsafeShiftL` bits0)+    a1 <- VUM.replicate (1 `unsafeShiftL` bits0) (-1)+    let go !t0 !t1 !bits !cnt !i+            | i >= e = pure cnt+            | otherwise = do+                let !k0 = VU.unsafeIndex w0 i+                    !k1 = VU.unsafeIndex w1 i+                    !mask = (1 `unsafeShiftL` bits) - 1+                    probe !sl = do+                        y <- VUM.unsafeRead t1 sl+                        if y < 0+                            then do+                                VUM.unsafeWrite t0 sl k0+                                VUM.unsafeWrite t1 sl k1+                                pure True+                            else do+                                x <- VUM.unsafeRead t0 sl+                                if eq x y k0 k1+                                    then pure False+                                    else probe ((sl + 1) .&. mask)+                new <- probe (slotOf pb bits (hashOf k0 k1))+                if not new+                    then go t0 t1 bits cnt (i + 1)+                    else do+                        onNew k0 k1+                        if 2 * (cnt + 1) > mask + 1+                            then do+                                (t0', t1') <- growPairs pb hashOf t0 t1 bits+                                go t0' t1' (bits + 1) (cnt + 1) (i + 1)+                            else go t0 t1 bits (cnt + 1) (i + 1)+    go a0 a1 bits0 0 s+{-# INLINE dedupPairs #-}++growPairs ::+    Int ->+    (Int -> Int -> Int) ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    Int ->+    IO (VUM.IOVector Int, VUM.IOVector Int)+growPairs pb hashOf t0 t1 bits = do+    let !bits' = bits + 1+        !mask' = (1 `unsafeShiftL` bits') - 1+    n0 <- VUM.unsafeNew (mask' + 1)+    n1 <- VUM.replicate (mask' + 1) (-1)+    forM_ [0 .. VUM.length t1 - 1] $ \j -> do+        y <- VUM.unsafeRead t1 j+        when (y >= 0) $ do+            x <- VUM.unsafeRead t0 j+            let place !sl = do+                    y' <- VUM.unsafeRead n1 sl+                    if y' < 0+                        then VUM.unsafeWrite n0 sl x >> VUM.unsafeWrite n1 sl y+                        else place ((sl + 1) .&. mask')+            place (slotOf pb bits' (hashOf x y))+    pure (n0, n1)+{-# INLINE growPairs #-}
+ src-internal/DataFrame/Internal/Aggregation/Kernel/Fused.hs view
@@ -0,0 +1,592 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | Fused multi-reduction kernels: several reductions in one pass.+module DataFrame.Internal.Aggregation.Kernel.Fused (+    FusedAgg,+    mkFusedAgg,+    runFusedAggs,+    GatherAgg,+    mkGatherAgg,+    runGatherAggs,+) where++import Control.Exception (evaluate)+import Control.Monad (replicateM)+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++-- | One fused reduction of a multi-reduction streaming pass.+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 skips the fused streaming pass.+fusedSeqFloatGroups :: Int+fusedSeqFloatGroups = 65536++-- | Build the fused-pass reduction for a pair, or 'Nothing' if it can't stream.+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)+                    _ -> Nothing+        _ -> Nothing++-- | Execute all fused reductions in one pass over @rowToGroup@ and values.+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: step functions and merge.+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)++-- | Mean aggs keep their (sum, count) pair interleaved in one array.+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)++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)++-- | One fused reduction of a multi-reduction gather pass.+data GatherAgg+    = forall s.+        GatherAgg+        (IO s)+        (s -> Int -> Int -> IO ())+        (s -> IO Column)++-- | Build the fused gather reduction for a pair, or 'Nothing' if it can't fuse.+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)++-- | Groups each fused-gather block hands to every reduction before moving on.+gatherBlock :: Int+gatherBlock = 32++-- | Execute all fused gather reductions in one traversal over group ranges.+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 #-}++-- Gather folds replicate their unfused kernels' per-group recurrence exactly.++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,350 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | The fused two-column moment kernel: six sufficient statistics per group.+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: @(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 filling the count and sums.+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'.+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 counterpart of 'momentScatterPar' over typed value columns.+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 for the 'testEquality' dispatch.+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+-- SPECIALIZE keeps per-element @realToFrac@ from going through 'Rational'.+{-# 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,599 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | The group-range scatter-accumulate reduction kernel over grouped layout.+module DataFrame.Internal.Aggregation.Kernel.Scatter (+    scatterReduce,+    scatterReducePar,+    maxMinusMinScatterPar,++    -- * Group-range helpers+    -- $shared+    groupRangeBounds,+    rtgFromVis,+    overGroupsAcc,+    groupSlices,+    streamGroupCap,+) where++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 the Moments kernel, which partitions groups alike.++-- | 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)+{-# 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 #-}++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 for the 'testEquality' dispatch above.+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+-- SPECIALIZE lets @realToFrac@ resolve at a concrete type (~4x faster).+{-# 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))+{-# INLINEABLE reduceParTyped #-}++-- | Fold each group's rows into a register accumulator, then hand it on.+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) is carried in registers.+        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 #-}++-- | Fused @max a - min b@ over the group-range layout in one traversal.+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 for the 'testEquality' dispatch.+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 live in registers; one traversal reads both 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++-- | Group-count cap for the fused streaming scatter kernels.+streamGroupCap :: Int+streamGroupCap = 1048576++-- | Near-equal contiguous slices of the group domain for parallel merges.+groupSlices :: Int -> [(Int, Int)]+groupSlices = chunksFor 4096
+ src-internal/DataFrame/Internal/Aggregation/Plan.hs view
@@ -0,0 +1,203 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | The aggregation fast-path planner: recognises supported aggregate shapes.+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.Kernel.Distinct (distinctSupported)+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 (..),+    SomeExpr (..),+ )+import Type.Reflection (Typeable, typeRep)++-- | The plan 'planAgg' produces for a recognised output expression.+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+    | -- | Exact distinct count over one named column.+      PlanCountDistinct T.Text++-- | Plan a named output expression, or 'Nothing' if the shape is unsupported.+planAgg :: GroupedDataFrame -> SomeExpr -> Maybe AggPlan+planAgg gdf (SomeExpr (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)+        "countDistinct" -> do+            outputType @Int+            c <- getColumn name (fullDataframe gdf)+            if distinctSupported c then Just (PlanCountDistinct name) else Nothing+        _ -> 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)+        "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 aggregate group: six sufficient-statistic columns.+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: linear in one base column, or a product.+data Term+    = Lin T.Text+    | Prod T.Text T.Text+    deriving (Eq, Ord, Show)++-- | Recognise the moment shape across a whole @aggregate@ list.+planMoments :: GroupedDataFrame -> [(T.Text, SomeExpr)] -> 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 SomeExpr -> (T.Text, SomeExpr) -> Maybe (T.Text, Role)+classify exprs (name, SomeExpr expr) = case expr of+    Agg (MergeAgg "count" _ _ _ _) _ -> Just (name, RoleN)+    Agg (FoldAgg "sum" _ _) arg -> (\t -> (name, termRole t)) <$> resolveTerm exprs (SomeExpr 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'.+resolveTerm :: M.Map T.Text SomeExpr -> SomeExpr -> Maybe Term+resolveTerm exprs = go (8 :: Int)+  where+    go 0 _ = Nothing+    go fuel (SomeExpr 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) (SomeExpr inner)+        Binary op l r+            | binaryName op == "mult" && binaryCommutative op -> do+                Lin a <- go (fuel - 1) (SomeExpr l)+                Lin b <- go (fuel - 1) (SomeExpr 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++-- | Find the unordered pair of base columns that the linear sums name.+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,33 @@+{-# LANGUAGE ScopedTypeVariables #-}++-- | Vocabulary shared by the aggregation planner and kernels.+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.+data Reduction+    = RSum+    | RCount+    | RMin+    | RMax+    | RMean+    | RStd+    | RVar+    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,91 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE MagicHash #-}++-- | A cheap FxHash-style hash used to bucket rows without @hashable@.+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, the initial accumulator.+fnvOffset :: Int+fnvOffset = fromIntegral (0xcbf29ce484222325 :: Word64)++-- | FNV-1a 64-bit prime.+fnvPrime :: Int+fnvPrime = 0x00000100000001b3++-- | Sentinel mixed in for a null slot so @Nothing@ differs from @Just 0@.+nullSalt :: Int+nullSalt = fromIntegral (0x9E3779B97F4A7C15 :: Word64)++-- | Mix an 'Int' into the accumulator with an FxHash-style step.+mixInt :: Int -> Int -> Int+mixInt acc x = (rotateL acc 13 `xor` x) * fnvPrime+{-# INLINE mixInt #-}++-- | Mix a 'Double' into the accumulator.+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' into the accumulator over its UTF-8 bytes, eight at a time.+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 of a text array, eight bytes at a time.+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 hash for arbitrary 'Show'-able values; slower.+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,95 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-}++-- | Stable rank of group representatives by ascending unsigned hash.+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, ordering like signed 'Int'.+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,280 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-}++-- | Parallel stable sort of row indices by unsigned per-row hash.+module DataFrame.Internal.Algorithms.Sort.Radix.Parallel (+    parSortByHash,+    parSortThreshold,+) where++import Control.Concurrent (getNumCapabilities)+import Control.Monad (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 sequential radix path is used instead.+parSortThreshold :: Int+parSortThreshold = 500000++-- | Top-bits partition index of a hash, in 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.+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 #-}++-- | Stable LSD radix sort of @[0, n)@ by ascending 'sortKey' of their hash.+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 buffers.+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, ping-ponging between buffers.+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++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 counting sort.+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 and chunks; histograms become cursors.+    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 its cursor.+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 into the output slots.+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,16 @@+-- | Umbrella re-export of the column implementation, split by concern.+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,297 @@+{-# 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 companions.+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 with optional bit-packed null bitmap.+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 - split this into a separate intermediate representation.++    -- Bit-packed Text: shared UTF-8 buffer, offsets, optional bitmap; CSV only.+    PackedText :: Maybe Bitmap -> {-# UNPACK #-} !PackedTextData -> Column+    -- A join's same-named non-key column pair; both sides keep native form.+    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 nulls.+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.+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 to dataframe columns, used as an I/O intermediate.+data MutableColumn where+    MBoxedColumn :: (Columnable a) => VBM.IOVector a -> MutableColumn+    MUnboxedColumn :: (Columnable a, VU.Unbox a) => VUM.IOVector a -> MutableColumn++-- | A type-erased 'Column' with a phantom element type, for expressions.+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@; identity otherwise.+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++-- | Selects which Column type to build from a vector for each 'Rep'.+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' in O(1), keeping native representations.+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).+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,179 @@+{-# 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 qualified Data.List as L+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.+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.+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 validity vector (1 = valid, 0 = null).+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 L.foldl' setBitIf (0 :: Word8) [0 .. 7]++-- | Build a bitmap from a list of null-row indices.+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.+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 #-}++-- | Whether any of the first @n@ rows is null.+bitmapHasNulls :: Int -> Bitmap -> Bool+bitmapHasNulls n bm = popCountUpTo n bm < n+{-# INLINE bitmapHasNulls #-}++-- | 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.+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,296 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE UndecidableInstances #-}++-- | Mutable, growable column builders for high-throughput ingest.+module DataFrame.Internal.Column.Builder (+    ColumnBuilder (..),+    NumBuilder,+    IntBuilder,+    DoubleBuilder,+    TextBuilder,+    TextChunk (..),+    newIntBuilder,+    newDoubleBuilder,+    newNumBuilder,+    newTextBuilder,+    appendInt,+    appendDouble,+    appendNum,+    appendText,+    appendTextSlice,+    appendTextSliceFromPtr,+    freezeTextChunk,+    concatColumns,+    concatColumnsIO,+    mergeTextChunks,+    mergeTextChunksIO,+) 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,+    concatColumnsIO,+    mergeTextChunks,+    mergeTextChunksIO,+ )+import Foreign.Ptr (Ptr)++-- | Operations shared by all column builders; 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, with a null-slot sentinel.+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 using one packed, growable byte array.+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.+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,422 @@+{-# 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 vectors or lists, plus typed extraction.+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.Int (Int16, Int32, Int64, Int8)+import Data.Kind (Type)+import Data.Type.Equality (TestEquality (..))+import Data.Word (Word16, Word32, Word64, 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, picking the best representation.+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.+fromUnboxedVector ::+    forall a. (Columnable a, VU.Unbox a) => VU.Vector a -> Column+fromUnboxedVector = UnboxedColumn Nothing++-- | O(n) Convert a list to a column, picking the best representation.+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.+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) Convert a column to a list; throws if the type is wrong.+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@.+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+                | hasNulls (VB.length inner) bm -> Left (nullsMismatch @c)+                | otherwise -> 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+                | hasNulls (VU.length inner) bm -> Left (nullsMismatch @c)+                | otherwise -> 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 #-}++hasNulls :: Int -> Maybe Bitmap -> Bool+hasNulls n = maybe False (bitmapHasNulls n)++-- | The column has nulls, so it must be read as @Maybe c@.+nullsMismatch :: forall c. (Columnable c) => DataFrameException+nullsMismatch =+    TypeMismatchException+        ( MkTypeErrorContext+            { userType = Right (typeRep @c)+            , expectedType = Right (typeRep @(Maybe c))+            , callingFunctionName = Just "toVector"+            , errorColumnName = Nothing+            }+        )++-- Some common types we will use for numerical computing.++-- | Convert a column to an unboxed 'Double' vector, coercing numeric types.+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.+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.+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 widenToInt f of+                    Just v -> Right v+                    Nothing ->+                        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+                            }+                        )++-- | Widen an integral vector to 'Int', or 'Nothing' if not integral.+widenToInt ::+    forall a. (Columnable a, VU.Unbox a) => VU.Vector a -> Maybe (VU.Vector Int)+widenToInt v+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = Just v+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int64) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int32) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int16) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int8) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Word) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Word64) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Word32) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Word16) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Word8) =+        Just (VU.map fromIntegral v)+    | otherwise = case sIntegral @a of+        STrue -> Just (VU.map fromIntegral v)+        SFalse -> 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 #-}++-- Freezes the data vector; builds the bitmap only if the column 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 -> case testEquality (typeRep @a) (typeRep @T.Text) of+            Just Refl ->+                VB.imap+                    (\i x -> if bitmapTestBit bitmap i then x else "null")+                    col'+            Nothing ->+                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,164 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | Dictionary-encode a text group key to dense first-appearance Int codes.+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.+dictMaxCardinality :: Int+dictMaxCardinality = 1048576++-- | Dictionary-encode a text-like column to dense codes plus cardinality.+dictEncodeColumn :: Column -> Maybe (VU.Vector Int, Int)+dictEncodeColumn = dictEncodeColumnUpTo dictMaxCardinality++-- | Like 'dictEncodeColumn' but bails to 'Nothing' past @maxCard@ values.+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 by hashing raw UTF-8 bytes.+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'.+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 over an open-addressing hash table.+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+                | next > maxCard = pure Nothing+                | i >= n = pure (Just next)+                | 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,183 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | Concatenation of per-chunk 'Column's, merging text at the byte level.+module DataFrame.Internal.Column.Merge (+    TextChunk (..),+    concatColumns,+    concatColumnsIO,+    mergeTextChunks,+    mergeTextChunksIO,+    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_, when)+import Control.Monad.ST (runST, stToIO)+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.Control.Concurrent (+    forkJoin_,+    parThreshold,+    shouldParallelize,+ )+import DataFrame.Internal.Data.PackedText (+    PackedOffsets (..),+    mkPackedContiguous,+    offAt,+    offCount,+    packedRowOffsets,+ )+import Type.Reflection (typeRep)++-- | A frozen text-builder chunk: UTF-8 bytes, row offsets and validity.+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.+packedFromTextChunk :: TextChunk -> Column+packedFromTextChunk (TextChunk arr _used offs bm) =+    PackedText bm (mkPackedContiguous arr offs)++-- | Merge text chunks into one packed-text 'Column' with one copy per chunk.+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))++-- | 'mergeTextChunks' with copies and offset rebases spread over workers.+mergeTextChunksIO :: [TextChunk] -> IO Column+mergeTextChunksIO [] = error "DataFrame.Internal.ColumnMerge.mergeTextChunksIO: empty list"+mergeTextChunksIO [c] = pure (packedFromTextChunk c)+mergeTextChunksIO cs+    | not (shouldParallelize parThreshold totalRows) = pure (mergeTextChunks cs)+    | otherwise = do+        arr <- stToIO (A.new (max 1 totalBytes))+        offs <- VUM.unsafeNew (totalRows + 1)+        VUM.unsafeWrite offs 0 0+        let byteBases = scanl (+) 0 (map tcUsed cs)+            rowBases = scanl (+) 0 (map tcRows cs)+            splice byteBase rowBase c = do+                let n = tcRows c+                    co = tcOffsets c+                    go !i = when (i <= n) $ do+                        VUM.unsafeWrite offs (rowBase + i) (byteBase + VU.unsafeIndex co i)+                        go (i + 1)+                stToIO (A.copyI (tcUsed c) arr byteBase (tcBytes c) 0)+                go 1+        forkJoin_ (zipWith3 splice byteBases rowBases cs)+        farr <- stToIO (A.unsafeFreeze arr)+        foffs <- VU.unsafeFreeze offs+        let !bm = concatValidity [Validity (tcBitmap c) (tcRows c) | c <- cs]+        pure (PackedText bm (mkPackedContiguous farr foffs))+  where+    totalBytes = sum (map tcUsed cs)+    totalRows = sum (map tcRows cs)++-- | 'concatColumns', merging packed text in parallel.+concatColumnsIO :: [Column] -> IO Column+concatColumnsIO cols+    | Just chunks <- mapM packedChunk cols = mergeTextChunksIO chunks+    | otherwise = pure (concatColumns cols)++-- | Merge per-chunk columns into one column.+concatColumns :: [Column] -> Column+concatColumns [] = error "DataFrame.Internal.Column.Builder.concatColumns: empty list"+concatColumns [c] = c+-- Normalize the whole list: a packed or merged chunk demotes all to boxed.+concatColumns cols+    | Just chunks <- mapM packedChunk cols = mergeTextChunks chunks+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"++-- | A packed-text column without a selection, viewed as a builder chunk.+packedChunk :: Column -> Maybe TextChunk+packedChunk (PackedText bm p) = case packedRowOffsets p of+    Just (arr, offs) ->+        let n = offCount offs - 1+            offsV = case offs of+                Offs64 v -> v+                Offs32 v -> VU.map fromIntegral v+         in Just (TextChunk arr (offAt offs n) offsV bm)+    Nothing -> Nothing+packedChunk _ = Nothing
+ src-internal/DataFrame/Internal/Column/Operations.hs view
@@ -0,0 +1,1529 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE BangPatterns #-}+{-# 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 TypeOperators #-}+{-# LANGUAGE UndecidableInstances #-}++-- | Bulk operations over columns: map, fold, slice, gather, zip, append.+module DataFrame.Internal.Column.Operations where++import qualified Data.Text as T+import qualified Data.Text.Array as A+import Data.Text.Internal (Text (Text))+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_, join, 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 (+    capabilities,+    forkJoin,+    forkJoin_,+    parThreshold,+    parallelChunks,+    parallelChunks_,+    shouldParallelize,+    splitChunkRange,+ )+import DataFrame.Internal.Data.PackedText (+    PackedSel (..),+    PackedTextData (..),+    mkPackedContiguous,+    offCount,+    packedGather,+    packedIndexText,+    packedLength,+ )+import System.IO.Unsafe (unsafePerformIO)+import Type.Reflection (+    TypeRep,+    Typeable,+    eqTypeRep,+    typeOf,+    typeRep,+    withTypeable,+    pattern App,+    type (:~:) (Refl),+    type (:~~:) (HRefl),+ )++-- | Force evaluation of all elements in a column.+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++-- | Surface the deferred both-sides-null error of a 'MergedColumn'.+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 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+    PackedText Nothing p+        | Just Refl <- testEquality (typeRep @b) (typeRep @T.Text) ->+            Right (mapPacked p)+    c@(PackedText _ _) -> mapColumn f (materializePacked c)+    c@(MergedColumn _ _) -> mapColumn f (materializeMerged c)+  where+    -- Rows decode one at a time from packed bytes; text results are re-packed.+    mapPacked :: (b ~ T.Text) => PackedTextData -> Column+    mapPacked p =+        let !n = packedLength p+         in case testEquality (typeRep @c) (typeRep @T.Text) of+                Just Refl -> packedMapText n (f . packedIndexText p)+                Nothing -> case sUnbox @c of+                    STrue -> UnboxedColumn Nothing (parGenerateUnboxed n (f . packedIndexText p))+                    SFalse -> fromVector @c (VB.generate n (f . packedIndexText p))+    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 ->+                let !n = VB.length col+                 in Right $ case sUnbox @c of+                        STrue -> UnboxedColumn bm $ case nullsIn n bm of+                            Nothing -> parGenerateUnboxed n (f . VB.unsafeIndex col)+                            Just nulls -> generateValidUnboxed nulls n (f . VB.unsafeIndex col)+                        SFalse -> case bm of+                            Nothing -> fromVector @c (VB.map f col)+                            Just b' -> generateValidBoxed b' n (f . VB.unsafeIndex 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 ->+                let !n = VU.length col+                 in Right $ case sUnbox @c of+                        STrue -> UnboxedColumn bm $ case nullsIn n bm of+                            Nothing -> fromMaybe mapLoop compiledMap f col+                            Just nulls -> generateValidUnboxed nulls n (f . VU.unsafeIndex col)+                        SFalse -> case bm of+                            Nothing -> fromVector @c (VB.generate n (f . VU.unsafeIndex col))+                            Just b' -> generateValidBoxed b' n (f . VU.unsafeIndex col)+            Nothing -> throwTypeMismatch @a @b+{-# INLINEABLE mapColumn #-}++-- | Apply a function returning an unboxed result to an unboxed vector.+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 ->+            let !n = VB.length col+             in Right $ case sUnbox @c of+                    STrue -> UnboxedColumn bm $ case nullsIn n bm of+                        Nothing -> VU.generate n (\i -> f i (VB.unsafeIndex col i))+                        Just nulls -> generateValidUnboxed nulls n (\i -> f i (VB.unsafeIndex col i))+                    SFalse -> case bm of+                        Nothing -> fromVector @c (VB.imap f col)+                        Just b' -> generateValidBoxed b' n (\i -> f i (VB.unsafeIndex col i))+        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 ->+            let !n = VU.length col+             in Right $ case sUnbox @c of+                    STrue -> UnboxedColumn bm $ case nullsIn n bm of+                        Nothing -> VU.imap f col+                        Just nulls -> generateValidUnboxed nulls n (\i -> f i (VU.unsafeIndex col i))+                    SFalse -> case bm of+                        Nothing -> fromVector @c (VB.imap f (VG.convert col))+                        Just b' -> generateValidBoxed b' n (\i -> f i (VU.unsafeIndex col i))+        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 reads out-of-range as "", so check the slice bounds 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 #-}++-- (n) Selects the elements at a given set of indices, keeping order.++-- | Parallel unboxed 'VU.generate', one contiguous chunk per capability.+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 its inline variant.+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@ is @v ! (ix ! i)@.+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 for monomorphic wrappers.+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 #-}++-- | Run @k@ with the element type known, or 'Nothing' for other types.+withKnownElem ::+    forall a r. (Typeable a) => (forall x. (VU.Unbox x) => a :~: x -> r) -> Maybe r+withKnownElem k+    | Just Refl <- testEquality (typeRep @a) (typeRep @Double) =+        Just (k @Double Refl)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = Just (k @Int Refl)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Bool) = Just (k @Bool Refl)+    | otherwise = Nothing+{-# INLINE withKnownElem #-}++-- | 'Nothing' for element types without a compiled loop.+compiledMap ::+    forall a c.+    (Typeable a, Typeable c) => Maybe ((a -> c) -> VU.Vector a -> VU.Vector c)+compiledMap = join (withKnownElem @a (\Refl -> withKnownElem @c (\Refl -> mapLoop)))++compiledZip ::+    forall a c.+    (Typeable a, Typeable c) =>+    Maybe ((a -> a -> c) -> VU.Vector a -> VU.Vector a -> VU.Vector c)+compiledZip = join (withKnownElem @a (\Refl -> withKnownElem @c (\Refl -> zipLoop)))++compiledFoldl ::+    forall a b.+    (Typeable a, Typeable b) => Maybe ((b -> a -> b) -> b -> VU.Vector a -> b)+compiledFoldl = join (withKnownElem @a (\Refl -> withKnownElem @b (\Refl -> VU.foldl')))++compiledFoldl1 ::+    forall a. (Typeable a) => Maybe ((a -> a -> a) -> VU.Vector a -> a)+compiledFoldl1 = withKnownElem @a (\Refl -> VU.foldl1')++mapLoop :: (VU.Unbox a, VU.Unbox c) => (a -> c) -> VU.Vector a -> VU.Vector c+mapLoop f v = parGenerateUnboxedInline (VU.length v) (f . VU.unsafeIndex v)+{-# INLINE mapLoop #-}++-- | Element-wise over two vectors, to the shorter length.+zipLoop ::+    (VU.Unbox a, VU.Unbox b, VU.Unbox c) =>+    (a -> b -> c) -> VU.Vector a -> VU.Vector b -> VU.Vector c+zipLoop f l r =+    parGenerateUnboxedInline+        (min (VU.length l) (VU.length r))+        (\i -> f (VU.unsafeIndex l i) (VU.unsafeIndex r i))+{-# INLINE zipLoop #-}++-- | The bitmap, if it marks any of the first @n@ rows null.+nullsIn :: Int -> Maybe Bitmap -> Maybe Bitmap+nullsIn n (Just bm) | bitmapHasNulls n bm = Just bm+nullsIn _ _ = Nothing+{-# INLINE nullsIn #-}++validAt :: Maybe Bitmap -> Int -> Bool+validAt bm i = maybe True (`bitmapTestBit` i) bm+{-# INLINE validAt #-}++-- | @g i@ for the valid rows only; null rows hold zero.+generateValidUnboxed ::+    (VU.Unbox c) => Bitmap -> Int -> (Int -> c) -> VU.Vector c+generateValidUnboxed bm n g = unsafePerformIO $ do+    mv <- VUM.new n+    let fill !lo !hi =+            let go !i+                    | i >= hi = pure ()+                    | bitmapTestBit bm i = VUM.unsafeWrite mv i (g i) >> go (i + 1)+                    | otherwise = go (i + 1)+             in go lo+    if shouldParallelize parThreshold n+        then parallelChunks_ parThreshold n fill+        else fill 0 n+    VU.unsafeFreeze mv+{-# NOINLINE generateValidUnboxed #-}++-- | A boxed column of @g i@ for the valid rows only; null rows stay null.+generateValidBoxed ::+    forall c. (Columnable c) => Bitmap -> Int -> (Int -> c) -> Column+generateValidBoxed bm n g = case typeRep @c of+    App tMaybe _+        | Just HRefl <- eqTypeRep tMaybe (typeRep @Maybe) ->+            fromVector @c+                (VB.generate n (\i -> if bitmapTestBit bm i then g i else Nothing))+    _ ->+        BoxedColumn+            (Just bm)+            ( VB.generate+                n+                (\i -> if bitmapTestBit bm i then g i else errorWithoutStackTrace "null slot")+            )++-- | 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)@.+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 that fetches array slots eagerly.+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 parallel pass.+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, built 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 use closure-free kernels; others take the generic 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.+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 a nullable source bitmap by raw indices.+        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 handles -1 sentinels itself; pass raw indices.+            PackedText srcBm p -> PackedText (withBm srcBm) (packedGather indices p)+            BoxedColumn srcBm v+                -- An empty source means every index is a sentinel.+                | 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 monomorphic Int/Double kernels.+                    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.+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 #-}++-- | Positions of the 'True' entries, in order, computed in parallel.+maskIndices :: VU.Vector Bool -> VU.Vector Int+maskIndices mask+    | not (shouldParallelize parThreshold n) = VU.findIndices id mask+    | otherwise = unsafePerformIO $ do+        let chunks = splitChunkRange capabilities n+            countChunk (lo, hi) = pure $! countFrom lo hi 0+            countFrom !i !hi !c+                | i >= hi = c+                | otherwise = countFrom (i + 1) hi (c + fromEnum (VU.unsafeIndex mask i))+        counts <- forkJoin (map countChunk chunks)+        let starts = scanl (+) 0 counts+        out <- VUM.unsafeNew (last starts)+        let fill (lo, hi) start =+                let go !i !k+                        | i >= hi = pure ()+                        | VU.unsafeIndex mask i = VUM.unsafeWrite out k i >> go (i + 1) (k + 1)+                        | otherwise = go (i + 1) k+                 in go lo start+        forkJoin_ (zipWith fill chunks starts)+        VU.unsafeFreeze out+  where+    !n = VU.length mask++findIndices ::+    forall a.+    (Columnable a) =>+    (a -> Bool) ->+    Column ->+    Either DataFrameException (VU.Vector Int)+findIndices predicate = \case+    BoxedColumn bm (v :: VB.Vector b) -> run bm v+    UnboxedColumn bm (v :: VU.Vector b) -> run bm v+    c@(PackedText _ _) -> findIndices predicate (materializePacked c)+    c@(MergedColumn _ _) -> findIndices predicate (materializeMerged c)+  where+    -- Null rows never match.+    run ::+        forall b v.+        (Typeable b, VG.Vector v b, VG.Vector v Int) =>+        Maybe Bitmap ->+        v b ->+        Either DataFrameException (VU.Vector Int)+    run bm column = case testEquality (typeRep @a) (typeRep @b) of+        Just Refl -> Right $ case nullsIn (VG.length column) bm of+            Nothing -> VG.convert (VG.findIndices predicate column)+            Just nulls ->+                VU.filter+                    (\i -> bitmapTestBit nulls i && predicate (VG.unsafeIndex column i))+                    (VU.enumFromN 0 (VG.length column))+        Nothing ->+            Left $+                TypeMismatchException+                    MkTypeErrorContext+                        { userType = Right (typeRep @a)+                        , expectedType = Right (typeRep @b)+                        , callingFunctionName = Just "findIndices"+                        , errorColumnName = Nothing+                        }++-- | Fold (right) column with index, skipping null rows.+ifoldrColumn ::+    forall a b.+    (Columnable a, Columnable b) =>+    (Int -> a -> b -> b) -> b -> Column -> Either DataFrameException b+ifoldrColumn f acc = \case+    BoxedColumn bm column -> foldrWorker bm column+    UnboxedColumn bm column -> foldrWorker bm 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) =>+        Maybe Bitmap ->+        v c ->+        Either DataFrameException b+    foldrWorker bm vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl -> pure $ case nullsIn (VG.length vec) bm of+            Nothing -> VG.ifoldr f acc vec+            Just nulls -> VG.ifoldr (\i x r -> if bitmapTestBit nulls i then f i x r else r) acc vec+        Nothing ->+            Left $+                TypeMismatchException+                    MkTypeErrorContext+                        { userType = Right (typeRep @a)+                        , expectedType = Right (typeRep @c)+                        , callingFunctionName = Just "ifoldrColumn"+                        , errorColumnName = Nothing+                        }++-- | Left fold over a column, skipping null rows.+foldlColumn ::+    forall a b.+    (Columnable a, Columnable b) =>+    (b -> a -> b) -> b -> Column -> Either DataFrameException b+foldlColumn f acc = \case+    BoxedColumn bm column -> foldlWorker bm column+    UnboxedColumn bm (column :: VU.Vector c)+        | Just Refl <- testEquality (typeRep @a) (typeRep @c)+        , Nothing <- nullsIn (VU.length column) bm ->+            pure (fromMaybe VU.foldl' compiledFoldl f acc column)+        | otherwise -> foldlWorker bm 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) =>+        Maybe Bitmap ->+        v c ->+        Either DataFrameException b+    foldlWorker bm vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl -> pure $ case nullsIn (VG.length vec) bm of+            Nothing -> VG.foldl' f acc vec+            Just nulls -> VG.ifoldl' (\r i x -> if bitmapTestBit nulls i then f r x else r) acc vec+        Nothing ->+            Left $+                TypeMismatchException+                    MkTypeErrorContext+                        { userType = Right (typeRep @a)+                        , expectedType = Right (typeRep @c)+                        , callingFunctionName = Just "foldlColumn"+                        , errorColumnName = Nothing+                        }++-- | Left fold seeded with the first valid row, skipping null rows.+foldl1Column ::+    forall a.+    (Columnable a) =>+    (a -> a -> a) -> Column -> Either DataFrameException a+foldl1Column f = \case+    BoxedColumn bm column -> foldl1Worker bm column+    UnboxedColumn bm (column :: VU.Vector c)+        | Just Refl <- testEquality (typeRep @a) (typeRep @c)+        , Nothing <- nullsIn (VU.length column) bm ->+            pure (fromMaybe VU.foldl1' compiledFoldl1 f column)+        | otherwise -> foldl1Worker bm column+    c@(PackedText _ _) -> foldl1Column f (materializePacked c)+    c@(MergedColumn _ _) -> foldl1Column f (materializeMerged c)+  where+    foldl1Worker ::+        forall c v.+        (Typeable c, VG.Vector v c) =>+        Maybe Bitmap ->+        v c ->+        Either DataFrameException a+    foldl1Worker bm vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl -> case nullsIn (VG.length vec) bm of+            Nothing -> pure (VG.foldl1' f vec)+            Just nulls ->+                let step r i x+                        | not (bitmapTestBit nulls i) = r+                        | otherwise = Just $! maybe x (`f` x) r+                 in maybe+                        (Left (EmptyDataSetException "foldl1Column"))+                        Right+                        (VG.ifoldl' step Nothing 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 each group's first element as seed.+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 bm (vec :: VU.Vector d) -> uncurry UnboxedColumn <$> foldl1Worker bm vec+        BoxedColumn bm (vec :: VB.Vector d) -> uncurry BoxedColumn <$> foldl1Worker bm vec+        PackedText _ _ -> foldl1DirectGroups f (materializePacked col) valueIndices offsets+        MergedColumn _ _ -> foldl1DirectGroups f (materializeMerged col) valueIndices offsets+  where+    nGroups = VU.length offsets - 1+    foldl1Worker ::+        forall c v.+        (Typeable c, VG.Vector v c) =>+        Maybe Bitmap ->+        v c ->+        Either DataFrameException (Maybe Bitmap, v c)+    foldl1Worker bm vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl -> Right $ case nullsIn (VG.length vec) bm of+            Nothing -> (Nothing, VG.generate nGroups foldGroup)+            Just nulls ->+                -- A group of all-null rows has no seed, so it is null.+                let groups = VG.generate nGroups (foldValidGroup nulls) :: VB.Vector (Maybe c)+                    empty = VU.filter (isNothing . VB.unsafeIndex groups) (VU.enumFromN 0 nGroups)+                    placeholder k = VG.unsafeIndex vec (VU.unsafeIndex valueIndices (VU.unsafeIndex offsets k))+                 in ( Just (buildBitmapFromNulls' nGroups empty)+                    , VG.generate nGroups (\k -> fromMaybe (placeholder k) (VB.unsafeIndex groups k))+                    )+          where+            row i = VG.unsafeIndex vec (VU.unsafeIndex valueIndices i)+            foldGroup k =+                let !s = VU.unsafeIndex offsets k+                    !e = VU.unsafeIndex offsets (k + 1)+                 in go (s + 1) e (row s)+            go !i !e !acc+                | i >= e = acc+                | otherwise = go (i + 1) e $! f acc (row i)+            foldValidGroup nulls k = goValid (VU.unsafeIndex offsets k) Nothing+              where+                !e = VU.unsafeIndex offsets (k + 1)+                goValid !i acc+                    | i >= e = acc+                    | not (bitmapTestBit nulls (VU.unsafeIndex valueIndices i)) =+                        goValid (i + 1) acc+                    | otherwise = goValid (i + 1) (Just $! maybe (row i) (`f` row i) acc)+        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.+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 bm (vec :: VU.Vector d) -> foldLinearWorker bm vec+        BoxedColumn bm (vec :: VB.Vector d) -> foldLinearWorker bm vec+        PackedText _ _ ->+            foldLinearGroups f seed (materializePacked col) rowToGroup nGroups+        MergedColumn _ _ ->+            foldLinearGroups f seed (materializeMerged col) rowToGroup nGroups+  where+    -- Null rows are skipped; a group with no valid rows keeps the seed.+    foldLinearWorker ::+        forall c v.+        (Typeable c, VG.Vector v c) =>+        Maybe Bitmap ->+        v c ->+        Either DataFrameException Column+    foldLinearWorker bm vec = case testEquality (typeRep @b) (typeRep @c) of+        Just Refl ->+            let nulls = nullsIn (VG.length vec) bm+             in Right $+                    unsafePerformIO $+                        runWith+                            ( \readAt writeAt ->+                                VG.iforM_ vec $ \row x -> when (validAt nulls row) $ 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 #-}++-- | The first valid row.+headColumn :: forall a. (Columnable a) => Column -> Either DataFrameException a+headColumn = \case+    BoxedColumn bm col -> headWorker bm col+    UnboxedColumn bm col -> headWorker bm col+    c@(PackedText _ _) -> headColumn (materializePacked c)+    c@(MergedColumn _ _) -> headColumn (mergedHead c)+  where+    headWorker ::+        forall c v.+        (Typeable c, VG.Vector v c) =>+        Maybe Bitmap ->+        v c ->+        Either DataFrameException a+    headWorker bm vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl ->+            case VU.find (validAt bm) (VU.enumFromN 0 (VG.length vec)) of+                Nothing -> Left (EmptyDataSetException "headColumn")+                Just i -> pure (VG.unsafeIndex vec i)+        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 bmL column) (BoxedColumn bmR other) = withBitmap (bothValid bmL bmR) (zipVectors column other)+zipColumns (BoxedColumn bmL column) (UnboxedColumn bmR other) = withBitmap (bothValid bmL bmR) (zipVectors column other)+zipColumns (UnboxedColumn bmL column) (BoxedColumn bmR other) = withBitmap (bothValid bmL bmR) (zipVectors column other)+zipColumns (UnboxedColumn bmL column) (UnboxedColumn bmR other) = withBitmap (bothValid bmL bmR) (zipVectors column other)+{-# INLINE zipColumns #-}++zipVectors ::+    forall a b va vb.+    (Columnable a, Columnable b, VG.Vector va a, VG.Vector vb b) =>+    va a -> vb b -> Column+zipVectors column other =+    case sUnbox @a of+        STrue -> case sUnbox @b of+            STrue -> UnboxedColumn Nothing (VU.generate n pairAt)+            SFalse -> BoxedColumn Nothing (VB.generate n pairAt)+        SFalse -> BoxedColumn Nothing (VB.generate n pairAt)+  where+    !n = min (VG.length column) (VG.length other)+    pairAt i = (VG.unsafeIndex column i, VG.unsafeIndex other i)+{-# INLINE zipVectors #-}++-- | 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 left right = case (left, right) of+    (PackedText _ _, _) -> zipWithColumns f (materializePacked left) right+    (_, PackedText _ _) -> zipWithColumns f left (materializePacked right)+    (MergedColumn _ _, _) -> zipWithColumns f (materializeMerged left) right+    (_, MergedColumn _ _) -> zipWithColumns f left (materializeMerged right)+    ( UnboxedColumn Nothing (l :: VU.Vector d)+        , UnboxedColumn Nothing (r :: VU.Vector e)+        )+            | Just Refl <- testEquality (typeRep @a) (typeRep @d)+            , Just Refl <- testEquality (typeRep @b) (typeRep @e)+            , STrue <- sUnbox @c ->+                pure . UnboxedColumn Nothing $ case testEquality (typeRep @a) (typeRep @b) of+                    Just Refl -> fromMaybe zipLoop compiledZip f l r+                    Nothing -> zipLoop f l r+    (BoxedColumn bmL l, BoxedColumn bmR r) -> elements bmL l bmR r+    (BoxedColumn bmL l, UnboxedColumn bmR r) -> elements bmL l bmR r+    (UnboxedColumn bmL l, BoxedColumn bmR r) -> elements bmL l bmR r+    (UnboxedColumn bmL l, UnboxedColumn bmR r) -> elements bmL l bmR r+  where+    -- A row is null if either side is; f never sees a null row.+    elements ::+        forall d e vl vr.+        (Typeable d, Typeable e, VG.Vector vl d, VG.Vector vr e) =>+        Maybe Bitmap -> vl d -> Maybe Bitmap -> vr e -> Either DataFrameException Column+    elements bmL l bmR r+        | Just Refl <- testEquality (typeRep @a) (typeRep @d)+        , Just Refl <- testEquality (typeRep @b) (typeRep @e) =+            let !n = min (VG.length l) (VG.length r)+                g i = f (VG.unsafeIndex l i) (VG.unsafeIndex r i)+             in pure $ case (bothValid bmL bmR, sUnbox @c) of+                    (Nothing, STrue) -> UnboxedColumn Nothing (parGenerateUnboxed n g)+                    (Nothing, SFalse) -> fromVector @c (VB.generate n g)+                    (Just bm, STrue) -> UnboxedColumn (Just bm) $ case nullsIn n (Just bm) of+                        Nothing -> parGenerateUnboxed n g+                        Just nulls -> generateValidUnboxed nulls n g+                    (Just bm, SFalse) -> generateValidBoxed bm n g+        | otherwise = zipWithColumnsGeneral f left right++-- | Valid where both are valid; a missing bitmap means every row is valid.+bothValid :: Maybe Bitmap -> Maybe Bitmap -> Maybe Bitmap+bothValid (Just l) (Just r) = Just (andBitmaps l r)+bothValid l Nothing = l+bothValid Nothing r = r++withBitmap :: Maybe Bitmap -> Column -> Column+withBitmap bm (BoxedColumn _ v) = BoxedColumn bm v+withBitmap bm (UnboxedColumn _ v) = UnboxedColumn bm v+withBitmap _ c = c++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 #-}++-- writeColumn and freezeColumn' moved to DataFrame.IO.Internal.MutableColumn.++-- | Freeze a mutable column into an @Either Text a@ column.+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 (all-valid bitmap).+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++-- | Fill the end of a column, up to n, with null rows.+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++-- | Fill the beginning of a column, up to n, with null rows.+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; fails if the column types differ.+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+                        }+                    )++-- ike 'mappendColumns' but wraps values in 'Either' across types.++-- | O(n) Concatenate a list of same-type columns in a single allocation.+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++-- | Gather one in-bounds index vector through several columns in one pass.+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 and an 'Int32'-sized base.+      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, one gathered column per spec.+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 finalizer.+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 #-}++-- | A packed text column of @f i@ per row, built chunk by chunk in parallel.+packedMapText :: Int -> (Int -> T.Text) -> Column+packedMapText n f = unsafePerformIO $ do+    chunks <-+        if shouldParallelize parThreshold n+            then parallelChunks parThreshold n (\lo hi -> pure $! packRows lo hi)+            else pure [packRows 0 n]+    let totalBytes = sum [VU.unsafeIndex o (VU.length o - 1) | (_, o) <- chunks]+        totalRows = sum [VU.length o - 1 | (_, o) <- chunks]+    pure $! runST $ do+        arr <- A.new (max 1 totalBytes)+        offs <- VUM.unsafeNew (totalRows + 1)+        VUM.unsafeWrite offs 0 0+        let splice !byteBase !rowBase ((a, o) : rest) = do+                let rows = VU.length o - 1+                    used = VU.unsafeIndex o rows+                A.copyI used arr byteBase a 0+                forM_ [1 .. rows] $ \i ->+                    VUM.unsafeWrite offs (rowBase + i) (byteBase + VU.unsafeIndex o i)+                splice (byteBase + used) (rowBase + rows) rest+            splice _ _ [] = pure ()+        splice 0 0 chunks+        farr <- A.unsafeFreeze arr+        foffs <- VU.unsafeFreeze offs+        pure (PackedText Nothing (mkPackedContiguous farr foffs))+  where+    -- One chunk's rows as (bytes, offsets); offsets are chunk-relative.+    packRows lo hi =+        let texts = VB.generate (hi - lo) (\i -> f (lo + i))+            lens = VU.generate (hi - lo) (textBytes . VB.unsafeIndex texts)+            offs = VU.scanl' (+) 0 lens+            bytes = VU.unsafeIndex offs (hi - lo)+            arr = runST $ do+                m <- A.new (max 1 bytes)+                VB.iforM_ texts $ \i (Text src off len) ->+                    A.copyI len m (VU.unsafeIndex offs i) src off+                A.unsafeFreeze m+         in (arr, offs)+    textBytes (Text _ _ len) = len
+ src-internal/DataFrame/Internal/Column/Properties.hs view
@@ -0,0 +1,146 @@+{-# 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'.+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@, looking through 'Maybe'.+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)++-- | Debugging helper: the type stored in a column's outermost vector.+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,153 @@+{-# 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@, avoiding the @these@ package.+data These a b = This a | That b | These a b+    deriving (Eq, Ord, Show, Read, Functor, Foldable, Traversable)++-- | Column representations used to select the right 'toColumn' behaviour.+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 (a, b) = If (Unboxable a) (Unboxable b) 'False+    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.+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 by integral gives Double, for './'.+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,168 @@+{-# 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.+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.+    !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@.+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,90 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}++-- | A flat open-addressing hash table mapping row hash to a dense group id.+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 with a power-of-two capacity.+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.+nextPow2Above :: Int -> Int+nextPow2Above n = go 2+  where+    go !p+        | p > n = p+        | otherwise = go (p * 2)+{-# INLINE nextPow2Above #-}++-- | Allocate an empty table for up to @n@ groups at load factor under 0.5.+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@ by hash and return its group id, adding one if new.+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 bits of a Fibonacci multiply.+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,304 @@+{-# LANGUAGE BangPatterns #-}++-- | Packed-text payload and byte-slice primitives over one shared buffer.+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, 'Int32' when every value fits and 'Int' otherwise.+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 null.+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 and optional selection.+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 to 'Int32' when the row count fits.+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.+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+        -- Only the selector is rebuilt; bytes and offsets stay shared.+        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; out-of-range picks give -1.++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 #-}++-- elector generate over 'parallelChunks_', one chunk per capability.++-- | 'parGenSel' with an INLINE body so each kernel inlines its pick function.+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+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, in O(k).+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)@.+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 and contiguous offsets when the payload is unselected.+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@, similar to freeze.+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, agreeing with 'Text' equality.+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 of two slices (memcmp semantics).+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,84 @@+{-# LANGUAGE BangPatterns #-}++-- | UTF-8 validation and lenient slice decoding for shared byte buffers.+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 (is not a continuation byte).+isUtf8Boundary :: Word8 -> Bool+isUtf8Boundary w = w < 0x80 || w >= 0xC0+{-# INLINE isUtf8Boundary #-}++-- | Decode a byte slice exactly like @decodeUtf8Lenient@.+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 by row offsets.+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,383 @@+{-# 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,+    fromColumns,+    getColumn,+    unsafeGetColumn,+    null,+    toCsv,+    toCsv',+    toSeparated,+    getRowAsText,+    showElement,+    stripJust,+    (=:),+) 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.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 (..),+    Columnable,+    columnLength,+    columnToTextVec,+    columnTypeString,+    expandColumn,+    forceColumn,+    fromList,+    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 SomeExpr+    }++-- | Force evaluation of all columns in a DataFrame.+forceDataFrame :: DataFrame -> DataFrame+forceDataFrame df@(DataFrame cols idx dims _exprs) =+    V.foldl' (\() c -> forceColumn c) () cols `seq` idx `seq` dims `seq` df++-- | Row grouping of a dataframe; can only be used with 'aggregate'.+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] is the group index for row i, built once in 'groupBy'.+    , groupRepRows :: VU.Vector Int+    -- ^ First original row of each group, length nGroups. See laziness note.+    }++-- | The historical five-field view of 'GroupedDataFrame'.+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; indices are in-bounds by construction.+                (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 truncated text.+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.+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, replacing any column of the same name.+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 @(name, column)@ pairs using 'insertColumn'.+fromColumns :: [(T.Text, Column)] -> DataFrame+fromColumns = L.foldl' (\df (name, column) -> insertColumn name column df) empty++mkColumn :: (Columnable a) => T.Text -> [a] -> (T.Text, Column)+mkColumn name column = (name, fromList column)++infix 1 =:+(=:) :: (Columnable a) => T.Text -> [a] -> (T.Text, Column)+(=:) = mkColumn++-- | Safely retrieve a column by name, or 'Nothing' if it does not exist.+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 absent.+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).+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,121 @@+-- | A minimal Wadler/Leijen-style document combinator and width-aware renderer.+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 space when its group fits the width, else a newline.+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, in parens when it must break.+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,507 @@+{-# 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 so the simplifier can match them by 'cast'.+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 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 SomeExpr where+    SomeExpr :: (Columnable a) => Expr a -> SomeExpr++instance Show SomeExpr where+    show :: SomeExpr -> String+    show (SomeExpr expr) = show expr++toSomeExpr :: (Columnable a) => Expr a -> SomeExpr+toSomeExpr = SomeExpr++fromSomeExpr :: forall a. (Columnable a) => SomeExpr -> Maybe (Expr a)+fromSomeExpr (SomeExpr (expr :: Expr b)) = do+    Refl <- testEquality (typeRep @a) (typeRep @b)+    pure expr++type NamedExpr = (T.Text, SomeExpr)++toNamedExpr :: (Columnable a) => T.Text -> Expr a -> NamedExpr+toNamedExpr exprName expr = (exprName, SomeExpr 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.+    (Columnable a, Columnable b) =>+    Expr a -> Expr a -> Expr b -> Expr b+replaceExpr old new expr = case testEquality (typeRep @a) (typeRep @b) of+    Just Refl -> if eqExpr old expr then new else replace'+    Nothing -> replace'+  where+    replace' = case expr of+        (Col _) -> expr+        (CastWith{}) -> expr+        (CastExprWith t f e) -> CastExprWith t f (replaceExpr old new e)+        (Lit _) -> expr+        (If cond l r) ->+            If (replaceExpr old new cond) (replaceExpr old new l) (replaceExpr old new r)+        (Unary op value) -> Unary op (replaceExpr old new value)+        (Binary op l r) -> Binary op (replaceExpr old new l) (replaceExpr old new r)+        (Agg op inner) -> Agg op (replaceExpr old new inner)+        (Over keys inner) -> Over keys (replaceExpr old new inner)++-- | Substitute 'Col' references from a name-to-expression map in one pass.+substituteColumns ::+    forall a. (Columnable a) => M.Map T.Text SomeExpr -> 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 (SomeExpr (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 pseudo-code at the default width.+prettyPrint :: Expr a -> String+prettyPrint = prettyPrintWidth P.defaultWidth++-- | Render an expression as readable, width-aware pseudo-code.+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.nest indentWidth (P.hardline <> 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 5 (P.group (toDoc 0 c))+            <> P.nest indentWidth (P.hardline <> toDoc 0 t)+            <> P.hardline+            <> renderElse e+    renderElse other =+        P.text "else" <> P.nest indentWidth (P.hardline <> toDoc 0 other)+    indentWidth :: Int+    indentWidth = 4
+ 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,+    SomeExpr (SomeExpr),+    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, SomeExpr 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++caseWhen :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a+caseWhen = If++-- | Closes a 'caseWhen' chain with its fallback value.+orElse :: Expr a -> Expr a+orElse = id++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 for nullable and non-nullable operands.+(.+) ::+    ( 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++-- Three-valued comparisons: Nothing if either operand is Nothing.++-- | Nullable-aware equality that widens numeric operands to a 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 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 that widens numeric operands.+(.<=) ::+    ( 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 that widens numeric operands.+(.>=) ::+    ( 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, giving @Maybe Bool@ if either side 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, giving @Maybe Bool@ if either side 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,378 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE FunctionalDependencies #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE UndecidableSuperClasses #-}++-- | Nullable-aware operators working across @Maybe a@ and @a@ operands.+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'.+type family BaseType a where+    BaseType (Maybe a) = a+    BaseType a = a++-- | Arithmetic binary operations over nullable and non-nullable operands.+class+    ( Columnable a+    , Columnable b+    , Columnable c+    ) =>+    NullableArithOp a b c+        | a b -> c+    where+    -- | Lift an arithmetic function over inner values; 'Nothing' wins.+    nullArithOp ::+        (BaseType a -> BaseType a -> BaseType a) ->+        a ->+        b ->+        c++-- | Compute the result type of a nullable comparison.+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.+class+    ( Columnable a+    , Columnable b+    , Columnable e+    ) =>+    NullableCmpOp a b e+    where+    -- | Lift a comparison over inner values with three-valued logic.+    nullCmpOp ::+        (BaseType a -> BaseType a -> Bool) ->+        a ->+        b ->+        e++-- | Non-nullable by non-nullable: apply directly, no wrapping.+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)++-- ift a unary function over a column expression, propagating 'Nothing'.++-- | Compute the result type of a nullable unary lift.+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)++-- ift a binary function over two column expressions, propagating 'Nothing'.++-- | Compute the result type of a nullable binary lift.+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)++-- | Widen two numeric base types to their promoted common type.+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.+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.+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 nullables.+type WidenResult a b = NullLift2Result a b (Promote (BaseType a) (BaseType b))++-- | Like 'NumericWidenOp' but uses 'PromoteDiv' (integral gives 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.+type WidenResultDiv a b =+    NullLift2Result a b (PromoteDiv (BaseType a) (BaseType b))
+ src-internal/DataFrame/Internal/Expression/Reflect.hs view
@@ -0,0 +1,111 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE UndecidableInstances #-}++-- | Class evidence for the supported column types, recovered from 'Typeable'.+module DataFrame.Internal.Expression.Reflect (+    Dict (..),+    RealUnbox,+    numDict,+    ordDict,+    fracDict,+    floatingDict,+    realDict,+    integralDict,+    realUnboxDict,+) where++import Data.Foldable (asum)+import Data.Int (Int16, Int32, Int64, Int8)+import Data.Kind (Constraint, Type)+import qualified Data.Text as T+import Data.Type.Equality (testEquality, (:~:) (Refl))+import qualified Data.Vector.Unboxed as VU+import Data.Word (Word16, Word32, Word64, Word8)+import Type.Reflection (Typeable, typeRep)++-- | A captured class dictionary.+data Dict (c :: Constraint) where+    Dict :: (c) => Dict c++-- | 'Real' and 'VU.Unbox' together, for variance, median and similar.+class (Real a, VU.Unbox a) => RealUnbox a++instance (Real a, VU.Unbox a) => RealUnbox a++-- | Evidence for @c a@ when @a@ is the concrete type @b@.+via ::+    forall b (c :: Type -> Constraint) a.+    (Typeable a, Typeable b, c b) => Maybe (Dict (c a))+via = case testEquality (typeRep @a) (typeRep @b) of+    Just Refl -> Just Dict+    Nothing -> Nothing++-- | The fixed-width integral types.+ints ::+    forall (c :: Type -> Constraint) a.+    ( Typeable a+    , c Int+    , c Int8+    , c Int16+    , c Int32+    , c Int64+    , c Word+    , c Word8+    , c Word16+    , c Word32+    , c Word64+    ) =>+    [Maybe (Dict (c a))]+ints =+    [ via @Int @c @a+    , via @Int8 @c @a+    , via @Int16 @c @a+    , via @Int32 @c @a+    , via @Int64 @c @a+    , via @Word @c @a+    , via @Word8 @c @a+    , via @Word16 @c @a+    , via @Word32 @c @a+    , via @Word64 @c @a+    ]++floats ::+    forall (c :: Type -> Constraint) a.+    (Typeable a, c Double, c Float) => [Maybe (Dict (c a))]+floats = [via @Double @c @a, via @Float @c @a]++numDict :: forall a. (Typeable a) => Maybe (Dict (Num a))+numDict = asum (ints @Num @a ++ [via @Integer @Num @a] ++ floats @Num @a)++ordDict :: forall a. (Typeable a) => Maybe (Dict (Ord a))+ordDict =+    asum+        ( ints @Ord @a+            ++ [via @Integer @Ord @a]+            ++ floats @Ord @a+            ++ [via @Bool @Ord @a, via @Char @Ord @a, via @T.Text @Ord @a, via @String @Ord @a]+        )++fracDict :: forall a. (Typeable a) => Maybe (Dict (Fractional a))+fracDict = asum (floats @Fractional @a)++floatingDict :: forall a. (Typeable a) => Maybe (Dict (Floating a))+floatingDict = asum (floats @Floating @a)++realDict :: forall a. (Typeable a) => Maybe (Dict (Real a))+realDict = asum (ints @Real @a ++ [via @Integer @Real @a] ++ floats @Real @a)++integralDict :: forall a. (Typeable a) => Maybe (Dict (Integral a))+integralDict = asum (ints @Integral @a ++ [via @Integer @Integral @a])++-- | 'Integer' is deliberately absent: it is 'Real' but not 'VU.Unbox'.+realUnboxDict :: forall a. (Typeable a) => Maybe (Dict (RealUnbox a))+realUnboxDict = asum (ints @RealUnbox @a ++ floats @RealUnbox @a)
+ src-internal/DataFrame/Internal/Expression/Simplify.hs view
@@ -0,0 +1,397 @@+{-# 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 for @Expr Bool@ and @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).+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.+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 over a @Double@ column.+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++-- | 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).+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++-- | Whether the path facts force @cond@ true, false, or leave it 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 a fact's solution set sit inside, apart from, or across @cond@?+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,889 @@+{-# 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,+    widenToInt,+ )+import DataFrame.Internal.Column.Bitmap (+    Bitmap,+    bitmapTestBit,+ )+import DataFrame.Internal.Column.Encode (dictEncodeColumnUpTo)+import DataFrame.Internal.Column.Types+import DataFrame.Internal.Control.Concurrent (+    chunksFor,+    forkJoin,+    parThreshold,+    shouldParallelize,+ )+import DataFrame.Internal.Data.HashTable (htInsert, newHashTable)+import DataFrame.Internal.Data.PackedText (+    PackedOffsets (..),+    PackedTextData (..),+    offCount,+    packedLength,+    packedSlice,+    selAt,+    sliceEqBytes,+ )+import DataFrame.Internal.DataFrame (DataFrame (..), GroupedDataFrame (..))+import DataFrame.Internal.Grouping.Direct (+    ascendingCodeGroups,+    directGroupThreshold,+    rangeOf,+ )+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 via a hash table.+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 for multi-key lists.+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.+data FusedKey = FusedKey+    { fkCode :: Int -> Int+    , fkDomain :: !Int+    }++-- | Classify a key column for the fused multi-key direct path.+fusedKey :: Column -> Maybe FusedKey+fusedKey col@(UnboxedColumn Nothing (_ :: VU.Vector a))+    | Just v <- asIntKeys col+    , 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++-- | An integral key column as 'Int's; floats are refused.+asIntKeys :: Column -> Maybe (VU.Vector Int)+asIntKeys (UnboxedColumn Nothing (v :: VU.Vector a)) = case sFloating @a of+    STrue -> Nothing+    SFalse -> widenToInt v+asIntKeys _ = Nothing++-- | Fuse per-key codes into one mixed-radix code per row and group directly.+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 use the radix engine; valueIndices stays deferred.+            (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, or 'Nothing' past the direct threshold.+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.+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, then group by counting sort.+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.+dictGroupEnabled :: Bool+dictGroupEnabled = False++-- | Cardinality ceiling for the single-key dict-encode probe.+dictSingleThreshold :: Int+dictSingleThreshold = 4096++-- | The sequential grouping path: one open-addressing table over all rows.+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, identical in output to sequential.+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+            hashes <- computeRowHashesWithIO True n selectedCols+            parallelAssignGroups n hashes eqRow+     in -- rowToGroup stays an unforced argument (the module is -XStrict).+        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 a dense group id in first-appearance order.+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 (by hash).+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 an unboxed vector.+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 :: DataFrame -> [Int] -> Int -> Int -> Bool+eqKeyRow df indicesToGroup = case map (colEqRow . (columns df V.!)) indicesToGroup of+    [p] -> p+    preds ->+        let go [] _ _ = True+            go (p : ps) a b = p a b && go ps a b+         in go preds++-- | Per-column row equality respecting nulls.+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 gives equal strings equal codes.+    let eqV = case ptSel p of+            Just sel+                | ptCanonicalSel p ->+                    \a b -> selAt sel a == selAt sel b+            _ -> packedEqRow p+     in withNulls bm eqV+{-# INLINE colEqRow #-}++-- | Byte equality of two packed-text rows, with the layout matched once.+packedEqRow :: PackedTextData -> (Int -> Int -> Bool)+packedEqRow (PackedTextData arr offs sel _) = case (sel, offs) of+    (Nothing, Offs32 o) -> \a b -> a == b || eqBase (off32 o) a b+    (Nothing, Offs64 o) -> \a b -> a == b || eqBase (VU.unsafeIndex o) a b+    (Just s, Offs32 o) -> eqSel s (off32 o)+    (Just s, Offs64 o) -> eqSel s (VU.unsafeIndex o)+  where+    off32 o i = fromIntegral (VU.unsafeIndex o i)+    -- Base rows; equal ones share bytes, and -1 is the empty string.+    eqSel s off a b =+        let !ra = selAt s a+            !rb = selAt s b+         in ra == rb || eqBase off ra rb+    eqBase off ra rb =+        let span r = if r < 0 then (0, 0) else let o = off r in (o, off (r + 1) - o)+            (!oa, !la) = span ra+            (!ob, !lb) = span rb+         in sliceEqBytes arr oa la arr ob lb+    {-# INLINE eqSel #-}+    {-# INLINE eqBase #-}++-- | Wrap a value-equality with null handling (both null counts as equal).+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@ by stable counting sort.+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)++-- | Contiguous per-worker row ranges, one chunk per capability when parallel.+directRowChunks :: Int -> [(Int, Int)]+directRowChunks = chunksFor parThreshold++-- | Like 'directRowChunks' but over a code/group domain.+directCodeSlices :: Int -> [(Int, Int)]+directCodeSlices = chunksFor 4096++-- | The eager phases of the direct counting-sort grouping.+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 row share one word per code.+        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 huge frames: separate count and firstOcc arrays.+        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': merge, normalise and build the layout.+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))++-- | Row limit for the packed count/first-row encoding (@row + 1 < 2^31@).+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.+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 count and first-row encoding.+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 first occurrences from the packed chunk histograms.+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 count plus first occurrences.+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 each code's first occurrence.+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 into the offsets array.+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++-- | Scatter representative rows for every occupied code, in parallel.+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 through @codeToGroup@.+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 building @valueIndices@ from per-row codes.+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 histograms, rewritten into 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 write cursor.+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 its advancing 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.+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.+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.+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,720 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | Low-cardinality direct-indexed grouping fast path for dense Int key codes.+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 the direct grouping path accepts.+directGroupThreshold :: Int+directGroupThreshold = 1048576++-- | The grouping layout the hash path also produces.+data DirectGrouping = DirectGrouping+    { dgRowToGroup :: !(VU.Vector Int)+    , dgValueIndices :: !(VU.Vector Int)+    , dgOffsets :: !(VU.Vector Int)+    , dgNGroups :: !Int+    }++-- | Take the direct path for a clean small-range Int key, else 'Nothing'.+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 when parallel.+rowChunks :: Int -> [(Int, Int)]+rowChunks = chunksFor parThreshold++-- | Like 'rowChunks' but over the code domain.+codeSlices :: Int -> [(Int, Int)]+codeSlices = chunksFor 4096++-- un each action on its own thread and collect results in order.++-- | Build the grouping by counting sort on a per-row code in @[0, card)@.+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)++-- | The layout every aggregation needs, without the O(n) stable placement.+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: per-chunk histograms with first-occurrence tracking.+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 count, recording first rows.+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)++-- | Merge per-chunk counts and first rows over one slice of the code domain.+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: a two-level radix over packed @(code, row)@ words.+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': mapping, offsets and representatives.+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+    -- An identity code map lets us skip the per-row table lookup.+    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 into the offsets array.+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)++-- | The unique stable counting-sort permutation of @rowToGroup@.+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 histograms, 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+            ]+    -- Each chunk histogram becomes its write cursor: offs[g] + earlier chunks.+    _ <-+        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 placement into @valueIndices@.+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 offset-seeded cursors.+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, split into sub-chunks.+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 ()++-- | Code-domain size above which parallel passes use two-level radix.+twoLevelCardThreshold :: Int+twoLevelCardThreshold = 1024++-- | The two-level engine packs @(code, row)@ into one machine word.+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.+newtype PackedBuckets = PackedBuckets (VUM.IOVector Int)++packedVec :: PackedBuckets -> VUM.IOVector Int+packedVec (PackedBuckets v) = v++-- | Partition rows into buckets by top code bits as packed words.+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 count.+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 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, with oversized buckets run first.+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 packed slice.+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: dense ids for non-empty codes.+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,411 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE Strict #-}++-- | Partitioned group-by: counting-sort rows by hash bits, group per partition.+module DataFrame.Internal.Grouping.Partitioned (+    parallelAssignGroups,+    rtgFromVisOffs,+    numPartitionsFor,+) where++import Control.Concurrent (getNumCapabilities)+import Control.Monad (when)+import Data.Bits (countLeadingZeros, unsafeShiftL, unsafeShiftR)+import qualified Data.List as L+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 of a hash, ordering like signed 'Int'.+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 scaled by capabilities and rows.+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, returning @(valueIndices, offsets)@ canonically.+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+    link <- VUM.new (max 1 n)+    (partStart, sortedRows, sortedHash, sortedLink) <-+        partitionRows n hashes eqRow link p shift+    localGid <- VUM.new (max 1 n)+    gidByRow <- 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+        sortedLink+        gidByRow+        localGid+        canonBoxes+        nLocalGroups+    (globalBase, canonOf, nGroups) <- canonicalize p canonBoxes nLocalGroups+    assemble n p partStart sortedRows localGid globalBase canonOf nGroups++-- | Bucket every row index into its partition by a counting sort.+partitionRows ::+    Int ->+    VU.Vector Int ->+    (Int -> Int -> Bool) ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    IO (VU.Vector Int, VU.Vector Int, VU.Vector Int, VU.Vector Int)+partitionRows n hashes eqRow link p shift = do+    caps <- getNumCapabilities+    let chunks = rowChunks caps n+    cursors <-+        forkJoin [histChunk hashes eqRow link p shift lo hi | (lo, hi) <- chunks]+    -- Exclusive prefix over partitions and chunks; histograms become cursors.+    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)+    sortedLinkM <- VUM.new (max 1 n)+    forkJoin_+        [ scatterChunk hashes link shift cur sortedM sortedHashM sortedLinkM lo hi+        | ((lo, hi), cur) <- zip chunks cursors+        ]+    partStart <- VU.unsafeFreeze partStartM+    sortedRows <- VU.unsafeFreeze sortedM+    sortedHash <- VU.unsafeFreeze sortedHashM+    sortedLink <- VU.unsafeFreeze sortedLinkM+    pure (partStart, sortedRows, sortedHash, sortedLink)++-- | 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, linking rows to recent equal keys.+histChunk ::+    VU.Vector Int ->+    (Int -> Int -> Bool) ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    Int ->+    Int ->+    IO (VUM.IOVector Int)+histChunk hashes eqRow link p shift lo hi = do+    acc <- VUM.replicate p (0 :: Int)+    -- Two ways per set, so two hot keys sharing a set do not evict each other.+    cacheHash <- VUM.new (2 * linkCacheSets)+    cacheRow <- VUM.replicate (2 * linkCacheSets) (-1 :: Int)+    let hit !w !h !i = do+            r <- VUM.unsafeRead cacheRow w+            h' <- VUM.unsafeRead cacheHash w+            pure (if r >= 0 && h' == h && eqRow r i then r else -1)+        go !i+            | i >= hi = pure acc+            | otherwise = do+                let !h = VU.unsafeIndex hashes i+                    !pp = partIx shift h+                    !w0 = 2 * cacheSet h+                r0 <- hit w0 h i+                r <- if r0 >= 0 then pure r0 else hit (w0 + 1) h i+                VUM.unsafeWrite link i r+                when (r < 0) $ do+                    -- Way 0 holds the newest key; the older one moves to way 1.+                    VUM.unsafeRead cacheHash w0 >>= VUM.unsafeWrite cacheHash (w0 + 1)+                    VUM.unsafeRead cacheRow w0 >>= VUM.unsafeWrite cacheRow (w0 + 1)+                    VUM.unsafeWrite cacheHash w0 h+                    VUM.unsafeWrite cacheRow w0 i+                c <- VUM.unsafeRead acc pp+                VUM.unsafeWrite acc pp (c + 1)+                go (i + 1)+    go lo++-- | Cache set of a hash, mixing the weak low bits first.+cacheSet :: Int -> Int+cacheSet h =+    fromIntegral+        ( (fromIntegral h * 0x9E3779B97F4A7C15 :: Word64)+            `unsafeShiftR` (64 - linkCacheBits)+        )+{-# INLINE cacheSet #-}++linkCacheBits :: Int+linkCacheBits = 12++-- | Sets in each chunk's cache of recent keys (a power of two).+linkCacheSets :: Int+linkCacheSets = 1 `unsafeShiftL` linkCacheBits++-- | Scatter one row chunk through the chunk's private cursor.+scatterChunk ::+    VU.Vector Int ->+    VUM.IOVector Int ->+    Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    IO ()+scatterChunk hashes link shift cursor sortedM sortedHashM sortedLinkM 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.unsafeRead link i >>= VUM.unsafeWrite sortedLinkM pos+            VUM.unsafeWrite cursor pp (pos + 1)+            go (i + 1)++-- | Partition indices, largest first, so skewed partitions start early.+largestFirst :: VU.Vector Int -> VU.Vector Int+largestFirst partStart =+    let !p = VU.length partStart - 1+        size i = VU.unsafeIndex partStart (i + 1) - VU.unsafeIndex partStart i+     in VU.fromList+            (map snd (L.sortBy (flip compare) [(size i, i) | i <- [0 .. p - 1]]))++-- | Group each partition with its own hash table and rank its local groups.+runPartitions ::+    Int ->+    Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    (Int -> Int -> Bool) ->+    VU.Vector Int ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    VM.IOVector (VU.Vector Int) ->+    VUM.IOVector Int ->+    IO ()+runPartitions caps p partStart sortedRows sortedHash eqRow sortedLink gidByRow localGid canonBoxes nLocalGroups =+    pooledIndices caps p (groupPartition . VU.unsafeIndex order)+  where+    !order = largestFirst partStart+    groupPartition !pp = do+        let !s = VU.unsafeIndex partStart pp+            !e = VU.unsafeIndex partStart (pp + 1)+            !sz = e - s+        when (sz > 0) $ do+            -- Linked rows skip the table, so size it by the unlinked ones.+            let countUnlinked !pos !k+                    | pos >= e = pure k+                    | otherwise =+                        countUnlinked (pos + 1) (if VU.unsafeIndex sortedLink pos < 0 then k + 1 else k)+            nUnlinked <- countUnlinked s 0+            ht <- newHashTable nUnlinked+            repHashM <- VUM.new nUnlinked+            let loop !pos !nextGid+                    | pos >= e = pure nextGid+                    | otherwise = do+                        let !row = VU.unsafeIndex sortedRows pos+                            !h = VU.unsafeIndex sortedHash pos+                            !r = VU.unsafeIndex sortedLink pos+                        if r >= 0+                            then do+                                -- r: an earlier unlinked row of this partition.+                                gid <- VUM.unsafeRead gidByRow r+                                VUM.unsafeWrite localGid pos gid+                                loop (pos + 1) nextGid+                            else do+                                (gid, isNew) <- htInsert ht eqRow nextGid row h+                                VUM.unsafeWrite localGid pos gid+                                VUM.unsafeWrite gidByRow row 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++-- | Exclusive prefix sum of per-partition group counts into @globalBase@.+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)@ from global group ids.+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; counts becomes the 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+    let !order = largestFirst partStart+    pooledIndices caps p (doPartition . VU.unsafeIndex order)+    VUM.unsafeWrite offsM nGroups n+    offs <- VU.unsafeFreeze offsM+    vis <- VU.unsafeFreeze visM+    pure (vis, offs)++-- | Deferred @rowToGroup@ from @(valueIndices, offsets)@.+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,1431 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TupleSections #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeOperators #-}+{-# 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.Set as S+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.Aggregation.Kernel.Distinct (+    countDistinctColumn,+    countDistinctGroups,+ )+import DataFrame.Internal.Column+import DataFrame.Internal.DataFrame+import DataFrame.Internal.Expression+import qualified DataFrame.Internal.Grouping as G+import qualified DataFrame.Internal.Interpreter.Kernels as K+import Type.Reflection (+    SomeTypeRep (..),+    Typeable,+    typeRep,+ )++import Data.Int (Int16, Int32, Int64, Int8)+import Data.Word (Word64)+import GHC.Float (castDoubleToWord64)++-- Specializations for common aggregation types to avoid dictionary overhead.+{-# SPECIALIZE foldLinearGroups ::+    ((Double, Int) -> Double -> (Double, Int)) ->+    (Double, Int) ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    ((Double, Int) -> Float -> (Double, Int)) ->+    (Double, Int) ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    ((Double, Int) -> Int -> (Double, Int)) ->+    (Double, Int) ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    ((Double, Int) -> Int8 -> (Double, Int)) ->+    (Double, Int) ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    ((Double, Int) -> Int16 -> (Double, Int)) ->+    (Double, Int) ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    ((Double, Int) -> Int32 -> (Double, Int)) ->+    (Double, Int) ->+    Column ->+    VU.Vector Int ->+    Int ->+    Either DataFrameException Column+    #-}+{-# SPECIALIZE foldLinearGroups ::+    ((Double, Int) -> Int64 -> (Double, Int)) ->+    (Double, Int) ->+    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 ::+    ((Double, Int) -> 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+    #-}+-- Bool-mapping unary ops (e.g. 'not')+{-# SPECIALIZE mapColumn ::+    (Bool -> Bool) -> Column -> Either DataFrameException Column+    #-}++-- | The result of interpreting an expression; literals stay scalar.+data Value a where+    -- | A single value, not yet broadcast to any length.+    Scalar ::+        (Columnable a) =>+        !a ->+        -- | A flat column, one element per row (or per group once aggregated).+        Value a+    Flat ::+        (Columnable a) =>+        !Column ->+        -- | A grouped column: one 'Column' slice per group.+        Value a+    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++-- | Force a 'Value' into a flat 'Column' of the given length.+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@ in the cheapest form.+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)++-- | 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, degenerating when one is scalar.+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 #-}++-- | If-then-else over values already computed for every row.+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 (Flat cc) l r = do+    lc <- alongside cc l+    rc <- alongside cc r+    Flat <$> chooseRows cc lc rc+branchValue (Group cgs) l r =+    Group+        <$> V.imapM+            ( \i cc -> do+                lc <- alongside cc =<< groupAt i l+                rc <- alongside cc =<< groupAt i r+                chooseRows cc lc rc+            )+            cgs++shapeMismatch :: DataFrameException+shapeMismatch =+    AggregatedAndNonAggregatedException "if-then-else branches" "mismatched shapes"++-- | A branch as a column with one value per row of the condition column.+alongside ::+    forall a.+    (Columnable a) => Column -> Value a -> Either DataFrameException Column+alongside cc (Scalar x) = Right (broadcastScalar (columnLength cc) x)+alongside _ (Flat c) = Right c+alongside _ (Group _) = Left shapeMismatch++-- | The part of a branch that belongs to group @i@.+groupAt :: Int -> Value a -> Either DataFrameException (Value a)+groupAt _ v@(Scalar _) = Right v+groupAt i (Group gs) | i < V.length gs = Right (Flat (gs V.! i))+groupAt _ _ = Left shapeMismatch++chooseRows :: Column -> Column -> Column -> Either DataFrameException Column+chooseRows cc lc rc = do+    cs <- conditionVector cc+    let n = VU.length cs+    both <- mappendColumns lc rc+    pure (atIndicesStable (VU.imap (\i c -> if c then i else n + i) cs) both)++-- | A condition with nulls is an error.+conditionVector :: Column -> Either DataFrameException (VU.Vector Bool)+conditionVector (UnboxedColumn bm (v :: VU.Vector x))+    | Just Refl <- testEquality (typeRep @x) (typeRep @Bool)+    , maybe True (not . bitmapHasNulls (VU.length v)) bm =+        Right v+conditionVector c = toVector @Bool @VU.Vector c++-- | A value as a kernel operand: a literal or a null-free unboxed column.+operand :: forall a. (Columnable a) => Value a -> Maybe (K.Operand a)+operand (Scalar s) = Just (K.Literal s)+operand (Flat (UnboxedColumn Nothing (v :: VU.Vector x)))+    | Just Refl <- testEquality (typeRep @x) (typeRep @a) = Just (K.Column v)+operand _ = Nothing++fromOperand :: (Columnable a, VU.Unbox a) => K.Operand a -> Value a+fromOperand (K.Literal x) = Scalar x+fromOperand (K.Column v) = Flat (UnboxedColumn Nothing v)++-- | A column loop for a built-in binary operator, if it can run.+builtinBinary ::+    forall c b a.+    (Columnable c, Columnable b, Columnable a) =>+    T.Text -> Maybe (Value c -> Value b -> Maybe (Value a))+builtinBinary name = do+    Refl <- testEquality (typeRep @c) (typeRep @b)+    case (comparison, arithmetic, integral, textEquality, textContains) of+        (Just run, _, _, _, _) -> Just run+        (_, Just run, _, _, _) -> Just run+        (_, _, Just run, _, _) -> Just run+        (_, _, _, Just run, _) -> Just run+        (_, _, _, _, found) -> found+  where+    -- Int-only integer division and remainder.+    integral :: (c ~ b) => Maybe (Value c -> Value c -> Maybe (Value a))+    integral = do+        Refl <- testEquality (typeRep @a) (typeRep @Int)+        Refl <- testEquality (typeRep @c) (typeRep @Int)+        op <- case name of+            "div" -> Just div+            "mod" -> Just mod+            _ -> Nothing+        -- Only non-zero literal divisors: lazy evaluation might skip a zero.+        Just $ \l r -> case r of+            Scalar k | k /= 0 -> fromOperand <$> (K.zipOperands op <$> operand l <*> operand r)+            _ -> Nothing+    textContains :: (c ~ b) => Maybe (Value c -> Value c -> Maybe (Value a))+    textContains = do+        Refl <- testEquality (typeRep @c) (typeRep @T.Text)+        Refl <- testEquality (typeRep @a) (typeRep @Bool)+        case name of+            "contains" -> Just packedContains+            _ -> Nothing+    comparison :: (c ~ b) => Maybe (Value c -> Value c -> Maybe (Value a))+    comparison = do+        Refl <- testEquality (typeRep @a) (typeRep @Bool)+        op <- K.cmpOpByName name+        run <- K.compareOperands @c+        Just (\l r -> fromOperand <$> (run op <$> operand l <*> operand r))+    arithmetic :: (c ~ b) => Maybe (Value c -> Value c -> Maybe (Value a))+    arithmetic = do+        Refl <- testEquality (typeRep @a) (typeRep @c)+        op <- K.arithOpByName name+        run <- K.arithOperands @c+        case sUnbox @c of+            STrue -> Just (\l r -> fromOperand <$> (run op <$> operand l <*> operand r))+            SFalse -> Nothing+    textEquality :: (c ~ b) => Maybe (Value c -> Value c -> Maybe (Value a))+    textEquality = do+        Refl <- testEquality (typeRep @c) (typeRep @T.Text)+        Refl <- testEquality (typeRep @a) (typeRep @Bool)+        wantEq <- case name of+            "eq" -> Just True+            "neq" -> Just False+            _ -> Nothing+        Just (packedEquals wantEq)++-- | Column on the left, needle on the right; not used if needle has U+FFFD.+packedContains :: Value T.Text -> Value T.Text -> Maybe (Value Bool)+packedContains (Flat (PackedText Nothing p)) (Scalar t)+    | not (T.any (== '\xFFFD') t) =+        Just (Flat (UnboxedColumn Nothing (K.packedContainsLit p t)))+packedContains _ _ = Nothing++-- | Not used when the literal contains U+FFFD (see 'K.packedEqLit').+packedEquals :: Bool -> Value T.Text -> Value T.Text -> Maybe (Value Bool)+packedEquals wantEq l r = do+    (p, t) <- case (l, r) of+        (Flat (PackedText Nothing p), Scalar t) -> Just (p, t)+        (Scalar t, Flat (PackedText Nothing p)) -> Just (p, t)+        _ -> Nothing+    if T.any (== '\xFFFD') t+        then Nothing+        else Just (Flat (UnboxedColumn Nothing (K.packedEqLit wantEq p t)))++-- | As 'builtinBinary', for unary operators.+builtinUnary ::+    forall b a.+    (Columnable b, Columnable a) => T.Text -> Maybe (Value b -> Maybe (Value a))+builtinUnary "toDouble"+    | Just Refl <- testEquality (typeRep @b) (typeRep @Int)+    , Just Refl <- testEquality (typeRep @a) (typeRep @Double) =+        Just $ \v -> case operand v of+            Just (K.Column x) -> Just (Flat (UnboxedColumn Nothing (K.intToDouble x)))+            _ -> Nothing+builtinUnary "textLength"+    | Just Refl <- testEquality (typeRep @b) (typeRep @T.Text)+    , Just Refl <- testEquality (typeRep @a) (typeRep @Int) =+        Just $ \case+            Flat (PackedText Nothing p) -> Just (Flat (UnboxedColumn Nothing (K.packedCharCount p)))+            _ -> Nothing+builtinUnary _ = Nothing++applyUnary ::+    forall op b a.+    (UnaryOp op, Columnable b, Columnable a) =>+    op b a -> Value b -> Either DataFrameException (Value a)+applyUnary op v = case builtinUnary @b @a (unaryName op) >>= ($ v) of+    Just out -> Right out+    Nothing -> liftValue (fastUnaryFn @b @a (unaryName op) (unaryFn op)) v++applyBinary ::+    forall op c b a.+    (BinaryOp op, Columnable c, Columnable b, Columnable a) =>+    op c b a -> Value c -> Value b -> Either DataFrameException (Value a)+applyBinary op l r = case builtinBinary @c @b @a (binaryName op) >>= (\run -> run l r) of+    Just out -> Right out+    Nothing -> liftValue2 (binaryFn op) l r++-- | The rows of the frame an evaluation covers.+data Rows = AllRows | Rows !(VU.Vector Int)++data Env = Env {envFrame :: DataFrame, envRows :: Rows}++frameRows :: Env -> Int+frameRows env = fst (dataframeDimensions (envFrame env))++rowCount :: Env -> Int+rowCount env = case envRows env of+    AllRows -> frameRows env+    Rows sel -> VU.length sel++-- | Keep only the covered rows of a value computed for the whole frame.+restrict :: Rows -> Value a -> Value a+restrict (Rows sel) (Flat col) = Flat (atIndicesStable sel col)+restrict _ v = v++-- | The rows at these positions of the rows already covered.+narrow :: Env -> VU.Vector Int -> Env+narrow env positions = env{envRows = Rows rows}+  where+    rows = case envRows env of+        AllRows -> positions+        Rows sel -> VU.unsafeBackpermute sel positions++-- | Identifies an expression built only from columns, literals and built-ins.+data Key+    = KeyColumn T.Text SomeTypeRep+    | KeyLiteral Literal+    | KeyUnary T.Text SomeTypeRep Key+    | KeyBinary T.Text SomeTypeRep Key Key+    deriving (Eq, Ord)++data Literal+    = LitDouble Word64+    | LitInt Int+    | LitBool Bool+    | LitText T.Text+    deriving (Eq, Ord)++-- | Keys are only built for expressions up to this many nodes.+maxKeyNodes :: Int+maxKeyNodes = 32++keyOf :: (Columnable a) => Expr a -> Maybe Key+keyOf expr = fst <$> go maxKeyNodes expr+  where+    -- Returns the key and the node budget left.+    go :: forall x. (Columnable x) => Int -> Expr x -> Maybe (Key, Int)+    go budget _ | budget <= 0 = Nothing+    go budget e = case e of+        Col name -> Just (KeyColumn name (SomeTypeRep (typeRep @x)), budget - 1)+        Lit v -> (\l -> (KeyLiteral l, budget - 1)) <$> literalKey v+        Unary op (inner :: Expr b) -> do+            _ <- builtinUnary @b @x (unaryName op)+            (k, left) <- go (budget - 1) inner+            Just (KeyUnary (unaryName op) (SomeTypeRep (typeRep @b)) k, left)+        Binary op (l :: Expr c) (r :: Expr b) -> do+            _ <- builtinBinary @c @b @x (binaryName op)+            -- A division is shareable only with a non-zero literal divisor.+            if binaryName op `elem` ["div", "mod"]+                then case r of+                    Lit (k :: b)+                        | Just Refl <- testEquality (typeRep @b) (typeRep @Int)+                        , k /= 0 ->+                            Just ()+                    _ -> Nothing+                else Just ()+            (kl, afterLeft) <- go (budget - 1) l+            (kr, left) <- go afterLeft r+            Just (KeyBinary (binaryName op) (SomeTypeRep (typeRep @c)) kl kr, left)+        _ -> Nothing++literalKey :: forall a. (Columnable a) => a -> Maybe Literal+literalKey v+    | Just Refl <- testEquality (typeRep @a) (typeRep @Double) =+        Just (LitDouble (castDoubleToWord64 v))+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = Just (LitInt v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Bool) = Just (LitBool v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) = Just (LitText v)+    | otherwise = Nothing++-- | The keys of operator nodes that occur more than once in an expression.+repeatedKeys :: (Columnable a) => Expr a -> S.Set Key+repeatedKeys expr = M.keysSet (M.filter (> 1) (count expr M.empty))+  where+    count :: forall x. (Columnable x) => Expr x -> M.Map Key Int -> M.Map Key Int+    count e seen = case e of+        Unary _ inner -> count inner (note e seen)+        Binary _ l r -> count r (count l (note e seen))+        If c l r -> count r (count l (count c seen))+        _ -> seen+    note :: forall x. (Columnable x) => Expr x -> M.Map Key Int -> M.Map Key Int+    note e seen = maybe seen (\k -> M.insertWith (+) k (1 :: Int) seen) (keyOf e)++data Sharing = Sharing+    { repeated :: !(S.Set Key)+    , rowsAsked :: !(M.Map Key Int)+    -- ^ How many rows each repeated key has been computed for so far.+    , computed :: !(M.Map Key Column)+    -- ^ Results kept for the whole frame.+    , bytesKept :: !Int+    }++-- | Results kept for reuse during one evaluation may total this many bytes.+maxBytesKept :: Int+maxBytesKept = 256 * 1024 * 1024++newtype Eval a = Eval+    {runEval :: Sharing -> Either DataFrameException (a, Sharing)}++instance Functor Eval where+    fmap f (Eval m) = Eval (fmap (first f) . m)++instance Applicative Eval where+    pure x = Eval (\s -> Right (x, s))+    Eval mf <*> Eval mx = Eval $ \s -> do+        (f, s') <- mf s+        (x, s'') <- mx s'+        Right (f x, s'')++instance Monad Eval where+    Eval m >>= k = Eval $ \s -> do+        (x, s') <- m s+        runEval (k x) s'++failing :: Either DataFrameException a -> Eval a+failing r = Eval (\s -> fmap (,s) r)++inContext :: (Show e) => Expr e -> Eval a -> Eval a+inContext expr (Eval m) = Eval (addContext expr . m)++evalFlat ::+    (Columnable a) => DataFrame -> Expr a -> Either DataFrameException (Value a)+evalFlat df expr =+    fst+        <$> runEval+            (evalRows (Env df AllRows) expr)+            (Sharing (repeatedKeys expr) M.empty M.empty 0)++evalRows :: forall a. (Columnable a) => Env -> Expr a -> Eval (Value a)+evalRows env expr = case expr of+    Lit v -> pure (Scalar v)+    Unary{} -> evalShared env expr+    Binary{} -> evalShared env expr+    If cond l r -> inContext expr $ do+        c <- evalRows env cond+        case c of+            Scalar holds -> do+                -- The other branch runs on no rows; it can only report columns.+                let (chosen, other) = if holds then (l, r) else (r, l)+                _ <- evalRows (narrow env VU.empty) other+                evalRows env chosen+            Flat cc -> do+                cs <- failing (conditionVector cc)+                let (yes, no) = K.partitionRows cs+                lv <- evalArm env yes l+                rv <- evalArm env no r+                failing (mergeArms cs yes no lv rv)+            Group _ -> failing (Left (InternalException "grouped condition in a flat frame"))+    -- Aggregations and windows need every row; then keep the covered ones.+    _ -> failing (restrict (envRows env) <$> eval (FlatCtx (envFrame env)) expr)++-- | Evaluate an operator node, reusing a result kept for the whole frame.+evalShared :: forall a. (Columnable a) => Env -> Expr a -> Eval (Value a)+evalShared env expr = Eval $ \s -> case keyOf expr of+    Just key+        | S.member key (repeated s) -> case M.lookup key (computed s) of+            Just col -> Right (restrict (envRows env) (Flat col), s)+            Nothing -> runEval (askFor key) s+    _ -> runEval (evalNode env expr) s+  where+    total = frameRows env+    askFor key = Eval $ \s ->+        let asked = M.findWithDefault 0 key (rowsAsked s) + rowCount env+         in if asked >= total && bytesKept s + 8 * total <= maxBytesKept+                then do+                    (full, s') <- runEval (evalNode env{envRows = AllRows} expr) s+                    Right (restrict (envRows env) full, keep key full s')+                else runEval (evalNode env expr) s{rowsAsked = M.insert key asked (rowsAsked s)}+    keep :: Key -> Value a -> Sharing -> Sharing+    keep key (Flat col) s =+        s+            { computed = M.insert key col (computed s)+            , bytesKept = bytesKept s + columnBytes col+            }+    keep _ _ s = s++evalNode :: forall a. (Columnable a) => Env -> Expr a -> Eval (Value a)+evalNode env expr = inContext expr $ case expr of+    Unary op inner -> evalRows env inner >>= failing . applyUnary op+    Binary op l r -> do+        lv <- evalRows env l+        rv <- evalRows env r+        failing (applyBinary op lv rv)+    _ -> evalRows env expr++columnBytes :: Column -> Int+columnBytes (UnboxedColumn _ (v :: VU.Vector x))+    | Just Refl <- testEquality (typeRep @x) (typeRep @Bool) = VU.length v+    | otherwise = 8 * VU.length v+columnBytes c = 16 * columnLength c++-- | A branch computed for every covered row, or only for the rows choosing it.+data Arm a = ForEveryRow (Value a) | ForChosenRows (Value a)++-- | A column or literal is already there for every row; else compute lazily.+evalArm :: (Columnable a) => Env -> VU.Vector Int -> Expr a -> Eval (Arm a)+evalArm env chosen expr = case expr of+    Col _ -> ForEveryRow <$> evalRows env expr+    Lit _ -> ForEveryRow <$> evalRows env expr+    _ -> ForChosenRows <$> evalRows (narrow env chosen) expr++-- | Each row's value from the branch its condition chooses.+mergeArms ::+    forall a.+    (Columnable a) =>+    VU.Vector Bool ->+    VU.Vector Int ->+    VU.Vector Int ->+    Arm a ->+    Arm a ->+    Either DataFrameException (Value a)+mergeArms cs yes no l r = case (K.mergeBranches @a, branch l, branch r, sUnbox @a) of+    (Just run, Just lb, Just rb, STrue) -> Right (Flat (UnboxedColumn Nothing (run cs lb rb)))+    _+        | VU.null no -> Right (Flat lc)+        | VU.null yes -> Right (Flat rc)+        | otherwise -> do+            both <- mappendColumns lc rc+            -- Row i reads the next unused value of its branch.+            let place (y, f) c = if c then (y + 1, f) else (y, f + 1)+                taken = VU.prescanl' place (0, VU.length yes) cs+                source = VU.zipWith (\c (y, f) -> if c then y else f) cs taken+            Right (Flat (atIndicesStable source both))+  where+    branch :: Arm a -> Maybe (K.Branch a)+    branch arm = case arm of+        ForEveryRow v ->+            (\case K.Literal x -> K.Constant x; K.Column xs -> K.EveryRow xs)+                <$> operand v+        ForChosenRows v ->+            (\case K.Literal x -> K.Constant x; K.Column xs -> K.ChosenRows xs)+                <$> operand v+    -- The generic path wants each branch as one value per row choosing it.+    chosenColumn :: VU.Vector Int -> Arm a -> Column+    chosenColumn rows arm = case arm of+        ForEveryRow v -> materialize @a (VU.length rows) (restrict (Rows rows) v)+        ForChosenRows v -> materialize @a (VU.length rows) v+    lc = chosenColumn yes l+    rc = chosenColumn no r++-- | Annotate any 'TypeMismatchException' with the expression 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++-- | Given a flat column and grouping metadata, produce one 'Column' 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 #-}++-- | Coerce a column to type @a@, then apply @onResult@ to each element.+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)++-- | Wrap each element in 'Just' when the output type is @Maybe c@.+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+                }++-- | Evaluate an expression in a given context, producing a 'Value'.+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 (FlatCtx df) expr@(Unary{}) = evalFlat df expr+eval (FlatCtx df) expr@(Binary{}) = evalFlat df expr+eval (FlatCtx df) expr@(If{}) = evalFlat df expr+eval ctx expr@(Unary op (inner :: Expr b)) = addContext expr $ do+    v <- eval @b ctx inner+    applyUnary 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+        applyBinary 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: seeded FoldAgg on a bare Col folds over permuted indices.++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: countDistinct on a bare Col goes to the partitioned hash kernel.+eval ctx (Agg (MergeAgg "countDistinct" _ _ _ _) (Col name))+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int)+    , Just col <- getColumn name (ctxFrame ctx)+    , Just v <- distinctValue col =+        Right v+  where+    ctxFrame (FlatCtx df) = df+    ctxFrame (GroupCtx gdf) = fullDataframe gdf+    distinctValue col = case ctx of+        FlatCtx _ -> Scalar <$> countDistinctColumn col+        GroupCtx gdf ->+            Flat . fromUnboxedVector+                <$> countDistinctGroups (rowToGroup gdf) (numGroups gdf) col+-- 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 (FlatCtx df) (Agg (MergeAgg "mean" _ _ _ (_ :: acc -> a)) (Col name))+    | Just Refl <- testEquality (typeRep @a) (typeRep @Double)+    , Just v <- fusedMean df name =+        Right (Scalar v)+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+-- Fast path: whole-frame sum of a clean Int/Double column, one parallel loop.+eval (FlatCtx df) (Agg (FoldAgg "sum" Nothing (_ :: a -> b -> a)) inner)+    | Just Refl <- testEquality (typeRep @a) (typeRep @b)+    , Just v <- fusedSum df inner =+        Right (Scalar v)+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@.+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.+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.+data AggregationResult a+    = UnAggregated Column+    | Aggregated (TypedColumn a)++-- | Interpret an expression against a flat 'DataFrame', producing a column.+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'.+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++-- | The mean of a clean Int or Double column. Empty gives NaN, like 'mean'.+fusedMean :: DataFrame -> T.Text -> Maybe Double+fusedMean df name = do+    UnboxedColumn Nothing (v :: VU.Vector t) <- getColumn name df+    total <- case testEquality (typeRep @t) (typeRep @Int) of+        Just Refl -> Just (K.parSumBy fromIntegral v)+        Nothing -> case testEquality (typeRep @t) (typeRep @Double) of+            Just Refl -> Just (K.parSumBy id v)+            Nothing -> Nothing+    pure $! if VU.null v then 0 / 0 else total / fromIntegral (VU.length v)++-- | The sum of a column, or of an arithmetic combination with a literal.+fusedSum :: forall a. (Columnable a) => DataFrame -> Expr a -> Maybe a+fusedSum df expr = case expr of+    Col name -> withClean name Nothing+    Binary op (Col name :: Expr c) (Lit (k :: b))+        | Just Refl <- testEquality (typeRep @c) (typeRep @a)+        , Just Refl <- testEquality (typeRep @b) (typeRep @a) ->+            withClean name (Just (binaryName op, k, False))+    Binary op (Lit (k :: c)) (Col name :: Expr b)+        | Just Refl <- testEquality (typeRep @c) (typeRep @a)+        , Just Refl <- testEquality (typeRep @b) (typeRep @a) ->+            withClean name (Just (binaryName op, k, True))+    _ -> Nothing+  where+    -- The operator is picked by name at a concrete element type.+    withClean :: T.Text -> Maybe (T.Text, a, Bool) -> Maybe a+    withClean name combo = do+        UnboxedColumn Nothing (v :: VU.Vector t) <- getColumn name df+        Refl <- testEquality (typeRep @t) (typeRep @a)+        case testEquality (typeRep @a) (typeRep @Int) of+            Just Refl -> sumInt combo v+            Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+                Just Refl -> sumDouble combo v+                Nothing -> Nothing++sumInt :: Maybe (T.Text, Int, Bool) -> VU.Vector Int -> Maybe Int+sumInt combo v = case combo of+    Nothing -> Just (K.parSumBy id v)+    Just ("add", k, _) -> Just (K.parSumBy (+ k) v)+    Just ("mult", k, _) -> Just (K.parSumBy (* k) v)+    Just ("sub", k, False) -> Just (K.parSumBy (subtract k) v)+    Just ("sub", k, True) -> Just (K.parSumBy (k -) v)+    _ -> Nothing++sumDouble :: Maybe (T.Text, Double, Bool) -> VU.Vector Double -> Maybe Double+sumDouble combo v = case combo of+    Nothing -> Just (K.parSumBy id v)+    Just ("add", k, _) -> Just (K.parSumBy (+ k) v)+    Just ("mult", k, _) -> Just (K.parSumBy (* k) v)+    Just ("sub", k, False) -> Just (K.parSumBy (subtract k) v)+    Just ("sub", k, True) -> Just (K.parSumBy (k -) v)+    _ -> Nothing
+ src-internal/DataFrame/Internal/Interpreter/Kernels.hs view
@@ -0,0 +1,270 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeOperators #-}++module DataFrame.Internal.Interpreter.Kernels (+    Operand (..),+    Branch (..),+    CmpOp (..),+    cmpOpByName,+    ArithOp (..),+    arithOpByName,+    compareOperands,+    arithOperands,+    zipOperands,+    mergeBranches,+    partitionRows,+    intToDouble,+    packedEqLit,+    packedContainsLit,+    packedCharCount,+    parSumBy,+) where++import Control.Monad.ST (runST)+import Data.Bits ((.&.))+import qualified Data.Text as T+import qualified Data.Text.Array as A+import Data.Text.Internal (Text (Text))+import Data.Type.Equality (testEquality, (:~:) (Refl))+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import Type.Reflection (Typeable, typeRep)++import DataFrame.Internal.Column.Operations (parGenerateUnboxedInline)+import DataFrame.Internal.Control.Concurrent (parThreshold, parallelChunks)+import DataFrame.Internal.Data.PackedText (+    PackedOffsets (..),+    PackedTextData (..),+    packedLength,+    selAt,+    sliceEqBytes,+ )+import System.IO.Unsafe (unsafePerformIO)++-- | One side of an operation: the same value in every row, or one per row.+data Operand a = Literal !a | Column !(VU.Vector a)++data CmpOp = CmpLt | CmpLeq | CmpGt | CmpGeq | CmpEq | CmpNeq+    deriving (Eq, Show)++cmpOpByName :: T.Text -> Maybe CmpOp+cmpOpByName name = case name of+    "lt" -> Just CmpLt+    "leq" -> Just CmpLeq+    "gt" -> Just CmpGt+    "geq" -> Just CmpGeq+    "eq" -> Just CmpEq+    "neq" -> Just CmpNeq+    _ -> Nothing++data ArithOp = ArithAdd | ArithSub | ArithMult+    deriving (Eq, Show)++arithOpByName :: T.Text -> Maybe ArithOp+arithOpByName name = case name of+    "add" -> Just ArithAdd+    "sub" -> Just ArithSub+    "mult" -> Just ArithMult+    _ -> Nothing++-- | Runs @k@ with the element type known, only for 'Double' or 'Int'.+withNumeric ::+    forall a r.+    (Typeable a) =>+    (forall x. (VU.Unbox x, Ord x, Num x) => a :~: x -> r) -> Maybe r+withNumeric k+    | Just Refl <- testEquality (typeRep @a) (typeRep @Double) =+        Just (k @Double Refl)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = Just (k @Int Refl)+    | otherwise = Nothing+{-# INLINE withNumeric #-}++compareOperands ::+    forall a.+    (Typeable a) => Maybe (CmpOp -> Operand a -> Operand a -> Operand Bool)+compareOperands = withNumeric @a (\Refl -> compareLoop)++arithOperands ::+    forall a. (Typeable a) => Maybe (ArithOp -> Operand a -> Operand a -> Operand a)+arithOperands = withNumeric @a (\Refl -> arithLoop)++-- | One branch of an if-then-else: constant, every-row, or chosen-rows values.+data Branch a = Constant !a | EveryRow !(VU.Vector a) | ChosenRows !(VU.Vector a)++-- | Each row's value from the branch its condition chooses.+mergeBranches ::+    forall a.+    (Typeable a) => Maybe (VU.Vector Bool -> Branch a -> Branch a -> VU.Vector a)+mergeBranches = withNumeric @a (\Refl -> mergeLoop)++-- The operator must be fully applied or it is not inlined into the loop.+compareLoop ::+    (VU.Unbox a, Ord a) => CmpOp -> Operand a -> Operand a -> Operand Bool+compareLoop op l r = case op of+    CmpLt -> zipOperands (<) l r+    CmpLeq -> zipOperands (<=) l r+    CmpGt -> zipOperands (>) l r+    CmpGeq -> zipOperands (>=) l r+    CmpEq -> zipOperands (==) l r+    CmpNeq -> zipOperands (/=) l r+{-# INLINE compareLoop #-}++arithLoop ::+    (VU.Unbox a, Num a) => ArithOp -> Operand a -> Operand a -> Operand a+arithLoop op l r = case op of+    ArithAdd -> zipOperands (+) l r+    ArithSub -> zipOperands (-) l r+    ArithMult -> zipOperands (*) l r+{-# INLINE arithLoop #-}++zipOperands ::+    (VU.Unbox a, VU.Unbox b, VU.Unbox c) =>+    (a -> b -> c) -> Operand a -> Operand b -> Operand c+zipOperands f (Literal x) (Literal y) = Literal (f x y)+zipOperands f (Column xs) (Literal y) =+    Column+        (parGenerateUnboxedInline (VU.length xs) (\i -> f (VU.unsafeIndex xs i) y))+zipOperands f (Literal x) (Column ys) =+    Column (parGenerateUnboxedInline (VU.length ys) (f x . VU.unsafeIndex ys))+zipOperands f (Column xs) (Column ys) =+    Column+        ( parGenerateUnboxedInline+            (min (VU.length xs) (VU.length ys))+            (\i -> f (VU.unsafeIndex xs i) (VU.unsafeIndex ys i))+        )+{-# INLINE zipOperands #-}++mergeLoop ::+    (VU.Unbox a) => VU.Vector Bool -> Branch a -> Branch a -> VU.Vector a+mergeLoop cs l r = case (l, r) of+    -- With a value per chosen row, the output is filled in row order.+    (ChosenRows _, _) -> inOrder+    (_, ChosenRows _) -> inOrder+    _ ->+        parGenerateUnboxedInline+            n+            (\i -> if VU.unsafeIndex cs i then at l i 0 else at r i 0)+  where+    n = VU.length cs+    -- Row i's value, given how many values this branch has given out.+    at (Constant x) _ _ = x+    at (EveryRow xs) i _ = VU.unsafeIndex xs i+    at (ChosenRows xs) _ used = VU.unsafeIndex xs used+    {-# INLINE at #-}+    inOrder = runST $ do+        out <- VUM.unsafeNew n+        let go !i !lUsed !rUsed+                | i >= n = pure ()+                | VU.unsafeIndex cs i =+                    VUM.unsafeWrite out i (at l i lUsed) >> go (i + 1) (lUsed + 1) rUsed+                | otherwise =+                    VUM.unsafeWrite out i (at r i rUsed) >> go (i + 1) lUsed (rUsed + 1)+        go 0 0 0+        VU.unsafeFreeze out+{-# INLINE mergeLoop #-}++-- | The positions where the condition holds and where it does not.+partitionRows :: VU.Vector Bool -> (VU.Vector Int, VU.Vector Int)+partitionRows cs = runST $ do+    let n = VU.length cs+    yes <- VUM.unsafeNew n+    no <- VUM.unsafeNew n+    -- Write to both and advance one side, so there is no data-dependent branch.+    let go !i !ny !nn+            | i >= n = pure (ny, nn)+            | otherwise = do+                let c = fromEnum (VU.unsafeIndex cs i)+                VUM.unsafeWrite yes ny i+                VUM.unsafeWrite no nn i+                go (i + 1) (ny + c) (nn + 1 - c)+    (ny, nn) <- go 0 0 0+    (,) <$> VU.unsafeFreeze (VUM.take ny yes) <*> VU.unsafeFreeze (VUM.take nn no)++intToDouble :: VU.Vector Int -> VU.Vector Double+intToDouble v = parGenerateUnboxedInline (VU.length v) (fromIntegral . VU.unsafeIndex v)+{-# NOINLINE intToDouble #-}++-- | Whether each row of a packed Text column equals a literal.+packedEqLit :: Bool -> PackedTextData -> T.Text -> VU.Vector Bool+packedEqLit wantEq p@(PackedTextData arr offs sel _) (Text la lo ll) =+    -- Match the layout once so each row costs two offset reads and a compare.+    case (sel, offs) of+        (Nothing, Offs32 o) -> gen id (off32 o)+        (Nothing, Offs64 o) -> gen id (VU.unsafeIndex o)+        (Just s, Offs32 o) -> gen (selAt s) (off32 o)+        (Just s, Offs64 o) -> gen (selAt s) (VU.unsafeIndex o)+  where+    off32 o i = fromIntegral (VU.unsafeIndex o i)+    -- A negative base row is the empty string.+    gen base off =+        parGenerateUnboxedInline (packedLength p) $ \i ->+            let !r = base i+                !o = if r < 0 then 0 else off r+                !l = if r < 0 then 0 else off (r + 1) - o+             in (l == ll && sliceEqBytes arr o l la lo ll) == wantEq+    {-# INLINE gen #-}+{-# NOINLINE packedEqLit #-}++-- | Whether each row of a packed Text column contains a literal substring.+packedContainsLit :: PackedTextData -> T.Text -> VU.Vector Bool+packedContainsLit p@(PackedTextData arr offs sel _) (Text la lo ll)+    | ll == 0 = VU.replicate (packedLength p) True+    | otherwise = case (sel, offs) of+        (Nothing, Offs32 o) -> gen id (off32 o)+        (Nothing, Offs64 o) -> gen id (VU.unsafeIndex o)+        (Just s, Offs32 o) -> gen (selAt s) (off32 o)+        (Just s, Offs64 o) -> gen (selAt s) (VU.unsafeIndex o)+  where+    off32 o i = fromIntegral (VU.unsafeIndex o i)+    !first = A.unsafeIndex la lo+    gen base off =+        parGenerateUnboxedInline (packedLength p) $ \i ->+            let !r = base i+             in r >= 0 && search (off r) (off (r + 1))+    {-# INLINE gen #-}+    -- Scan for the needle's first byte, then compare the rest.+    search !o !end = go o+      where+        !lastStart = end - ll+        go !k+            | k > lastStart = False+            | A.unsafeIndex arr k == first && sliceEqBytes arr k ll la lo ll = True+            | otherwise = go (k + 1)+{-# NOINLINE packedContainsLit #-}++-- | Character count of each row of a packed Text column (UTF-8 lead bytes).+packedCharCount :: PackedTextData -> VU.Vector Int+packedCharCount p@(PackedTextData arr offs sel _) = case (sel, offs) of+    (Nothing, Offs32 o) -> gen id (off32 o)+    (Nothing, Offs64 o) -> gen id (VU.unsafeIndex o)+    (Just s, Offs32 o) -> gen (selAt s) (off32 o)+    (Just s, Offs64 o) -> gen (selAt s) (VU.unsafeIndex o)+  where+    off32 o i = fromIntegral (VU.unsafeIndex o i)+    gen base off =+        parGenerateUnboxedInline (packedLength p) $ \i ->+            let !r = base i+             in if r < 0 then 0 else count (off r) (off (r + 1)) 0+    {-# INLINE gen #-}+    count !k !end !n+        | k >= end = n+        | A.unsafeIndex arr k .&. 0xC0 == 0x80 = count (k + 1) end n+        | otherwise = count (k + 1) end (n + 1)+{-# NOINLINE packedCharCount #-}++-- | @sum (map f v)@ as one parallel loop with no intermediate column.+parSumBy :: (VU.Unbox a, Num b) => (a -> b) -> VU.Vector a -> b+parSumBy f v = unsafePerformIO $ do+    parts <-+        parallelChunks parThreshold (VU.length v) (\lo hi -> pure $! go lo hi 0)+    pure $! sum parts+  where+    go !i !hi !acc+        | i >= hi = acc+        | otherwise = go (i + 1) hi (acc + f (VU.unsafeIndex v i))+{-# INLINE parSumBy #-}
+ src-internal/DataFrame/Internal/Row.hs view
@@ -0,0 +1,194 @@+{-# 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+    -- Null is its own constructor to avoid unpacking a Value (Maybe a).+    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 to represent rows as heterogeneous 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.+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, in natural order.+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 with only the named columns.+toRowVector :: [T.Text] -> DataFrame -> V.Vector Row+toRowVector names df = V.generate (fst (dataframeDimensions df)) (mkRowRep df names)++-- | Given a row, get the value associated with a field.+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"++-- Rows follow the caller's requested column order, not 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,207 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | Row-hash kernels with a parallel driver, feeding grouping and joins.+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,+ )++-- | Row count from which the parallel hash beats the sequential single range.+parRowHashThreshold :: Int+parRowHashThreshold = 200000++-- | Compute the per-row key hash over the selected key columns.+computeRowHashesIO :: Int -> [Column] -> IO (VU.Vector Int)+computeRowHashesIO = computeRowHashesWithIO False++-- | 'computeRowHashesIO' with an explicit dictionary-code switch.+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 a row range into @mv@.+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.+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.+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.+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,109 @@+-- | The curated public surface of dataframe-core: types, constraints, verbs.+module DataFrame.Core (+    -- * The DataFrame+    DataFrame,+    GroupedDataFrame,+    empty,+    fromColumns,+    (=:),+    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,+    toSomeExpr,+    fromSomeExpr,+    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,+    fromColumns,+    insertColumn,+    null,+    toCsv,+    toCsv',+    toMarkdown,+    toMarkdown',+    toSeparated,+    (=:),+ )+import DataFrame.Internal.Expression (+    Expr,+    NamedExpr,+    eSize,+    fromSomeExpr,+    prettyPrint,+    prettyPrintWidth,+    toNamedExpr,+    toSomeExpr,+ )+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,6 @@+-- | Public name for dataframe's expression operators (re-export shim).+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,25 +9,29 @@     -- * 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 import DataFrame.Typed.Schema (KnownSchema (..)) import DataFrame.Typed.Types (TypedDataFrame (..)) -{- | Validate that an untyped 'DataFrame' matches the expected schema @cols@,-then wrap it. Returns 'Nothing' on mismatch.--}+-- | Validate an untyped 'DataFrame' against schema @cols@ and wrap it. freeze ::     forall cols. (KnownSchema cols) => D.DataFrame -> Maybe (TypedDataFrame cols) freeze df = case validateSchema @cols df of@@ -43,15 +47,25 @@     Left err -> Left err     Right _ -> Right (TDF df) -{- | Unwrap a typed DataFrame back to the untyped representation.-Always safe; discards type information.--}+-- | Like 'freezeWithError' but throws a 'DataFrameException' on mismatch.+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. thaw :: TypedDataFrame cols -> D.DataFrame thaw (TDF df) = df -{- | Wrap an untyped DataFrame without any validation.-Used internally after delegation where the library guarantees schema correctness.--}+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. unsafeFreeze :: D.DataFrame -> TypedDataFrame cols unsafeFreeze = TDF @@ -82,17 +96,13 @@                             <> ", got "                             <> T.pack (C.columnTypeString col) -{- | Check if a Column's element type matches the expected SomeTypeRep.-For nullable columns (those with a bitmap), @Maybe a@ in the schema matches-a column whose inner type is @a@, since we store nullable data as-@BoxedColumn (Just bm) a@ or @UnboxedColumn (Just bm) a@ rather than-@Column (Maybe a)@.--}+-- | Check that a Column's element type matches the expected SomeTypeRep. matchesType :: SomeTypeRep -> C.Column -> Bool matchesType expected col =     let expectedStr = show expected         colTypeStr = C.columnTypeString col      in expectedStr == colTypeStr-            || ( C.hasMissing col -- nullable column: schema says "Maybe X", column stores "X" with a bitmap+            -- A nullable column stores "X" under a "Maybe X" schema type.+            || ( C.hasMissing col                     && Just colTypeStr == stripPrefix "Maybe " expectedStr                )
src/DataFrame/Typed/Generic.hs view
@@ -9,37 +9,7 @@ {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} -{- |-Module      : DataFrame.Typed.Generic-License     : MIT--Generic-based opt-in for record-to-schema derivation. Mirrors the Template-Haskell splice in "DataFrame.Typed.TH" but builds the schema type from a-@GHC.Generics.Generic@ instance instead of @reify@.--Use it like this:--@-data Order = Order-  { orderId :: Int64-  , region  :: Text-  , amount  :: Double-  } deriving (Show, Eq, Generic)--type OrderSchema = SchemaOf Order--instance HasSchema Order OrderSchema where-  toColumns   = genericToColumns-  fromColumns = genericFromColumns-@--Field names are translated with the @CamelCase -> snake_case@ rule-(matching 'DataFrame.Typed.TH.camelToSnake'); use 'SchemaOfRaw' if you-want the schema to keep the record selector names verbatim — in that-case you cannot use 'genericToColumns' \/ 'genericFromColumns' and must-either hand-roll the instance or use the TH splice with a custom name-transform.--}+-- | Generic-based opt-in for record-to-schema derivation. module DataFrame.Typed.Generic (     -- * Type-level schema derivation     NameCase (..),@@ -86,32 +56,24 @@ 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-'RepToSchema'.--* 'SnakeCase' — translate @camelCaseField@ to @\"camel_case_field\"@.-* 'IdentityCase' — keep the selector name verbatim.--}+-- | Field-name policy applied to record selectors for 'RepToSchema'. data NameCase = SnakeCase | IdentityCase -{- | The schema type @[Column 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+-- | The schema type derived from the 'Rep' of a record type.+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     TransformName 'IdentityCase s = s --- | Type-level camelCase -> snake_case. Matches 'camelToSnake' at the value level.+-- | Type-level camelCase to snake_case, matching 'camelToSnake'. type family CamelToSnake (s :: Symbol) :: Symbol where     CamelToSnake s = SnakeStart (UnconsSymbol s) @@ -142,9 +104,7 @@ -- | Identity-cased schema derived from @a@'s 'Generic' representation. type SchemaOfRaw a = RepToSchema 'IdentityCase (Rep a) -{- | Walks the 'Rep' tree of a record, producing or consuming a list of-named columns. Used by 'genericToColumns' \/ 'genericFromColumns'.--}+-- | Walks the 'Rep' tree of a record, producing or consuming named columns. class GHasColumns (r :: Type -> Type) where     gToColumns :: [r p] -> [(T.Text, C.Column)]     gFromColumns :: D.DataFrame -> Either T.Text [r p]@@ -184,22 +144,12 @@         v <- requireColumn @a colName df         pure (map (M1 . K1) (VB.toList v)) -{- | Default implementation of 'DataFrame.Typed.Record.toColumns' for any-@Generic@ record. Field names are translated with @camelCase -> snake_case@.--@-instance HasSchema Order (SchemaOf Order) where-  toColumns   = genericToColumns-  fromColumns = genericFromColumns-@--}+-- | Default 'DataFrame.Typed.Record.recordsToColumns' for any Generic record. genericToColumns ::     forall a. (Generic a, GHasColumns (Rep a)) => [a] -> [(T.Text, C.Column)] genericToColumns = gToColumns . map from -{- | Default implementation of 'DataFrame.Typed.Record.fromColumns' for any-@Generic@ record.--}+-- | Default 'DataFrame.Typed.Record.recordsFromFrame' for any Generic record. genericFromColumns ::     forall a. (Generic a, GHasColumns (Rep a)) => D.DataFrame -> Either T.Text [a] genericFromColumns df = fmap (map to) (gFromColumns df)
src/DataFrame/Typed/Record.hs view
@@ -5,15 +5,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} -{- |-Module      : DataFrame.Typed.Record-License     : MIT--Bridge a Haskell record type to a typed dataframe schema. Instances are-typically generated by the @deriveSchemaFromType@ Template Haskell splice-(see "DataFrame.Typed.TH"), but can be written by hand as well, or-plugged into the @Generic@-derived defaults from "DataFrame.Typed.Generic".--}+-- | Bridge a Haskell record type to a typed dataframe schema. module DataFrame.Typed.Record (     -- * Class     HasSchema (..),@@ -35,48 +27,24 @@ import qualified Data.Vector as VB  import qualified DataFrame.Internal.Column as C-import DataFrame.Internal.DataFrame (fromNamedColumns)+import DataFrame.Internal.DataFrame (fromColumns) import qualified DataFrame.Internal.DataFrame as D import DataFrame.Typed.Types (TypedDataFrame (..))--{- | Bridge a Haskell record type @a@ to a typed-dataframe schema.--The schema is exposed as an associated type family 'Schema' so that-instances can pick it up from a 'GHC.Generics.Rep' computation (see-'DataFrame.Typed.Generic.SchemaOf') or from an explicit list emitted by-'DataFrame.Typed.TH.deriveSchemaFromType'.+import GHC.TypeLits -@toColumns@ explodes a list of records into a list of named columns.-@fromColumns@ reconstructs the records from a 'D.DataFrame', returning-@Left err@ if a column is missing or has the wrong type.--}+-- | Bridge a Haskell record type to a typed-dataframe schema. class HasSchema a where-    type Schema a :: [Type]-    toColumns :: [a] -> [(T.Text, C.Column)]-    fromColumns :: D.DataFrame -> Either T.Text [a]--{- | Build an untyped 'D.DataFrame' from a list of records.--@-data Order = Order { orderId :: Int64, region :: Text, amount :: Double }-\$(deriveSchemaFromType ''Order)--xs :: [Order]-xs = [Order 1 "us" 10.0, Order 2 "eu" 20.0]+    type Schema a :: [(Symbol, Type)]+    recordsToColumns :: [a] -> [(T.Text, C.Column)]+    recordsFromFrame :: D.DataFrame -> Either T.Text [a] -df :: DataFrame-df = fromRecords xs-@--}+-- | Build an untyped 'D.DataFrame' from a list of records. fromRecords :: (HasSchema a) => [a] -> D.DataFrame-fromRecords = fromNamedColumns . toColumns--{- | Parse a list of records out of an untyped 'D.DataFrame'.+fromRecords = fromColumns . recordsToColumns -Returns @Left err@ on schema mismatch (missing column, wrong type).--}+-- | Parse a list of records out of an untyped 'D.DataFrame'. toRecords :: (HasSchema a) => D.DataFrame -> Either T.Text [a]-toRecords = fromColumns+toRecords = recordsFromFrame  -- | Like 'fromRecords' but returns a 'TypedDataFrame' tagged with the schema. fromRecordsTyped :: forall a. (HasSchema a) => [a] -> TypedDataFrame (Schema a)@@ -85,13 +53,9 @@ -- | Like 'toRecords' but accepts a 'TypedDataFrame'. toRecordsTyped ::     forall a. (HasSchema a) => TypedDataFrame (Schema a) -> Either T.Text [a]-toRecordsTyped (TDF df) = fromColumns df--{- | Extract a column as a boxed vector by name, returning a 'T.Text' error-on missing column or type mismatch.+toRecordsTyped (TDF df) = recordsFromFrame df -Used by code generated by 'DataFrame.Typed.TH.deriveSchemaFromType'.--}+-- | Extract a column as a boxed vector by name, or a 'T.Text' error. requireColumn ::     forall a.     (C.Columnable a) => T.Text -> D.DataFrame -> Either T.Text (VB.Vector 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,73 +62,83 @@      -- * 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") -{- | Like 'Lookup', but returns a harmless fallback ('Int') instead of-'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+-- | Like 'Lookup' but falls back to 'Int' instead of a 'TypeError'.+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 +147,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 +161,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 +186,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 +219,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 +232,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 +242,7 @@         (found :: Bool)         (name :: Symbol)         (rest :: [Symbol])-        (cols :: [Type]) ::+        (cols :: [(Symbol, Type)]) ::         Constraint     where     AssertAllPresentHelper 'True name rest cols = AssertAllPresent rest cols@@ -224,14 +250,12 @@         TypeError             ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema") -{- | Assert that each join key has the same element type in both schemas,-modulo 'Maybe'-wrapping on either side (the runtime join matches a nullable-key column against a plain one). Use together with 'AssertAllPresent', which-reports keys missing from either schema; absent keys are skipped here so the-error fires exactly once.--}+-- | Assert each join key has the same element type in both schemas. 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 +282,134 @@                 ':<>: 'Text " in the right table"             ) -{- | Strip 'Maybe' from all columns. Used by 'filterAllJust'.+type family+    AssertDisjoint (left :: [(Symbol, Type)]) (right :: [(Symbol, Type)]) ::+        Constraint+    where+    AssertDisjoint left right =+        AssertDisjointHelper (SharedNames left right) left right -@Column "x" (Maybe Double)@ becomes @Column "x" Double@.-@Column "y" Int@ stays @Column "y" Int@.--}-type family StripAllMaybe (cols :: [Type]) :: [Type] where-    StripAllMaybe '[] = '[]-    StripAllMaybe (Column n (Maybe a) ': rest) = Column n a ': StripAllMaybe rest-    StripAllMaybe (Column n a ': rest) = Column n a ': StripAllMaybe rest+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"+            ) -{- | Strip 'Maybe' from a single named column. Used by 'filterJust'.+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+        ) -@StripMaybeAt "x" '[Column "x" (Maybe Double), Column "y" Int]@-  = @'[Column "x" Double, Column "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+-- | Is @a@ a real, unboxed numeric type?+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 a column is not a real-number type.+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 unboxed type.+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'.+type family StripAllMaybe (cols :: [(Symbol, Type)]) :: [(Symbol, Type)] where+    StripAllMaybe '[] = '[]+    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'.+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 +417,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 +448,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 +473,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 +488,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 +518,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 +537,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 +557,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 +575,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 +589,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,13 +605,17 @@  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 --- | A class that provides a list of 'Text' values for a type-level list of Symbols.+-- | The column names a schema declares, in schema order.+schemaColumnNames :: forall cols. (KnownSchema cols) => [T.Text]+schemaColumnNames = map fst (schemaEvidence @cols)++-- | A class giving a list of 'Text' values for a type-level list of Symbols. class AllKnownSymbol (names :: [Symbol]) where     symbolVals :: [T.Text] 
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,12 @@  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'.+import DataFrame.Internal.Expression (Expr, NamedExpr, SomeExpr (..)) -The type parameter @cols@ is a type-level list of @Column name ty@ entries-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}+-- | A phantom-typed wrapper over the untyped 'DataFrame'.+newtype TypedDataFrame (cols :: [(Symbol, Type)]) = TDF {unTDF :: D.DataFrame}  instance Show (TypedDataFrame cols) where     show (TDF df) = show df@@ -55,43 +53,35 @@ 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@, yielding @a@.+newtype TExpr (cols :: [(Symbol, Type)]) a = TExpr {unTExpr :: Expr a} -{- | A typed expression validated against schema @cols@, producing values of type @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 -Unlike the untyped 'Expr a', a 'TExpr' can only be constructed through-type-safe combinators ('col', 'lit', arithmetic operations) that verify-column references exist in the schema with the correct type.+-- | The typed counterpart of an untyped expression type for schema @cols@.+type family AsTExpr (cols :: [(Symbol, Type)]) (e :: Type) :: Type where+    AsTExpr cols (Expr r) = TExpr cols r -Use 'unTExpr' to extract the underlying 'Expr' for delegation to the untyped API.--}-newtype TExpr (cols :: [Type]) a = TExpr {unTExpr :: Expr a}+-- | 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-@aggs@ phantom list. End users never construct this directly — they-compose 'DataFrame.Typed.Aggregate.as' entries with @(.)@ and let-'DataFrame.Typed.Aggregate.aggregate' apply the composition to-'TAggNil'.--@-as \@\"total\"   (F.sum  salary)-  . as \@\"avg_age\" (F.mean age)-@--}-data TAgg (keys :: [Symbol]) (cols :: [Type]) (aggs :: [Type]) where+-- | Internal aggregation chain; users compose it via the Aggregators API.+data TAgg (keys :: [Symbol]) (cols :: [(Symbol, Type)]) (aggs :: [(Symbol, Type)]) where     TAggNil :: TAgg keys cols '[]     TAggCons ::         (Columnable a) =>@@ -101,14 +91,12 @@         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).--}+-- | Extract the runtime 'NamedExpr' list from a 'TAgg' in declaration order. taggToNamedExprs :: TAgg keys cols aggs -> [NamedExpr] taggToNamedExprs = reverse . go   where     go :: TAgg keys cols aggs -> [NamedExpr]     go TAggNil = []-    go (TAggCons name (TExpr expr) rest) = (name, UExpr expr) : go rest+    go (TAggCons name (TExpr expr) rest) = (name, SomeExpr expr) : go rest
src/DataFrame/Typed/Util.hs view
@@ -4,15 +4,7 @@  import Data.Char (isUpper, toLower) -{- | Convert a camelCase identifier to snake_case.-->>> camelToSnake "orderId"-"order_id"->>> camelToSnake "amountUS"-"amount_u_s"->>> camelToSnake "region"-"region"--}+-- | Convert a camelCase identifier to snake_case. camelToSnake :: String -> String camelToSnake [] = [] camelToSnake (c : cs) = toLower c : go cs