packages feed

dataframe-core 1.0.2.0 → 1.1.0.0

raw patch · 24 files changed

+4302/−81 lines, 24 filesdep +primitivedep ~textPVP ok

version bump matches the API change (PVP)

Dependencies added: primitive

Dependency ranges changed: text

API changes (from Hackage documentation)

+ DataFrame.Internal.AggKernel: RCount :: Reduction
+ DataFrame.Internal.AggKernel: RMax :: Reduction
+ DataFrame.Internal.AggKernel: RMean :: Reduction
+ DataFrame.Internal.AggKernel: RMin :: Reduction
+ DataFrame.Internal.AggKernel: RStd :: Reduction
+ DataFrame.Internal.AggKernel: RSum :: Reduction
+ DataFrame.Internal.AggKernel: RTop2Sum :: Reduction
+ DataFrame.Internal.AggKernel: RVar :: Reduction
+ DataFrame.Internal.AggKernel: data Reduction
+ DataFrame.Internal.AggKernel: instance GHC.Classes.Eq DataFrame.Internal.AggKernel.Reduction
+ DataFrame.Internal.AggKernel: instance GHC.Show.Show DataFrame.Internal.AggKernel.Reduction
+ DataFrame.Internal.AggKernel: scatterColumnToDouble :: Column -> Maybe (Vector Double)
+ DataFrame.Internal.AggKernel: scatterReduce :: Reduction -> Vector Int -> Int -> Column -> Maybe Column
+ DataFrame.Internal.AggKernelDirect: directReduce :: Reduction -> Vector Int -> Int -> Column -> Maybe Column
+ DataFrame.Internal.AggKernelDirect: directThreshold :: Int
+ DataFrame.Internal.AggKernelPar: momentScatterPar :: Vector Int -> Vector Int -> Int -> Column -> Column -> Maybe Moments
+ DataFrame.Internal.AggKernelPar: scatterReducePar :: Reduction -> Vector Int -> Vector Int -> Int -> Column -> Maybe Column
+ DataFrame.Internal.AggPlan: MomentPlan :: Text -> Text -> Text -> Text -> Text -> Text -> Text -> Text -> MomentPlan
+ DataFrame.Internal.AggPlan: Moments :: Column -> Column -> Column -> Column -> Column -> Column -> Moments
+ DataFrame.Internal.AggPlan: PlanMaxMinusMin :: Text -> Text -> AggPlan
+ DataFrame.Internal.AggPlan: PlanMedian :: Text -> AggPlan
+ DataFrame.Internal.AggPlan: PlanScatter :: Reduction -> Text -> AggPlan
+ DataFrame.Internal.AggPlan: [mN] :: Moments -> Column
+ DataFrame.Internal.AggPlan: [mSx] :: Moments -> Column
+ DataFrame.Internal.AggPlan: [mSxx] :: Moments -> Column
+ DataFrame.Internal.AggPlan: [mSxy] :: Moments -> Column
+ DataFrame.Internal.AggPlan: [mSy] :: Moments -> Column
+ DataFrame.Internal.AggPlan: [mSyy] :: Moments -> Column
+ DataFrame.Internal.AggPlan: [mpColX] :: MomentPlan -> Text
+ DataFrame.Internal.AggPlan: [mpColY] :: MomentPlan -> Text
+ DataFrame.Internal.AggPlan: [mpNName] :: MomentPlan -> Text
+ DataFrame.Internal.AggPlan: [mpSxName] :: MomentPlan -> Text
+ DataFrame.Internal.AggPlan: [mpSxxName] :: MomentPlan -> Text
+ DataFrame.Internal.AggPlan: [mpSxyName] :: MomentPlan -> Text
+ DataFrame.Internal.AggPlan: [mpSyName] :: MomentPlan -> Text
+ DataFrame.Internal.AggPlan: [mpSyyName] :: MomentPlan -> Text
+ DataFrame.Internal.AggPlan: data AggPlan
+ DataFrame.Internal.AggPlan: data MomentPlan
+ DataFrame.Internal.AggPlan: data Moments
+ DataFrame.Internal.AggPlan: instance GHC.Classes.Eq DataFrame.Internal.AggPlan.Role
+ DataFrame.Internal.AggPlan: instance GHC.Classes.Eq DataFrame.Internal.AggPlan.Term
+ DataFrame.Internal.AggPlan: instance GHC.Classes.Ord DataFrame.Internal.AggPlan.Role
+ DataFrame.Internal.AggPlan: instance GHC.Classes.Ord DataFrame.Internal.AggPlan.Term
+ DataFrame.Internal.AggPlan: instance GHC.Show.Show DataFrame.Internal.AggPlan.Role
+ DataFrame.Internal.AggPlan: instance GHC.Show.Show DataFrame.Internal.AggPlan.Term
+ DataFrame.Internal.AggPlan: momentScatter :: Vector Int -> Int -> Column -> Column -> Maybe Moments
+ DataFrame.Internal.AggPlan: planAgg :: GroupedDataFrame -> UExpr -> Maybe AggPlan
+ DataFrame.Internal.AggPlan: planMoments :: GroupedDataFrame -> [(Text, UExpr)] -> Maybe MomentPlan
+ DataFrame.Internal.Column: [PackedText] :: Maybe Bitmap -> {-# UNPACK #-} !PackedTextData -> Column
+ DataFrame.Internal.Column: eqPackedCols :: Maybe Bitmap -> PackedTextData -> Maybe Bitmap -> PackedTextData -> Bool
+ DataFrame.Internal.Column: isPackedText :: Column -> Bool
+ DataFrame.Internal.Column: materializePacked :: Column -> Column
+ DataFrame.Internal.ColumnBuilder: TextChunk :: !Array -> !Int -> !Vector Int -> !Maybe Bitmap -> TextChunk
+ DataFrame.Internal.ColumnBuilder: [tcBitmap] :: TextChunk -> !Maybe Bitmap
+ DataFrame.Internal.ColumnBuilder: [tcBytes] :: TextChunk -> !Array
+ DataFrame.Internal.ColumnBuilder: [tcOffsets] :: TextChunk -> !Vector Int
+ DataFrame.Internal.ColumnBuilder: [tcUsed] :: TextChunk -> !Int
+ DataFrame.Internal.ColumnBuilder: appendDouble :: DoubleBuilder s -> Double -> ST s ()
+ DataFrame.Internal.ColumnBuilder: appendInt :: IntBuilder s -> Int -> ST s ()
+ DataFrame.Internal.ColumnBuilder: appendNull :: ColumnBuilder b => b s -> ST s ()
+ DataFrame.Internal.ColumnBuilder: appendNum :: Unbox a => NumBuilder a s -> a -> ST s ()
+ DataFrame.Internal.ColumnBuilder: appendText :: TextBuilder s -> Text -> ST s ()
+ DataFrame.Internal.ColumnBuilder: appendTextSlice :: TextBuilder s -> Array -> Int -> Int -> ST s ()
+ DataFrame.Internal.ColumnBuilder: appendTextSliceFromPtr :: TextBuilder s -> Ptr Word8 -> Int -> ST s ()
+ DataFrame.Internal.ColumnBuilder: builderLength :: ColumnBuilder b => b s -> ST s Int
+ DataFrame.Internal.ColumnBuilder: class ColumnBuilder (b :: Type -> Type)
+ DataFrame.Internal.ColumnBuilder: data NumBuilder a s
+ DataFrame.Internal.ColumnBuilder: data TextBuilder s
+ DataFrame.Internal.ColumnBuilder: data TextChunk
+ DataFrame.Internal.ColumnBuilder: freezeBuilder :: ColumnBuilder b => b s -> ST s Column
+ DataFrame.Internal.ColumnBuilder: freezeTextChunk :: TextBuilder s -> ST s TextChunk
+ DataFrame.Internal.ColumnBuilder: instance (DataFrame.Internal.Column.Columnable a, Data.Vector.Unboxed.Base.Unbox a) => DataFrame.Internal.ColumnBuilder.ColumnBuilder (DataFrame.Internal.ColumnBuilder.NumBuilder a)
+ DataFrame.Internal.ColumnBuilder: instance DataFrame.Internal.ColumnBuilder.ColumnBuilder DataFrame.Internal.ColumnBuilder.TextBuilder
+ DataFrame.Internal.ColumnBuilder: mergeColumns :: [Column] -> Column
+ DataFrame.Internal.ColumnBuilder: mergeTextChunks :: [TextChunk] -> Column
+ DataFrame.Internal.ColumnBuilder: newDoubleBuilder :: Int -> ST s (DoubleBuilder s)
+ DataFrame.Internal.ColumnBuilder: newIntBuilder :: Int -> ST s (IntBuilder s)
+ DataFrame.Internal.ColumnBuilder: newNumBuilder :: Unbox a => a -> Int -> ST s (NumBuilder a s)
+ DataFrame.Internal.ColumnBuilder: newTextBuilder :: Int -> Int -> ST s (TextBuilder s)
+ DataFrame.Internal.ColumnBuilder: type DoubleBuilder = NumBuilder Double
+ DataFrame.Internal.ColumnBuilder: type IntBuilder = NumBuilder Int
+ DataFrame.Internal.ColumnMerge: TextChunk :: !Array -> !Int -> !Vector Int -> !Maybe Bitmap -> TextChunk
+ DataFrame.Internal.ColumnMerge: [tcBitmap] :: TextChunk -> !Maybe Bitmap
+ DataFrame.Internal.ColumnMerge: [tcBytes] :: TextChunk -> !Array
+ DataFrame.Internal.ColumnMerge: [tcOffsets] :: TextChunk -> !Vector Int
+ DataFrame.Internal.ColumnMerge: [tcUsed] :: TextChunk -> !Int
+ DataFrame.Internal.ColumnMerge: data TextChunk
+ DataFrame.Internal.ColumnMerge: mergeColumns :: [Column] -> Column
+ DataFrame.Internal.ColumnMerge: mergeTextChunks :: [TextChunk] -> Column
+ DataFrame.Internal.ColumnMerge: packValidity :: Int -> MVector s Word8 -> ST s Bitmap
+ DataFrame.Internal.ColumnMerge: packedFromTextChunk :: TextChunk -> Column
+ DataFrame.Internal.ColumnMerge: spliceBitmaps :: [(Maybe Bitmap, Int)] -> Maybe Bitmap
+ DataFrame.Internal.ColumnMerge: tcRows :: TextChunk -> Int
+ DataFrame.Internal.DictEncode: dictEncodeColumn :: Column -> Maybe (Vector Int, Int)
+ DataFrame.Internal.DictEncode: dictEncodeColumnUpTo :: Int -> Column -> Maybe (Vector Int, Int)
+ DataFrame.Internal.DictEncode: dictMaxCardinality :: Int
+ DataFrame.Internal.Expression: substituteColumns :: Columnable a => Map Text UExpr -> Expr a -> Expr a
+ DataFrame.Internal.Grouping: groupByPar :: [Text] -> DataFrame -> GroupedDataFrame
+ DataFrame.Internal.Grouping: groupBySeq :: [Text] -> DataFrame -> GroupedDataFrame
+ DataFrame.Internal.GroupingDirect: DirectGrouping :: !Vector Int -> !Vector Int -> !Vector Int -> !Int -> DirectGrouping
+ DataFrame.Internal.GroupingDirect: [dgNGroups] :: DirectGrouping -> !Int
+ DataFrame.Internal.GroupingDirect: [dgOffsets] :: DirectGrouping -> !Vector Int
+ DataFrame.Internal.GroupingDirect: [dgRowToGroup] :: DirectGrouping -> !Vector Int
+ DataFrame.Internal.GroupingDirect: [dgValueIndices] :: DirectGrouping -> !Vector Int
+ DataFrame.Internal.GroupingDirect: data DirectGrouping
+ DataFrame.Internal.GroupingDirect: directGroupThreshold :: Int
+ DataFrame.Internal.GroupingDirect: tryDirectGroupColumn :: Column -> Maybe DirectGrouping
+ DataFrame.Internal.GroupingPar: numPartitionsFor :: Int -> Int
+ DataFrame.Internal.GroupingPar: parThreshold :: Int
+ DataFrame.Internal.GroupingPar: parallelAssignGroups :: Int -> Vector Int -> (Int -> Int -> Bool) -> IO (Vector Int, Vector Int, Vector Int)
+ DataFrame.Internal.GroupingPar: shouldParallelize :: Int -> Bool
+ DataFrame.Internal.Hash: mixBytes :: Int -> Array -> Int -> Int -> Int
+ DataFrame.Internal.HashTable: HashTable :: !MVector s Int -> !MVector s Int -> !MVector s Int -> !Int -> HashTable s
+ DataFrame.Internal.HashTable: [htGroup] :: HashTable s -> !MVector s Int
+ DataFrame.Internal.HashTable: [htHash] :: HashTable s -> !MVector s Int
+ DataFrame.Internal.HashTable: [htMask] :: HashTable s -> !Int
+ DataFrame.Internal.HashTable: [htRep] :: HashTable s -> !MVector s Int
+ DataFrame.Internal.HashTable: data HashTable s
+ DataFrame.Internal.HashTable: htInsert :: PrimMonad m => HashTable (PrimState m) -> (Int -> Int -> Bool) -> Int -> Int -> Int -> m (Int, Bool)
+ DataFrame.Internal.HashTable: newHashTable :: PrimMonad m => Int -> m (HashTable (PrimState m))
+ DataFrame.Internal.HashTable: nextPow2Above :: Int -> Int
+ DataFrame.Internal.PackedText: PackedTextData :: {-# UNPACK #-} !Array -> {-# UNPACK #-} !Vector Int -> !Maybe (Vector Int) -> PackedTextData
+ DataFrame.Internal.PackedText: [ptBytes] :: PackedTextData -> {-# UNPACK #-} !Array
+ DataFrame.Internal.PackedText: [ptOffsets] :: PackedTextData -> {-# UNPACK #-} !Vector Int
+ DataFrame.Internal.PackedText: [ptSel] :: PackedTextData -> !Maybe (Vector Int)
+ DataFrame.Internal.PackedText: data PackedTextData
+ DataFrame.Internal.PackedText: mkPackedContiguous :: Array -> Vector Int -> PackedTextData
+ DataFrame.Internal.PackedText: packedGather :: Vector Int -> PackedTextData -> PackedTextData
+ DataFrame.Internal.PackedText: packedIndexText :: PackedTextData -> Int -> Text
+ DataFrame.Internal.PackedText: packedLength :: PackedTextData -> Int
+ DataFrame.Internal.PackedText: packedRowOffsetVec :: PackedTextData -> Maybe (Array, Vector Int)
+ DataFrame.Internal.PackedText: packedSlice :: PackedTextData -> Int -> (Array, Int, Int)
+ DataFrame.Internal.PackedText: packedTake :: Int -> PackedTextData -> PackedTextData
+ DataFrame.Internal.PackedText: sliceCmpBytes :: Array -> Int -> Int -> Array -> Int -> Int -> Ordering
+ DataFrame.Internal.PackedText: sliceEqBytes :: Array -> Int -> Int -> Array -> Int -> Int -> Bool
+ DataFrame.Internal.ParRadixSort: parSortByHash :: Int -> Vector Int -> (Vector Int, Vector Int)
+ DataFrame.Internal.ParRadixSort: parSortThreshold :: Int
+ DataFrame.Internal.RadixRank: rankByHash :: PrimMonad m => (Int -> m Int) -> Int -> m (Vector Int)
+ DataFrame.Internal.RadixRank: sortKey :: Int -> Int
+ DataFrame.Internal.RowHash: computeRowHashesIO :: Int -> [Column] -> IO (Vector Int)
+ DataFrame.Internal.RowHash: hashRowRange :: IOVector Int -> Int -> Int -> [Column] -> IO ()
+ DataFrame.Internal.RowHash: parRowHashThreshold :: Int
+ DataFrame.Internal.Utf8: isUtf8Boundary :: Word8 -> Bool
+ DataFrame.Internal.Utf8: isValidUtf8Slice :: Array -> Int -> Int -> Bool
+ DataFrame.Internal.Utf8: lenientDecodeSlice :: Array -> Int -> Int -> Text
+ DataFrame.Internal.Utf8: sliceTextVector :: Array -> Vector Int -> Vector Text

Files

dataframe-core.cabal view
@@ -1,7 +1,6 @@ cabal-version:      2.4 name:               dataframe-core-version:            1.0.2.0-+version:            1.1.0.0 synopsis:           Core data structures for the dataframe library. description:     Minimal interchange-format types for the @dataframe@ ecosystem:@@ -36,16 +35,31 @@                         DataFrame.Operators                         DataFrame.Display.Terminal.Colours                         DataFrame.Display.Terminal.PrettyPrint+                        DataFrame.Internal.AggKernel+                        DataFrame.Internal.AggKernelDirect+                        DataFrame.Internal.AggKernelPar+                        DataFrame.Internal.AggPlan+                        DataFrame.Internal.GroupingDirect                         DataFrame.Internal.Column+                        DataFrame.Internal.ColumnBuilder+                        DataFrame.Internal.ColumnMerge                         DataFrame.Internal.DataFrame+                        DataFrame.Internal.DictEncode                         DataFrame.Internal.Expression                         DataFrame.Internal.Grouping+                        DataFrame.Internal.GroupingPar                         DataFrame.Internal.Hash+                        DataFrame.Internal.HashTable                         DataFrame.Internal.Interpreter                         DataFrame.Internal.Nullable+                        DataFrame.Internal.PackedText+                        DataFrame.Internal.ParRadixSort+                        DataFrame.Internal.RadixRank+                        DataFrame.Internal.RowHash                         DataFrame.Internal.Row                         DataFrame.Internal.Simplify                         DataFrame.Internal.Types+                        DataFrame.Internal.Utf8                         DataFrame.Typed.Freeze                         DataFrame.Typed.Generic                         DataFrame.Typed.Record@@ -54,8 +68,9 @@                         DataFrame.Typed.Util     build-depends:      base >= 4 && < 5,                         containers >= 0.6.7 && < 0.9,+                        primitive >= 0.7 && < 0.10,                         random >= 1 && < 2,-                        text >= 2.0 && < 3,+                        text >= 2.1 && < 3,                         vector ^>= 0.13     hs-source-dirs:     src     default-language:   Haskell2010
+ src/DataFrame/Internal/AggKernel.hs view
@@ -0,0 +1,308 @@+{-# 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 view
@@ -0,0 +1,373 @@+{-# 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 view
@@ -0,0 +1,454 @@+{-# 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 view
@@ -0,0 +1,311 @@+{-# 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 (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) = 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" -> require name (PlanScatter RMean name)+        "count" -> require name (PlanScatter RCount name)+        _ -> 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 view
@@ -42,7 +42,18 @@ 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@@ -60,6 +71,12 @@     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. @@ -187,13 +204,34 @@ 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 -- ---------------------------------------------------------------------------@@ -223,12 +261,14 @@ 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.@@ -239,6 +279,7 @@ 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 _)@),@@ -248,6 +289,7 @@ 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@@ -271,6 +313,8 @@     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.@@ -281,6 +325,8 @@     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 =@@ -304,6 +350,7 @@             | 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@@ -323,6 +370,7 @@                 | 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.@@ -364,8 +412,33 @@                             )                             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. -}@@ -497,6 +570,7 @@ 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.@@ -568,6 +642,7 @@ 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.@@ -596,6 +671,7 @@ 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.@@ -604,6 +680,8 @@ 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.@@ -612,6 +690,8 @@     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.@@ -625,6 +705,7 @@     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.@@ -652,6 +733,18 @@             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.@@ -663,6 +756,22 @@         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@@ -718,6 +827,7 @@ 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.@@ -745,6 +855,7 @@ ifoldrColumn f acc = \case     BoxedColumn _ column -> foldrWorker column     UnboxedColumn _ column -> foldrWorker column+    c@(PackedText _ _) -> ifoldrColumn f acc (materializePacked c)   where     foldrWorker ::         forall c v.@@ -771,6 +882,7 @@ foldlColumn f acc = \case     BoxedColumn _ column -> foldlWorker column     UnboxedColumn _ column -> foldlWorker column+    c@(PackedText _ _) -> foldlColumn f acc (materializePacked c)   where     foldlWorker ::         forall c v.@@ -797,6 +909,7 @@ foldl1Column f = \case     BoxedColumn _ column -> foldl1Worker column     UnboxedColumn _ column -> foldl1Worker column+    c@(PackedText _ _) -> foldl1Column f (materializePacked c)   where     foldl1Worker ::         forall c v.@@ -832,6 +945,7 @@     | 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.@@ -885,6 +999,8 @@     | 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.@@ -933,6 +1049,7 @@ headColumn = \case     BoxedColumn _ col -> headWorker col     UnboxedColumn _ col -> headWorker col+    c@(PackedText _ _) -> headColumn (materializePacked c)   where     headWorker ::         forall c v.@@ -957,6 +1074,8 @@  -- | 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@@ -978,6 +1097,8 @@ -- | 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 ->@@ -1096,9 +1217,15 @@ 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 =@@ -1128,6 +1255,9 @@  -- | 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 =@@ -1162,6 +1292,8 @@ -} 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@@ -1220,6 +1352,9 @@ 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@@ -1263,8 +1398,11 @@                              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@@ -1325,6 +1463,7 @@     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 ()@@ -1336,6 +1475,8 @@     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" @@ -1394,6 +1535,7 @@     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@@ -1464,6 +1606,7 @@ 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@@ -1546,6 +1689,7 @@ 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@@ -1629,6 +1773,7 @@ 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
+ src/DataFrame/Internal/ColumnBuilder.hs view
@@ -0,0 +1,299 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE UndecidableInstances #-}++{- | Mutable, growable column builders for high-throughput ingest. No+per-append @IORef@ traffic: hot counters live in an unboxed vector, payloads+double on demand, and validity is only materialized once a null is seen.+-}+module DataFrame.Internal.ColumnBuilder (+    ColumnBuilder (..),+    NumBuilder,+    IntBuilder,+    DoubleBuilder,+    TextBuilder,+    TextChunk (..),+    newIntBuilder,+    newDoubleBuilder,+    newNumBuilder,+    newTextBuilder,+    appendInt,+    appendDouble,+    appendNum,+    appendText,+    appendTextSlice,+    appendTextSliceFromPtr,+    freezeTextChunk,+    mergeColumns,+    mergeTextChunks,+) where++import qualified Data.Text as T+import qualified Data.Text.Array as A+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM++import Control.Monad (when)+import Control.Monad.ST (ST)+import Data.Bits (shiftR)+import Data.STRef+import Data.Text.Internal (Text (..))+import Data.Word (Word8)+import DataFrame.Internal.Column hiding (mergeColumns)+import DataFrame.Internal.ColumnMerge (+    TextChunk (..),+    mergeColumns,+    mergeTextChunks,+    packValidity,+ )+import Foreign.Ptr (Ptr)++{- | Operations shared by all column builders. A builder must not be used+again after 'freezeBuilder' (its storage is frozen in place, not copied).+-}+class ColumnBuilder b where+    -- | Append a null row (sentinel payload + invalid bit).+    appendNull :: b s -> ST s ()++    -- | Rows appended so far.+    builderLength :: b s -> ST s Int++    -- | Freeze into a fully-forced 'Column'; bitmap only when a null was seen.+    freezeBuilder :: b s -> ST s Column++-- Counter slots shared by the builders: rows, any-null flag, text bytes used.+cRows, cAnyNull, cBytes :: Int+cRows = 0+cAnyNull = 1+cBytes = 2++{- | Builder for unboxed numeric payloads ('Int', 'Double', ...). 'nbNull'+is the sentinel written into null slots (protected by the bitmap).+-}+data NumBuilder a s = NumBuilder+    { nbNull :: !a+    , nbCounters :: !(VUM.MVector s Int)+    , nbArrays :: !(STRef s (NumArrays a s))+    }++data NumArrays a s = NumArrays+    { naData :: !(VUM.MVector s a)+    , naValid :: !(VUM.MVector s Word8)+    }++type IntBuilder = NumBuilder Int++type DoubleBuilder = NumBuilder Double++-- | New numeric builder with a row-capacity hint and a null sentinel.+newNumBuilder :: (VU.Unbox a) => a -> Int -> ST s (NumBuilder a s)+newNumBuilder nullValue hint = do+    let cap = max 16 hint+    counters <- VUM.replicate 2 0+    dat <- VUM.unsafeNew cap+    val <- VUM.unsafeNew cap+    NumBuilder nullValue counters <$> newSTRef (NumArrays dat val)++newIntBuilder :: Int -> ST s (IntBuilder s)+newIntBuilder = newNumBuilder 0++newDoubleBuilder :: Int -> ST s (DoubleBuilder s)+newDoubleBuilder = newNumBuilder 0++appendNum :: (VU.Unbox a) => NumBuilder a s -> a -> ST s ()+appendNum b !x = do+    n <- VUM.unsafeRead (nbCounters b) cRows+    anyNull <- VUM.unsafeRead (nbCounters b) cAnyNull+    NumArrays dat val <- reserveNum b n+    VUM.unsafeWrite dat n x+    when (anyNull /= 0) $ VUM.unsafeWrite val n 1+    VUM.unsafeWrite (nbCounters b) cRows (n + 1)+{-# INLINE appendNum #-}++appendInt :: IntBuilder s -> Int -> ST s ()+appendInt = appendNum+{-# INLINE appendInt #-}++appendDouble :: DoubleBuilder s -> Double -> ST s ()+appendDouble = appendNum+{-# INLINE appendDouble #-}++-- Fetch the arrays, growing (doubling) first if row @n@ would not fit.+reserveNum :: (VU.Unbox a) => NumBuilder a s -> Int -> ST s (NumArrays a s)+reserveNum b n = do+    arrs <- readSTRef (nbArrays b)+    if n < VUM.length (naData arrs) then pure arrs else growNum b arrs+{-# INLINE reserveNum #-}++growNum ::+    (VU.Unbox a) => NumBuilder a s -> NumArrays a s -> ST s (NumArrays a s)+growNum b (NumArrays dat val) = do+    let cap = VUM.length dat+    dat' <- VUM.unsafeGrow dat cap+    val' <- VUM.unsafeGrow val cap+    let arrs = NumArrays dat' val'+    writeSTRef (nbArrays b) arrs+    pure arrs++instance (Columnable a, VU.Unbox a) => ColumnBuilder (NumBuilder a) where+    appendNull b = do+        n <- VUM.unsafeRead (nbCounters b) cRows+        anyNull <- VUM.unsafeRead (nbCounters b) cAnyNull+        NumArrays dat val <- reserveNum b n+        VUM.unsafeWrite dat n (nbNull b)+        when (anyNull == 0) $ do+            VUM.set (VUM.slice 0 n val) 1+            VUM.unsafeWrite (nbCounters b) cAnyNull 1+        VUM.unsafeWrite val n 0+        VUM.unsafeWrite (nbCounters b) cRows (n + 1)+    {-# INLINE appendNull #-}++    builderLength b = VUM.unsafeRead (nbCounters b) cRows++    freezeBuilder b = do+        n <- VUM.unsafeRead (nbCounters b) cRows+        anyNull <- VUM.unsafeRead (nbCounters b) cAnyNull+        NumArrays dat val <- readSTRef (nbArrays b)+        !vs <- freezeTrimmed n dat+        if anyNull /= 0+            then do+                !bm <- packValidity n val+                pure $! UnboxedColumn (Just bm) vs+            else pure $! UnboxedColumn Nothing vs++-- Zero-copy freeze; copies to exact size when slack exceeds a quarter of n.+freezeTrimmed :: (VU.Unbox a) => Int -> VUM.MVector s a -> ST s (VU.Vector a)+freezeTrimmed n mv+    | VUM.length mv - n <= n `shiftR` 2 = VU.unsafeFreeze (VUM.slice 0 n mv)+    | otherwise = VU.freeze (VUM.slice 0 n mv)++{- | Builder for 'Text' columns. All field bytes go into one exponentially+grown byte array; rows are recorded as offsets, so an append is a memcpy+and freezing slices 'Text' values off the shared array without copying.+-}+data TextBuilder s = TextBuilder+    { tbCounters :: !(VUM.MVector s Int)+    , tbArrays :: !(STRef s (TextArrays s))+    }++data TextArrays s = TextArrays+    { taBytes :: !(A.MArray s)+    , taByteCap :: !Int+    , taOffsets :: !(VUM.MVector s Int)+    -- ^ Row @i@ spans bytes @[offsets!i, offsets!(i+1))@.+    , taValid :: !(VUM.MVector s Word8)+    }++-- | New text builder with row-count and total-byte capacity hints.+newTextBuilder :: Int -> Int -> ST s (TextBuilder s)+newTextBuilder rowHint byteHint = do+    let rcap = max 16 rowHint+        bcap = max 64 byteHint+    counters <- VUM.replicate 3 0+    bytes <- A.new bcap+    offsets <- VUM.unsafeNew (rcap + 1)+    VUM.unsafeWrite offsets 0 0+    val <- VUM.unsafeNew rcap+    TextBuilder counters <$> newSTRef (TextArrays bytes bcap offsets val)++-- | Append @len@ raw bytes at @off@ in @src@ as one field (one memcpy).+appendTextSlice :: TextBuilder s -> A.Array -> Int -> Int -> ST s ()+appendTextSlice b src off len = do+    (n, pos, arrs) <- reserveText b len+    A.copyI len (taBytes arrs) pos src off+    finishTextAppend b arrs n (pos + len)+{-# INLINE appendTextSlice #-}++-- | 'appendTextSlice' from foreign memory (e.g. an mmapped file buffer).+appendTextSliceFromPtr :: TextBuilder s -> Ptr Word8 -> Int -> ST s ()+appendTextSliceFromPtr b ptr len = do+    (n, pos, arrs) <- reserveText b len+    A.copyFromPointer (taBytes arrs) pos ptr len+    finishTextAppend b arrs n (pos + len)+{-# INLINE appendTextSliceFromPtr #-}++-- | Append an already-decoded 'Text' (its bytes are UTF-8 already).+appendText :: TextBuilder s -> T.Text -> ST s ()+appendText b (Text src off len) = appendTextSlice b src off len+{-# INLINE appendText #-}++finishTextAppend :: TextBuilder s -> TextArrays s -> Int -> Int -> ST s ()+finishTextAppend b arrs n endPos = do+    anyNull <- VUM.unsafeRead (tbCounters b) cAnyNull+    when (anyNull /= 0) $ VUM.unsafeWrite (taValid arrs) n 1+    VUM.unsafeWrite (taOffsets arrs) (n + 1) endPos+    VUM.unsafeWrite (tbCounters b) cRows (n + 1)+    VUM.unsafeWrite (tbCounters b) cBytes endPos+{-# INLINE finishTextAppend #-}++reserveText :: TextBuilder s -> Int -> ST s (Int, Int, TextArrays s)+reserveText b extra = do+    n <- VUM.unsafeRead (tbCounters b) cRows+    pos <- VUM.unsafeRead (tbCounters b) cBytes+    arrs <- readSTRef (tbArrays b)+    arrs' <-+        if n < VUM.length (taValid arrs) && pos + extra <= taByteCap arrs+            then pure arrs+            else growText b arrs (n + 1) (pos + extra)+    pure (n, pos, arrs')+{-# INLINE reserveText #-}++growText :: TextBuilder s -> TextArrays s -> Int -> Int -> ST s (TextArrays s)+growText b (TextArrays bytes bcap offsets val) needRows needBytes = do+    let rcap = VUM.length val+    (offsets', val') <-+        if needRows > rcap+            then do+                let rcap' = max (2 * rcap) needRows+                o <- VUM.unsafeGrow offsets (rcap' - rcap)+                v <- VUM.unsafeGrow val (rcap' - rcap)+                pure (o, v)+            else pure (offsets, val)+    (bytes', bcap') <-+        if needBytes > bcap+            then do+                let cap' = max (2 * bcap) needBytes+                bs <- A.resizeM bytes cap'+                pure (bs, cap')+            else pure (bytes, bcap)+    let arrs = TextArrays bytes' bcap' offsets' val'+    writeSTRef (tbArrays b) arrs+    pure arrs++{- | Freeze a 'TextBuilder' into a raw 'TextChunk' for byte-level merging+('mergeTextChunks'): no 'T.Text' values are created until chunks merge.+-}+freezeTextChunk :: TextBuilder s -> ST s TextChunk+freezeTextChunk b = do+    n <- VUM.unsafeRead (tbCounters b) cRows+    anyNull <- VUM.unsafeRead (tbCounters b) cAnyNull+    used <- VUM.unsafeRead (tbCounters b) cBytes+    TextArrays bytes bcap offsets val <- readSTRef (tbArrays b)+    when (used < bcap) (A.shrinkM bytes used)+    arr <- A.unsafeFreeze bytes+    offs <- VU.unsafeFreeze (VUM.slice 0 (n + 1) offsets)+    bm <-+        if anyNull /= 0+            then Just <$> packValidity n val+            else pure Nothing+    pure (TextChunk arr used offs bm)++instance ColumnBuilder TextBuilder where+    appendNull b = do+        (n, pos, arrs) <- reserveText b 0+        anyNull <- VUM.unsafeRead (tbCounters b) cAnyNull+        when (anyNull == 0) $ do+            VUM.set (VUM.slice 0 n (taValid arrs)) 1+            VUM.unsafeWrite (tbCounters b) cAnyNull 1+        VUM.unsafeWrite (taValid arrs) n 0+        VUM.unsafeWrite (taOffsets arrs) (n + 1) pos+        VUM.unsafeWrite (tbCounters b) cRows (n + 1)+        VUM.unsafeWrite (tbCounters b) cBytes pos+    {-# INLINE appendNull #-}++    builderLength b = VUM.unsafeRead (tbCounters b) cRows++    freezeBuilder b = do+        chunk <- freezeTextChunk b+        pure $! mergeTextChunks [chunk]
+ src/DataFrame/Internal/ColumnMerge.hs view
@@ -0,0 +1,179 @@+{-# 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 view
@@ -28,6 +28,7 @@ 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)@@ -181,6 +182,8 @@         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.@@ -203,6 +206,7 @@                     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@@ -366,6 +370,9 @@ 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 view
@@ -0,0 +1,159 @@+{-# 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 view
@@ -15,6 +15,7 @@  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))@@ -369,6 +370,44 @@         (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
src/DataFrame/Internal/Grouping.hs view
@@ -9,11 +9,12 @@  module DataFrame.Internal.Grouping (     groupBy,+    groupBySeq,+    groupByPar,     buildRowToGroup,     changingPoints, ) where -import qualified Data.IntMap.Strict as IM import qualified Data.List as L import qualified Data.Map as M import qualified Data.Text as T@@ -32,12 +33,35 @@     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] ->@@ -57,83 +81,295 @@             VU.empty             (VU.fromList [0])             VU.empty-    | otherwise =-        let !vis = VU.map fst valIndices-            !os = changingPoints valIndices-            !n = nRows df-         in Grouped-                df-                names-                vis-                os-                (buildRowToGroup n vis os)+    | Just dg <- tryDirectGroup names df = dg+    | shouldParallelize n = groupByPar names df+    | otherwise = groupBySeq names df   where-    indicesToGroup = M.elems $ M.filterWithKey (\k _ -> k `elem` names) (columnIndices df)-    -- valIndices is a vector of (rowIdx, rowHash) sorted by rowHash so that-    -- runs of equal-hash rows are adjacent. Hashes are computed with the-    -- in-tree FNV combinators from "DataFrame.Internal.Hash", and bucketed-    -- with 'Data.IntMap.Strict' to keep the dep set minimal (no hashable,-    -- no vector-algorithms).-    valIndices = runST $ do-        let n = nRows df-        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-        hashes <- VU.unsafeFreeze mh-        -- Bucket row indices by hash using an IntMap, then walk it in-        -- ascending key order to emit (rowIdx, hash) pairs grouped by-        -- hash. Each bucket's accumulated list is reversed so rows come-        -- out in the original row order.-        let buckets =-                VU.ifoldl'-                    (\acc i h -> IM.insertWith (++) h [i] acc)-                    IM.empty-                    hashes-            ordered =-                [ (i, h)-                | (h, is) <- IM.toAscList buckets-                , i <- reverse is-                ]-        return (VU.fromList ordered)+    !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. -}@@ -163,6 +399,26 @@             )             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 
+ src/DataFrame/Internal/GroupingDirect.hs view
@@ -0,0 +1,260 @@+{-# 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 view
@@ -0,0 +1,355 @@+{-# 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 view
@@ -19,12 +19,14 @@     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@@ -90,8 +92,17 @@ bytes are folded in individually. -} mixText :: Int -> T.Text -> Int-mixText !acc (Text (ByteArray ba) off len) = goBytes (goWords acc off) wordsEnd+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@@ -107,7 +118,7 @@             let !(I# i#) = i                 !b = fromIntegral (W8# (indexWord8Array# ba i#)) :: Int              in goBytes (mixInt h b) (i + 1)-{-# INLINE mixText #-}+{-# INLINE mixBytes #-}  {- | Fallback for arbitrary 'Show'-able values. Slower but covers types without a dedicated combinator (e.g. 'Day', 'UTCTime').
+ src/DataFrame/Internal/HashTable.hs view
@@ -0,0 +1,113 @@+{-# 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 view
@@ -473,6 +473,7 @@ -} 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@@ -578,6 +579,7 @@                             )                 SFalse -> castMismatch @c @b     BoxedColumn _ _ -> tryParseWith @Double onResult col+    PackedText _ _ -> promoteToDoubleWith onResult (materializePacked col)  promoteToFloatWith ::     forall b.@@ -617,6 +619,7 @@                             )                 SFalse -> castMismatch @c @b     BoxedColumn _ _ -> tryParseWith @Float onResult col+    PackedText _ _ -> promoteToFloatWith onResult (materializePacked col)  promoteToIntWith ::     forall b.@@ -656,6 +659,7 @@                             )                 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@@ -670,6 +674,7 @@     (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
+ src/DataFrame/Internal/PackedText.hs view
@@ -0,0 +1,169 @@+{-# 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 view
@@ -0,0 +1,287 @@+{-# 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 !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 view
@@ -0,0 +1,106 @@+{-# LANGUAGE BangPatterns #-}++{- |+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 ::+    (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 !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 view
@@ -24,6 +24,7 @@ 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@@ -177,6 +178,7 @@                     (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@@ -200,5 +202,8 @@         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 view
@@ -0,0 +1,232 @@+{-# 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/Utf8.hs view
@@ -0,0 +1,98 @@+{-# 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/Typed/Schema.hs view
@@ -31,6 +31,7 @@     AssertAbsent,     AssertPresent,     AssertAllPresent,+    AssertKeyTypesMatch,     IsElem,      -- * Maybe-stripping families@@ -222,6 +223,40 @@     AssertAllPresentHelper 'False name rest cols =         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.+-}+type family+    AssertKeyTypesMatch (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) ::+        Constraint+    where+    AssertKeyTypesMatch '[] left right = ()+    AssertKeyTypesMatch (k ': ks) left right =+        ( KeyTypeMatchHelper k (SafeLookup k left) (SafeLookup k right)+        , AssertKeyTypesMatch ks left right+        )++type family+    KeyTypeMatchHelper (k :: Symbol) (l :: Type) (r :: Type) ::+        Constraint+    where+    KeyTypeMatchHelper k a a = ()+    KeyTypeMatchHelper k (Maybe a) a = ()+    KeyTypeMatchHelper k a (Maybe a) = ()+    KeyTypeMatchHelper k l r =+        TypeError+            ( 'Text "Join key '"+                ':<>: 'Text k+                ':<>: 'Text "' has type "+                ':<>: 'ShowType l+                ':<>: 'Text " in the left table but "+                ':<>: 'ShowType r+                ':<>: 'Text " in the right table"+            )  {- | Strip 'Maybe' from all columns. Used by 'filterAllJust'.