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 +18/−3
- src/DataFrame/Internal/AggKernel.hs +308/−0
- src/DataFrame/Internal/AggKernelDirect.hs +373/−0
- src/DataFrame/Internal/AggKernelPar.hs +454/−0
- src/DataFrame/Internal/AggPlan.hs +311/−0
- src/DataFrame/Internal/Column.hs +145/−0
- src/DataFrame/Internal/ColumnBuilder.hs +299/−0
- src/DataFrame/Internal/ColumnMerge.hs +179/−0
- src/DataFrame/Internal/DataFrame.hs +7/−0
- src/DataFrame/Internal/DictEncode.hs +159/−0
- src/DataFrame/Internal/Expression.hs +39/−0
- src/DataFrame/Internal/Grouping.hs +332/−76
- src/DataFrame/Internal/GroupingDirect.hs +260/−0
- src/DataFrame/Internal/GroupingPar.hs +355/−0
- src/DataFrame/Internal/Hash.hs +13/−2
- src/DataFrame/Internal/HashTable.hs +113/−0
- src/DataFrame/Internal/Interpreter.hs +5/−0
- src/DataFrame/Internal/PackedText.hs +169/−0
- src/DataFrame/Internal/ParRadixSort.hs +287/−0
- src/DataFrame/Internal/RadixRank.hs +106/−0
- src/DataFrame/Internal/Row.hs +5/−0
- src/DataFrame/Internal/RowHash.hs +232/−0
- src/DataFrame/Internal/Utf8.hs +98/−0
- src/DataFrame/Typed/Schema.hs +35/−0
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'.