dataframe-core 2.5.0.2 → 2.6.0.0

raw patch · 48 files changed

+2642/−2840 lines, 48 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• dataframe-core.cabal +4/−1
• src-internal/DataFrame/Display/Terminal/PrettyPrint.hs +5/−15
• src-internal/DataFrame/Internal/Aggregation/Kernel/Dense.hs +16/−276
• src-internal/DataFrame/Internal/Aggregation/Kernel/Distinct.hs +609/−0
• src-internal/DataFrame/Internal/Aggregation/Kernel/Fused.hs +13/−178
• src-internal/DataFrame/Internal/Aggregation/Kernel/Moments.hs +7/−46
• src-internal/DataFrame/Internal/Aggregation/Kernel/Scatter.hs +11/−230
• src-internal/DataFrame/Internal/Aggregation/Plan.hs +27/−46
• src-internal/DataFrame/Internal/Aggregation/Reduction.hs +2/−13
• src-internal/DataFrame/Internal/Algorithms/Hash.hs +8/−30
• src-internal/DataFrame/Internal/Algorithms/Rank/Radix.hs +2/−8
• src-internal/DataFrame/Internal/Algorithms/Sort/Radix/Parallel.hs +12/−64
• src-internal/DataFrame/Internal/Column.hs +1/−14
• src-internal/DataFrame/Internal/Column/Base.hs +14/−50
• src-internal/DataFrame/Internal/Column/Bitmap.hs +13/−22
• src-internal/DataFrame/Internal/Column/Builder.hs +9/−19
• src-internal/DataFrame/Internal/Column/Conversion.hs +64/−88
• src-internal/DataFrame/Internal/Column/Encode.hs +7/−28
• src-internal/DataFrame/Internal/Column/Merge.hs +70/−25
• src-internal/DataFrame/Internal/Column/Operations.hs +408/−251
• src-internal/DataFrame/Internal/Column/Properties.hs +3/−10
• src-internal/DataFrame/Internal/Column/Types.hs +4/−13
• src-internal/DataFrame/Internal/Control/Concurrent.hs +3/−12
• src-internal/DataFrame/Internal/Data/HashTable.hs +6/−24
• src-internal/DataFrame/Internal/Data/PackedText.hs +16/−73
• src-internal/DataFrame/Internal/Data/PackedText/Utf8.hs +4/−15
• src-internal/DataFrame/Internal/DataFrame.hs +25/−68
• src-internal/DataFrame/Internal/Display/Pretty.hs +3/−11
• src-internal/DataFrame/Internal/Expression.hs +17/−29
• src-internal/DataFrame/Internal/Expression/Operators.hs +16/−16
• src-internal/DataFrame/Internal/Expression/Operators/Nullable.hs +18/−107
• src-internal/DataFrame/Internal/Expression/Reflect.hs +111/−0
• src-internal/DataFrame/Internal/Expression/Simplify.hs +7/−27
• src-internal/DataFrame/Internal/Grouping.hs +83/−247
• src-internal/DataFrame/Internal/Grouping/Direct.hs +31/−170
• src-internal/DataFrame/Internal/Grouping/Partitioned.hs +119/−114
• src-internal/DataFrame/Internal/Interpreter.hs +537/−204
• src-internal/DataFrame/Internal/Interpreter/Kernels.hs +270/−0
• src-internal/DataFrame/Internal/Row.hs +7/−30
• src-internal/DataFrame/Internal/Row/RowHash.hs +8/−41
• src/DataFrame/Core.hs +9/−10
• src/DataFrame/Expression/Operators.hs +1/−7
• src/DataFrame/Typed/Freeze.hs +7/−20
• src/DataFrame/Typed/Generic.hs +7/−56
• src/DataFrame/Typed/Record.hs +11/−48
• src/DataFrame/Typed/Schema.hs +9/−41
• src/DataFrame/Typed/Types.hs +7/−34
• src/DataFrame/Typed/Util.hs +1/−9

@@ -1,6 +1,6 @@ cabal-version:      3.4 name:               dataframe-core-version:            2.5.0.2+version:            2.6.0.0 synopsis:           Core data structures for the dataframe library. description:     Minimal interchange-format types for the @dataframe@ ecosystem:@@ -46,6 +46,7 @@                         DataFrame.Internal.Control.Concurrent                         -- Aggregation kernels and planning.                         DataFrame.Internal.Aggregation.Kernel.Dense+                        DataFrame.Internal.Aggregation.Kernel.Distinct                         DataFrame.Internal.Aggregation.Kernel.Fused                         DataFrame.Internal.Aggregation.Kernel.Moments                         DataFrame.Internal.Aggregation.Kernel.Scatter@@ -75,12 +76,14 @@                         -- Expressions: syntax, operators, simplification.                         DataFrame.Internal.Expression                         DataFrame.Internal.Expression.Operators+                        DataFrame.Internal.Expression.Reflect                         DataFrame.Internal.Expression.Operators.Nullable                         DataFrame.Internal.Expression.Simplify                         DataFrame.Internal.Grouping                         DataFrame.Internal.Grouping.Direct                         DataFrame.Internal.Grouping.Partitioned                         DataFrame.Internal.Interpreter+                        DataFrame.Internal.Interpreter.Kernels                         DataFrame.Internal.Row                         DataFrame.Internal.Row.RowHash     build-depends:      base >= 4 && < 5,

@@ -5,17 +5,13 @@ import qualified Data.Text as T import qualified Data.Vector as V -{- | Output format for 'showTable'. 'Plain' renders a terminal-style table with-ASCII borders; 'Markdown' renders a GitHub-flavoured pipe table suitable for-notebooks.--}+-- | Output format for 'showTable': 'Plain' or 'Markdown'. data RenderFormat = Plain | Markdown     deriving (Show, Eq)  -- Utility functions to show a DataFrame as a Markdown-ish table. --- Fill functions, adapted from:--- https://stackoverflow.com/questions/5929377/format-list-output-in-haskell+-- Fill functions, adapted from StackOverflow question 5929377. type Filler = Int -> T.Text -> T.Text  -- A description of one table column.@@ -25,7 +21,7 @@     , colValueFill :: Filler     } --- Fill text to a given width with a pad character, aligning left, right, or center.+-- Pad text to a given width with a character, aligned left, right or center. fillLeft :: Char -> Int -> T.Text -> T.Text fillLeft c n s = s <> T.replicate (n - T.length s) (T.singleton c) @@ -50,10 +46,7 @@ center :: Int -> T.Text -> T.Text center = fillCenter ' ' -{- | Render a table from column-major data. @columns@ has one 'V.Vector' per-column; widths are computed in one pass per column (no row-major transpose),-and row lines are built by indexing each column at row @i@.--}+-- | Render a table from column-major data. showTable ::     RenderFormat ->     [T.Text] ->@@ -106,9 +99,6 @@                         : rowLines      in T.unlines outputLines -{- | Escape a value for a GitHub-flavoured Markdown table cell: a bare @|@ would-end the cell and a newline would end the row, so both are neutralised (the pipe-backslash-escaped, the newline turned into a @\<br\>@).--}+-- | Escape a value for a GitHub-flavoured Markdown table cell. escapeMarkdownCell :: T.Text -> T.Text escapeMarkdownCell = T.replace "\n" "<br>" . T.replace "|" "\\|"

@@ -5,25 +5,12 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | The low-cardinality DENSE reduction kernel: when the group domain is small,-the grouping layer's @rowToGroup@ already maps row -> group, so the reduction-scatters straight off it with no @valueIndices@ gather.--Parallel by ROW range with a private per-worker accumulator of @nGroups@ slots,-merged afterwards — which is why it needs a small domain, and why it admits only-order-independent reductions: the merge must be exact for the result to stay-byte-identical to @-N1@. Anything it rejects falls back to-"DataFrame.Internal.Aggregation.Kernel.Scatter".--The caller decides whether the domain is small enough; see @denseThreshold@ in-the operations layer.--}+-- | The low-cardinality dense reduction kernel (direct-indexed by group). module DataFrame.Internal.Aggregation.Kernel.Dense (     denseReduce,     denseMaxMinusMin, ) where -import Control.Monad (when) import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM@@ -43,24 +30,11 @@     shouldParallelize,  ) -{- | Group count at or below which the float (Double sum/mean) direct-reductions run as ONE sequential row-order pass: the accumulators fit in L2, the-pass is memory-bandwidth bound, and adding each value to its group in ascending-row order is exactly the order the group-range gather kernel uses (the grouping-layer's @valueIndices@ is a stable counting sort), so the result is-byte-identical to it. Above this the parallel chunked variant runs instead (see-'sumDblDense'). Var/std are never taken directly: any chunked merge of-variance state changes the float recurrence, and the group-range gather kernel-already runs them in parallel while replaying the interpreter's per-group-update order bit-for-bit.--}+-- | Group count at or below which float reductions run one sequential pass. seqFloatGroups :: Int seqFloatGroups = 65536 -{- | Run a recognised reduction through the direct-indexed path. 'Nothing' (so-the caller falls back to the order-preserving kernel) unless the reduction is-admitted at this element type AND the column is a clean unboxed Int/Double.--}+-- | Run a recognised reduction through the direct-indexed path, if admitted. denseReduce :: Reduction -> VU.Vector Int -> Int -> Column -> Maybe Column denseReduce red g nGroups col = case col of     UnboxedColumn Nothing (v :: VU.Vector a) ->@@ -73,13 +47,7 @@     _ -> Nothing {-# INLINEABLE denseReduce #-} -{- | The reductions admitted over an Int column. Sum/min/max/mean/count are-exact in the Int domain (any merge order gives the same bits); top2sum selects-the two largest values (order-independent as a multiset selection) and only-adds them once at finalize. Var/std stay with the group-range gather kernel:-its per-group Welford recurrence replays the interpreter's update order-bit-for-bit, which no chunk-merged direct pass can.--}+-- | The reductions admitted over an Int column. denseInt :: Reduction -> VU.Vector Int -> Int -> VU.Vector Int -> Maybe Column denseInt red g nGroups v = case red of     RCount -> Just (fromUnboxedVector (countDense g nGroups (VU.length v)))@@ -87,16 +55,9 @@     RMin -> Just (fromUnboxedVector (extremaIntDense True g nGroups v))     RMax -> Just (fromUnboxedVector (extremaIntDense False g nGroups v))     RMean -> Just (fromUnboxedVector (meanIntDense g nGroups v))-    RTop2Sum -> Just (fromUnboxedVector (top2Dense g nGroups v))-    RTop2Snd -> Just (fromUnboxedVector (top2SndDense g nGroups v))     _ -> Nothing -{- | The reductions admitted over a Double column. Count/min/max/top2sum are-order-independent (exact per-worker merge, byte-identical at any @-N@). The-float sum/mean run sequentially in row order below 'seqFloatGroups' (matching-the gather kernel's per-group addition order exactly) and as deterministic-chunked partials above it. Var/std keep the gather kernel (see 'denseInt').--}+-- | The reductions admitted over a Double column. denseDouble ::     Reduction -> VU.Vector Int -> Int -> VU.Vector Double -> Maybe Column denseDouble red g nGroups v = case red of@@ -105,17 +66,9 @@     RMean -> Just (fromUnboxedVector (meanDblDense g nGroups v))     RMin -> Just (fromUnboxedVector (extremaDblDense True g nGroups v))     RMax -> Just (fromUnboxedVector (extremaDblDense False g nGroups v))-    RTop2Sum -> Just (fromUnboxedVector (top2Dense g nGroups v))-    RTop2Snd -> Just (fromUnboxedVector (top2SndDense g nGroups v))     _ -> Nothing -{- | The fused @max a - min b@ direct pass: BOTH extrema accumulate in one-streaming loop over the rows (min/max are order-independent, so the per-worker-merge is exact and the result byte-identical to the two gather passes it-replaces). 'Nothing' unless both columns are clean unboxed and same-typed-(Int/Int keeps the Int result of the interpreter; Double/Double the Double one);-mixed pairs keep the gather fallback.--}+-- | The fused @max a - min b@ direct pass over both extrema at once. denseMaxMinusMin :: VU.Vector Int -> Int -> Column -> Column -> Maybe Column denseMaxMinusMin g nGroups ca cb = case (ca, cb) of     ( UnboxedColumn Nothing (va :: VU.Vector x)@@ -130,12 +83,7 @@     _ -> Nothing {-# INLINEABLE denseMaxMinusMin #-} -{- | Monomorphic entry points: the 'testEquality' dispatch above only yields an-unsafe coercion, so a direct call to the polymorphic 'maxMinusMinDense' there-would stay at the abstract element type and never meet its SPECIALIZE rules-(measured ~3x on the whole pass); calling through these fixed-type wrappers-(the coercion lands on the argument) does.--}+-- | Monomorphic entry points for the 'testEquality' dispatch above. maxMinusMinDenseInt ::     VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int maxMinusMinDenseInt = maxMinusMinDense minBound maxBound@@ -150,20 +98,11 @@ shouldPar :: Int -> Bool shouldPar = shouldParallelize parThreshold -{- | Fork @caps@ workers over disjoint contiguous row ranges of @[0, n)@, each-producing its own private accumulator (no shared array, no sync). Returns the-partials in worker order for the caller's merge; rethrows the first failure.-The chunking is a fixed function of @n@ and @caps@, so any merge over the-partials is deterministic at a given @-N@.--}+-- | Fork workers over disjoint row ranges, each with a private accumulator. runPartialsOver ::     Int -> Int -> (Int -> Int -> IO acc) -> IO [acc] runPartialsOver n _caps = parallelChunks parThreshold n ----------------------------------------------------------------------------------- Count (order-independent: per-group row count)--------------------------------------------------------------------------------- countDense :: VU.Vector Int -> Int -> Int -> VU.Vector Int countDense g nGroups n = unsafePerformIO $ do     parts <- runPartialsOver n capabilities (countChunk g nGroups)@@ -183,10 +122,6 @@     go lo     pure acc ----------------------------------------------------------------------------------- Integer sum (exact: merge order irrelevant)--------------------------------------------------------------------------------- sumIntDense :: VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int sumIntDense g nGroups v = unsafePerformIO $ do     parts <- runPartialsOver (VU.length v) capabilities (sumIntChunk g v nGroups)@@ -207,17 +142,7 @@     go lo     pure acc ----------------------------------------------------------------------------------- Double sum / mean (streaming scatter; chunked partials above seqFloatGroups)----------------------------------------------------------------------------------{- | Double group sums. At or below 'seqFloatGroups' a single sequential pass in-ascending row order — each group's additions happen in exactly the order the-group-range gather kernel performs them, so the result is byte-identical to it.-Above that, per-worker chunk partials merged in worker order: still-deterministic at a fixed @-N@, but the float summation order differs from the-sequential pass.--}+-- | Double group sums, sequential at small group counts. sumDblDense :: VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double sumDblDense g nGroups v     | nGroups <= seqFloatGroups || not (shouldPar n) =@@ -248,10 +173,7 @@     go lo     pure acc -{- | Double mean: fused (Double sum, count) per group, divided once at finalize.-Same order policy as 'sumDblDense' (sequential row order is byte-identical to-the gather kernel; the chunked variant changes the float summation order).--}+-- | Double mean: fused (sum, count) per group, divided once at finalize. meanDblDense :: VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double meanDblDense g nGroups v     | nGroups <= seqFloatGroups || not (shouldPar n) = unsafePerformIO $ do@@ -304,10 +226,6 @@     go 0     VU.unsafeFreeze out ----------------------------------------------------------------------------------- Integer min / max (order-independent)--------------------------------------------------------------------------------- extremaIntDense ::     Bool -> VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int extremaIntDense isMin g nGroups v = unsafePerformIO $ do@@ -338,10 +256,6 @@     go lo     pure acc ----------------------------------------------------------------------------------- Double min / max (order-independent: exact per-worker merge)--------------------------------------------------------------------------------- extremaDblDense ::     Bool -> VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double extremaDblDense isMin g nGroups v@@ -376,14 +290,7 @@     go lo     pure acc ----------------------------------------------------------------------------------- Fused max(a) - min(b) (order-independent: exact per-worker merge)----------------------------------------------------------------------------------{- | One streaming pass accumulating @max a@ and @min b@ together. @maxSeed@ is-the identity of @max@ (the type's least value), @minSeed@ of @min@ (its-greatest).--}+-- | One streaming pass accumulating @max a@ and @min b@ together. maxMinusMinDense ::     (VU.Unbox a, Num a, Ord a) =>     a ->@@ -425,9 +332,7 @@         finalizeMaxMinusMin nGroups mx mn   where     !n = VU.length va-{- INLINEABLE (not NOINLINE) so the SPECIALIZE pragmas above take effect; the-kernel is a pure function of its arguments, so the usual unsafePerformIO-sharing concern does not apply. -}+-- INLINEABLE (not NOINLINE) so the SPECIALIZE pragmas above take effect. {-# INLINEABLE maxMinusMinDense #-}  maxMinusMinChunk ::@@ -497,161 +402,7 @@     go 0     VU.unsafeFreeze out ----------------------------------------------------------------------------------- Top-2 sum (order-independent multiset selection; one float add at finalize)----------------------------------------------------------------------------------{- | Sum of the two largest values per group. Each accumulator holds the-(largest, second-largest) pair seen so far; merging two pairs keeps the top two-of the four candidates. No float ADDITION happens until the single @m1 + m2@ at-finalize, so the result is byte-identical to the gather kernel regardless of-chunking. Mirrors the gather kernel exactly, including the @realToFrac@ per-element and the @-inf -> 0@ guards at finalize.--}-top2Dense ::-    (VU.Unbox a, Real a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double-{-# SPECIALIZE top2Dense ::-    VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Double-    #-}-{-# SPECIALIZE top2Dense ::-    VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double-    #-}-top2Dense g nGroups v-    | not (shouldPar n) = unsafePerformIO $ do-        (m1, m2) <- top2Chunk g v nGroups 0 n-        finalizeTop2 nGroups m1 m2-    | otherwise = unsafePerformIO $ do-        parts <- runPartialsOver n capabilities (top2Chunk g v nGroups)-        (m1, m2) <- mergeTop2 nGroups parts-        finalizeTop2 nGroups m1 m2-  where-    !n = VU.length v-{- INLINEABLE (not NOINLINE) so the SPECIALIZE pragmas above take effect; pure-function of its arguments, so unsafePerformIO sharing is not a concern. -}-{-# INLINEABLE top2Dense #-}--{- | Second-largest value per group: the exact same per-worker-(largest, second-largest) accumulator and merge as 'top2Dense'-('top2Chunk'/'mergeTop2'), finalized to the second max alone. A group of-size < 2 finalizes its @-inf@ seed to NaN — documented behaviour (the-db-benchmark Q8 data has no size-1 @id6@ groups). Order-independent multiset-selection, so byte-identical at any @-N@.--}-top2SndDense ::-    (VU.Unbox a, Real a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double-{-# SPECIALIZE top2SndDense ::-    VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Double-    #-}-{-# SPECIALIZE top2SndDense ::-    VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double-    #-}-top2SndDense g nGroups v-    | not (shouldPar n) = unsafePerformIO $ do-        (m1, m2) <- top2Chunk g v nGroups 0 n-        finalizeTop2Snd nGroups m1 m2-    | otherwise = unsafePerformIO $ do-        parts <- runPartialsOver n capabilities (top2Chunk g v nGroups)-        (m1, m2) <- mergeTop2 nGroups parts-        finalizeTop2Snd nGroups m1 m2-  where-    !n = VU.length v-{- INLINEABLE (not NOINLINE) so the SPECIALIZE pragmas above take effect; pure-function of its arguments, so unsafePerformIO sharing is not a concern. -}-{-# INLINEABLE top2SndDense #-}--top2Chunk ::-    (VU.Unbox a, Real a) =>-    VU.Vector Int ->-    VU.Vector a ->-    Int ->-    Int ->-    Int ->-    IO (VUM.IOVector Double, VUM.IOVector Double)-top2Chunk g v nGroups lo hi = do-    let ninf = negate (1 / 0) :: Double-    m1 <- VUM.replicate nGroups ninf-    m2 <- VUM.replicate nGroups ninf-    let go !i-            | i >= hi = pure ()-            | otherwise = do-                let !k = VU.unsafeIndex g i-                    !x = realToFrac (VU.unsafeIndex v i)-                a1 <- VUM.unsafeRead m1 k-                if x > a1-                    then do-                        VUM.unsafeWrite m1 k x-                        VUM.unsafeWrite m2 k a1-                    else do-                        a2 <- VUM.unsafeRead m2 k-                        when (x > a2) (VUM.unsafeWrite m2 k x)-                go (i + 1)-    go lo-    pure (m1, m2)--{- | Top two of the four candidates @{a1, a2, b1, b2}@ per group (each pair-already ordered @m1 >= m2@, @-inf@ seeds included).--}-mergeTop2 ::-    Int ->-    [(VUM.IOVector Double, VUM.IOVector Double)] ->-    IO (VUM.IOVector Double, VUM.IOVector Double)-mergeTop2 nGroups parts = case parts of-    [] -> error "mergeTop2: no partials"-    ((m10, m20) : rest) -> do-        let add (m1, m2) = do-                let go !k-                        | k >= nGroups = pure ()-                        | otherwise = do-                            a1 <- VUM.unsafeRead m10 k-                            a2 <- VUM.unsafeRead m20 k-                            b1 <- VUM.unsafeRead m1 k-                            b2 <- VUM.unsafeRead m2 k-                            if b1 > a1-                                then do-                                    VUM.unsafeWrite m10 k b1-                                    VUM.unsafeWrite m20 k (max a1 b2)-                                else VUM.unsafeWrite m20 k (max a2 b1)-                            go (k + 1)-                go 0-        mapM_ add rest-        pure (m10, m20)--finalizeTop2 ::-    Int -> VUM.IOVector Double -> VUM.IOVector Double -> IO (VU.Vector Double)-finalizeTop2 nGroups m1 m2 = do-    out <- VUM.new nGroups-    let go !k-            | k >= nGroups = pure ()-            | otherwise = do-                a1 <- VUM.unsafeRead m1 k-                a2 <- VUM.unsafeRead m2 k-                let s = (if isInfinite a1 then 0 else a1) + (if isInfinite a2 then 0 else a2)-                VUM.unsafeWrite out k s-                go (k + 1)-    go 0-    VU.unsafeFreeze out--finalizeTop2Snd ::-    Int -> VUM.IOVector Double -> VUM.IOVector Double -> IO (VU.Vector Double)-finalizeTop2Snd nGroups _m1 m2 = do-    out <- VUM.new nGroups-    let go !k-            | k >= nGroups = pure ()-            | otherwise = do-                a2 <- VUM.unsafeRead m2 k-                VUM.unsafeWrite out k (if isInfinite a2 then 0 / 0 else a2)-                go (k + 1)-    go 0-    VU.unsafeFreeze out------------------------------------------------------------------------------------ Integer mean (exact integer sum + count, divided once -> order-independent)----------------------------------------------------------------------------------{- | Integer mean in ONE fused pass: a running integer sum and count per group,-divided once at finalize. The integer sum is exact, so the parallel partial-merge is byte-identical to the sequential single pass at any @-N@.--}+-- | Integer mean in one fused pass: running sum and count per group. meanIntDense :: VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Double meanIntDense g nGroups v     | not (shouldPar n) = unsafePerformIO $ do@@ -704,10 +455,6 @@     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@@ -724,10 +471,7 @@         mapM_ add rest         VU.unsafeFreeze p0 -{- | Sum the per-worker Double partials into the first worker's accumulator IN-WORKER ORDER: deterministic at a fixed @-N@, but the float summation order is-chunk-major rather than the sequential row order.--}+-- | Sum per-worker Double partials into the first worker's, in worker order. mergeDblSum :: Int -> [VUM.IOVector Double] -> IO (VU.Vector Double) mergeDblSum nGroups parts = case parts of     [] -> VU.unsafeFreeze =<< VUM.replicate nGroups 0@@ -744,9 +488,7 @@         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.--}+-- | Merge per-worker (sum, count) partials by exact integer addition. mergePair ::     Int ->     [(VUM.IOVector Int, VUM.IOVector Int)] ->@@ -769,9 +511,7 @@         mapM_ add rest         pure (s0, c0) -{- | As 'mergePair' but for the Double (sum, count) partials of the Double mean;-worker-order float sums (see 'mergeDblSum').--}+-- | As 'mergePair' but for the Double (sum, count) partials of the mean. mergeMeanDbl ::     Int ->     [(VUM.IOVector Double, VUM.IOVector Int)] ->

@@ -0,0 +1,609 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE UnboxedTuples #-}++module DataFrame.Internal.Aggregation.Kernel.Distinct (+    distinctSupported,+    countDistinctColumn,+    countDistinctGroups,+) where++import Control.Monad (forM_, replicateM, void, when)+import Data.Bits (+    countLeadingZeros,+    finiteBitSize,+    unsafeShiftL,+    unsafeShiftR,+    xor,+    (.&.),+ )+import Data.Char (ord)+import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef, writeIORef)+import Data.Int (Int16, Int32, Int64, Int8)+import Data.Primitive.ByteArray (MutableByteArray (..))+import qualified Data.Text.Array as A+import Data.Text.Internal (Text (Text))+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))+import qualified Data.Vector as V+import qualified Data.Vector.Primitive as VP+import qualified Data.Vector.Primitive.Mutable as VPM+import qualified Data.Vector.Unboxed as VU+import Data.Vector.Unboxed.Base (+    MVector (MV_Int),+    Vector (V_Int, V_Int64, V_Word64),+ )+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.Word (Word16, Word32, Word64, Word8)+import GHC.Exts (Int (I#), fetchAddIntArray#, (+#))+import GHC.Float (castDoubleToWord64)+import GHC.IO (IO (IO))+import System.IO.Unsafe (unsafePerformIO)+import Type.Reflection (typeRep)++import DataFrame.Internal.Algorithms.Hash (fnvOffset, mixBytes)+import DataFrame.Internal.Column (Column (..))+import DataFrame.Internal.Control.Concurrent (+    capabilities,+    forkJoin_,+    parThreshold,+    shouldParallelize,+    splitChunkRange,+ )+import DataFrame.Internal.Data.PackedText (+    PackedOffsets (..),+    PackedTextData,+    packedLength,+    packedRowOffsets,+    packedSlice,+ )++-- | How a column's rows become set keys.+data Keys+    = -- | One 64-bit word per row; equal words iff equal values.+      WordKeys !(VU.Vector Int)+    | -- | A UTF-8 byte slice per row.+      ByteKeys !Int !TextRows++-- | Where a text row's bytes live, with kernel instances for each layout.+data TextRows+    = Contig32 !A.Array !(VU.Vector Int32)+    | Contig64 !A.Array !(VU.Vector Int)+    | AnyRows (Int -> (A.Array, Int, Int))++-- | Whether 'countDistinctColumn' / 'countDistinctGroups' accept this column.+distinctSupported :: Column -> Bool+distinctSupported c = case columnKeys c of+    Just _ -> True+    Nothing -> False++columnKeys :: Column -> Maybe Keys+columnKeys col = case col of+    UnboxedColumn Nothing (v :: VU.Vector a)+        | Just Refl <- testEquality (typeRep @a) (typeRep @Int) -> Just (WordKeys v)+        | Just Refl <- testEquality (typeRep @a) (typeRep @Int64) ->+            Just (WordKeys (int64Words v))+        | Just Refl <- testEquality (typeRep @a) (typeRep @Word64) ->+            Just (WordKeys (word64Words v))+        | Just Refl <- testEquality (typeRep @a) (typeRep @Int32) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Int16) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Int8) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Word) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Word32) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Word16) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Word8) -> widen v+        | Just Refl <- testEquality (typeRep @a) (typeRep @Double) ->+            Just (WordKeys (VU.map doubleWord v))+        | Just Refl <- testEquality (typeRep @a) (typeRep @Char) ->+            Just (WordKeys (VU.map ord v))+        | Just Refl <- testEquality (typeRep @a) (typeRep @Bool) ->+            Just (WordKeys (VU.map fromEnum v))+    BoxedColumn Nothing (v :: V.Vector a)+        | Just Refl <- testEquality (typeRep @a) (typeRep @Text) ->+            Just+                ( ByteKeys+                    (V.length v)+                    (AnyRows (\i -> case V.unsafeIndex v i of Text arr o l -> (arr, o, l)))+                )+    PackedText Nothing p -> Just (ByteKeys (packedLength p) (packedRows p))+    _ -> Nothing+  where+    widen :: (Integral b, VU.Unbox b) => VU.Vector b -> Maybe Keys+    widen = Just . WordKeys . VU.map fromIntegral++packedRows :: PackedTextData -> TextRows+packedRows p = case packedRowOffsets p of+    Just (arr, Offs32 offs) -> Contig32 arr offs+    Just (arr, Offs64 offs) -> Contig64 arr offs+    Nothing -> AnyRows (packedSlice p)++-- | Reinterpret 64-bit words as 'Int' without copying (same width and layout).+int64Words :: VU.Vector Int64 -> VU.Vector Int+int64Words v@(V_Int64 (VP.Vector o l ba))+    | finiteBitSize (0 :: Int) == 64 = V_Int (VP.Vector o l ba)+    | otherwise = VU.map fromIntegral v++word64Words :: VU.Vector Word64 -> VU.Vector Int+word64Words v@(V_Word64 (VP.Vector o l ba))+    | finiteBitSize (0 :: Int) == 64 = V_Int (VP.Vector o l ba)+    | otherwise = VU.map fromIntegral v++-- | Equal doubles share a word: @-0.0@ joins @0.0@ and every NaN is one value.+doubleWord :: Double -> Int+doubleWord d+    | isNaN d = 0x7ff8000000000000+    | d == 0 = 0+    | otherwise = fromIntegral (castDoubleToWord64 d)++-- | Number of distinct values in a column, or 'Nothing' if not admitted.+countDistinctColumn :: Column -> Maybe Int+countDistinctColumn col = case columnKeys col of+    Nothing -> Nothing+    Just (WordKeys v) -> Just $! unsafePerformIO (wordGlobal v)+    Just (ByteKeys n rows) -> Just $! unsafePerformIO $ case rows of+        Contig32 arr offs -> bytesGlobal n (contigSlice arr offs)+        Contig64 arr offs -> bytesGlobal n (contigSlice arr offs)+        AnyRows slice -> bytesGlobal n slice+{-# NOINLINE countDistinctColumn #-}++-- | Distinct values per group, or 'Nothing' if the column is not admitted.+countDistinctGroups :: VU.Vector Int -> Int -> Column -> Maybe (VU.Vector Int)+countDistinctGroups rtg nGroups col = case columnKeys col of+    Nothing -> Nothing+    Just (WordKeys v) -> Just $! unsafePerformIO (wordGrouped rtg nGroups v)+    Just (ByteKeys _ rows) -> Just $! unsafePerformIO $ case rows of+        Contig32 arr offs -> bytesGrouped rtg nGroups (contigSlice arr offs)+        Contig64 arr offs -> bytesGrouped rtg nGroups (contigSlice arr offs)+        AnyRows slice -> bytesGrouped rtg nGroups slice+{-# NOINLINE countDistinctGroups #-}++contigSlice ::+    (Integral o, VU.Unbox o) => A.Array -> VU.Vector o -> Int -> (A.Array, Int, Int)+contigSlice arr offs i =+    let !o = fromIntegral (VU.unsafeIndex offs i)+     in (arr, o, fromIntegral (VU.unsafeIndex offs (i + 1)) - o)+{-# INLINE contigSlice #-}++wordGlobal :: VU.Vector Int -> IO Int+wordGlobal v = do+    let !n = VU.length v+        !pb = partitionBits n+        key = VU.unsafeIndex v+    Scattered w0 _ starts <-+        scatterRows+            n+            pb+            False+            False+            (fmix . key)+            (\r r' -> key r == key r')+            (\i _ -> key i)+            (const 0)+    counts <- forPartitions (VU.length starts - 1) (newIORef (0 :: Int)) $ \acc p -> do+        c <- dedupWords pb w0 (VU.unsafeIndex starts p) (VU.unsafeIndex starts (p + 1))+        modifyStrict acc (+ c)+    sum <$> mapM readIORef counts++wordGrouped :: VU.Vector Int -> Int -> VU.Vector Int -> IO (VU.Vector Int)+wordGrouped rtg nGroups v = do+    let !n = VU.length v+        !pb = partitionBits n+        key = VU.unsafeIndex v+        gid = VU.unsafeIndex rtg+    Scattered w0 w1 starts <-+        scatterRows+            n+            pb+            False+            True+            (\i -> pairHash (gid i) (key i))+            (\r r' -> key r == key r' && gid r == gid r')+            (\i _ -> key i)+            gid+    groupCounts nGroups (VU.length starts - 1) $ \counts p ->+        void $+            dedupPairs+                pb+                w0+                w1+                (VU.unsafeIndex starts p)+                (VU.unsafeIndex starts (p + 1))+                pairHash'+                (\a0 a1 b0 b1 -> a0 == b0 && a1 == b1)+                (\_ g -> bumpGroup counts g)+  where+    pairHash' k g = pairHash g k++bytesGlobal :: Int -> (Int -> (A.Array, Int, Int)) -> IO Int+bytesGlobal n slice = do+    let !pb = partitionBits n+        hashRow i = case slice i of (arr, o, l) -> fmix (mixBytes fnvOffset arr o l)+    Scattered w0 w1 starts <-+        scatterRows n pb True True hashRow (rowsEq slice) (\_ h -> h) id+    counts <- forPartitions (VU.length starts - 1) (newIORef (0 :: Int)) $ \acc p -> do+        c <-+            dedupPairs+                pb+                w0+                w1+                (VU.unsafeIndex starts p)+                (VU.unsafeIndex starts (p + 1))+                const+                (\h r h' r' -> h == h' && rowsEq slice r r')+                (\_ _ -> pure ())+        modifyStrict acc (+ c)+    sum <$> mapM readIORef counts++bytesGrouped ::+    VU.Vector Int -> Int -> (Int -> (A.Array, Int, Int)) -> IO (VU.Vector Int)+bytesGrouped rtg nGroups slice = do+    let !n = VU.length rtg+        !pb = partitionBits n+        gid = VU.unsafeIndex rtg+        hashRow i = case slice i of+            (arr, o, l) -> pairHash (gid i) (mixBytes fnvOffset arr o l)+    Scattered w0 w1 starts <-+        scatterRows+            n+            pb+            True+            True+            hashRow+            (\r r' -> gid r == gid r' && rowsEq slice r r')+            (\_ h -> h)+            id+    groupCounts nGroups (VU.length starts - 1) $ \counts p ->+        void $+            dedupPairs+                pb+                w0+                w1+                (VU.unsafeIndex starts p)+                (VU.unsafeIndex starts (p + 1))+                const+                (\h r h' r' -> h == h' && gid r == gid r' && rowsEq slice r r')+                (\_ r -> bumpGroup counts (gid r))++rowsEq :: (Int -> (A.Array, Int, Int)) -> Int -> Int -> Bool+rowsEq slice r r' = case slice r of+    (a, o, l) -> case slice r' of+        (a', o', l') -> l == l' && A.compare a o a' o' l == EQ+{-# INLINE rowsEq #-}++{-# INLINE bytesGlobal #-}+{-# INLINE bytesGrouped #-}++-- | Murmur3's 64-bit finaliser: every input bit reaches every output bit.+fmix :: Int -> Int+fmix k =+    let !h0 = fromIntegral k :: Word64+        !h1 = (h0 `xor` (h0 `unsafeShiftR` 33)) * 0xff51afd7ed558ccd+        !h2 = (h1 `xor` (h1 `unsafeShiftR` 33)) * 0xc4ceb9fe1a85ec53+     in fromIntegral (h2 `xor` (h2 `unsafeShiftR` 33))+{-# INLINE fmix #-}++pairHash :: Int -> Int -> Int+pairHash g k = fmix (k `xor` (g * 0x9E3779B97F4A7C1))+{-# INLINE pairHash #-}++-- | Partition of a hash: its top @pb@ bits (@pb >= 1@).+partOf :: Int -> Int -> Int+partOf pb h = fromIntegral ((fromIntegral h :: Word64) `unsafeShiftR` (64 - pb))+{-# INLINE partOf #-}++-- | Slot of a hash in a @2^bits@ table: the bits just below the partition bits.+slotOf :: Int -> Int -> Int -> Int+slotOf pb bits h =+    fromIntegral+        (((fromIntegral h :: Word64) `unsafeShiftL` pb) `unsafeShiftR` (64 - bits))+{-# INLINE slotOf #-}++-- | Partition count as a bit width, sized for cache-resident sets.+partitionBits :: Int -> Int+partitionBits n+    | not (shouldParallelize parThreshold n) = 1+    | otherwise =+        max+            1+            ( min+                maxPartitionBits+                (max (ceilLog2 (n `div` rowsPerPartition)) (ceilLog2 (4 * capabilities)))+            )++maxPartitionBits :: Int+maxPartitionBits = 8++rowsPerPartition :: Int+rowsPerPartition = 65536++ceilLog2 :: Int -> Int+ceilLog2 x+    | x <= 1 = 0+    | otherwise = finiteBitSize x - countLeadingZeros (x - 1)++-- | Entries in each chunk's cache of recently emitted keys.+cacheSize :: Int+cacheSize = 4096++-- | Partition mark of a row dropped as a cached duplicate.+droppedRow :: Word16+droppedRow = maxBound++-- | Rows scattered by partition: key words, optional second words, starts.+data Scattered = Scattered !(VU.Vector Int) !(VU.Vector Int) !(VU.Vector Int)++-- | Hash every row to a partition, then write each row's words into it.+scatterRows ::+    Int ->+    Int ->+    -- | keep the hash between passes+    Bool ->+    -- | write a second word per row+    Bool ->+    -- | row hash+    (Int -> Int) ->+    -- | do two rows hold the same key?+    (Int -> Int -> Bool) ->+    -- | first word, from row and hash+    (Int -> Int -> Int) ->+    -- | second word, from row+    (Int -> Int) ->+    IO Scattered+scatterRows n pb keepHash twoWords rowHash rowEq rowW0 rowW1 = do+    let !np = 1 `unsafeShiftL` pb+        chunks+            | shouldParallelize parThreshold n = splitChunkRange capabilities n+            | otherwise = [(0, n)]+        !nc = length chunks+    hashes <- if keepHash then VUM.unsafeNew n else VUM.unsafeNew 0+    parts <- VUM.unsafeNew n :: IO (VUM.IOVector Word16)+    hist <- VUM.replicate (nc * np) (0 :: Int)+    forkJoin_+        [ do+            cacheHash <- VUM.unsafeNew cacheSize+            cacheRow <- VUM.replicate cacheSize (-1 :: Int)+            let !base = ci * np+                go !i+                    | i >= hi = pure ()+                    | otherwise = do+                        let !h = rowHash i+                            !c = h .&. (cacheSize - 1)+                        when keepHash (VUM.unsafeWrite hashes i h)+                        r <- VUM.unsafeRead cacheRow c+                        dup <-+                            if r < 0+                                then pure False+                                else do+                                    h' <- VUM.unsafeRead cacheHash c+                                    pure (h' == h && rowEq r i)+                        if dup+                            then VUM.unsafeWrite parts i droppedRow+                            else do+                                VUM.unsafeWrite cacheHash c h+                                VUM.unsafeWrite cacheRow c i+                                let !p = partOf pb h+                                VUM.unsafeWrite parts i (fromIntegral p)+                                VUM.unsafeModify hist (+ 1) (base + p)+                        go (i + 1)+            go lo+        | (ci, (lo, hi)) <- zip [0 ..] chunks+        ]+    -- Exclusive prefix sum, partition-major, turns counts into write cursors.+    starts <- VUM.unsafeNew (np + 1)+    let prefix !p !run+            | p >= np = VUM.unsafeWrite starts np run+            | otherwise = do+                VUM.unsafeWrite starts p run+                let walk !ci !r+                        | ci >= nc = pure r+                        | otherwise = do+                            c <- VUM.unsafeRead hist (ci * np + p)+                            VUM.unsafeWrite hist (ci * np + p) r+                            walk (ci + 1) (r + c)+                run' <- walk 0 run+                prefix (p + 1) run'+    prefix 0 0+    kept <- VUM.unsafeRead starts np+    w0 <- VUM.unsafeNew kept+    w1 <- if twoWords then VUM.unsafeNew kept else VUM.unsafeNew 0+    forkJoin_+        [ let !base = ci * np+              go !i+                | i >= hi = pure ()+                | otherwise = do+                    p <- VUM.unsafeRead parts i+                    when (p /= droppedRow) $ do+                        h <- if keepHash then VUM.unsafeRead hashes i else pure 0+                        let !slot = base + fromIntegral p+                        pos <- VUM.unsafeRead hist slot+                        VUM.unsafeWrite hist slot (pos + 1)+                        VUM.unsafeWrite w0 pos (rowW0 i h)+                        when twoWords (VUM.unsafeWrite w1 pos (rowW1 i))+                    go (i + 1)+           in go lo+        | (ci, (lo, hi)) <- zip [0 ..] chunks+        ]+    Scattered+        <$> VU.unsafeFreeze w0+        <*> VU.unsafeFreeze w1+        <*> VU.unsafeFreeze starts+{-# INLINE scatterRows #-}++-- | Run @body@ over partitions on a worker pool pulling from a shared counter.+forPartitions :: Int -> IO s -> (s -> Int -> IO ()) -> IO [s]+forPartitions np newState body = do+    next <- newIORef 0+    let !width = if np > 1 && capabilities > 1 then min capabilities np else 1+    states <- replicateM width newState+    let worker st = do+            p <- atomicModifyIORef' next (\j -> (j + 1, j))+            when (p < np) (body st p >> worker st)+    forkJoin_ (map worker states)+    pure states+{-# INLINE forPartitions #-}++-- | Per-group counts from a body that bumps a counter once per new pair.+groupCounts :: Int -> Int -> (Counts -> Int -> IO ()) -> IO (VU.Vector Int)+groupCounts nGroups np body+    | nGroups * capabilities <= 1 `unsafeShiftL` 22 = do+        arrs <- forPartitions np (VUM.replicate nGroups 0) $ \arr p ->+            body (Counts False arr) p+        total <- VUM.replicate nGroups 0+        forM_ arrs $ \arr ->+            forM_ [0 .. nGroups - 1] $ \g -> do+                c <- VUM.unsafeRead arr g+                VUM.unsafeModify total (+ c) g+        VU.unsafeFreeze total+    | otherwise = do+        shared <- VUM.replicate nGroups 0+        _ <- forPartitions np (pure ()) $ \() p -> body (Counts True shared) p+        VU.unsafeFreeze shared+{-# INLINE groupCounts #-}++-- | A group-count array and whether other workers share it.+data Counts = Counts !Bool !(VUM.IOVector Int)++bumpGroup :: Counts -> Int -> IO ()+bumpGroup (Counts shared arr) g+    | shared = atomicIncr arr g+    | otherwise = VUM.unsafeModify arr (+ 1) g+{-# INLINE bumpGroup #-}++atomicIncr :: VUM.IOVector Int -> Int -> IO ()+atomicIncr (MV_Int (VPM.MVector (I# off) _ (MutableByteArray mba))) (I# i) =+    IO+        (\s -> case fetchAddIntArray# mba (off +# i) 1# s of (# s', _ #) -> (# s', () #))+{-# INLINE atomicIncr #-}++modifyStrict :: IORef Int -> (Int -> Int) -> IO ()+modifyStrict ref f = readIORef ref >>= \x -> writeIORef ref $! f x+{-# INLINE modifyStrict #-}++-- | Starting table width for a partition of @rows@ entries.+initialBits :: Int -> Int+initialBits rows = max 4 (min 12 (ceilLog2 (2 * rows)))++-- | Distinct words in @w0[s, e)@, with 'minBound' tracked by a flag.+dedupWords :: Int -> VU.Vector Int -> Int -> Int -> IO Int+dedupWords pb w0 s e = do+    let !bits0 = initialBits (e - s)+    t0 <- VUM.replicate (1 `unsafeShiftL` bits0) emptyWord+    let go !t !bits !cnt !sawEmpty !i+            | i >= e = pure (if sawEmpty then cnt + 1 else cnt)+            | otherwise = do+                let !k = VU.unsafeIndex w0 i+                if k == emptyWord+                    then go t bits cnt True (i + 1)+                    else do+                        let !mask = (1 `unsafeShiftL` bits) - 1+                            probe !sl = do+                                x <- VUM.unsafeRead t sl+                                if x == emptyWord+                                    then VUM.unsafeWrite t sl k >> pure True+                                    else+                                        if x == k+                                            then pure False+                                            else probe ((sl + 1) .&. mask)+                        new <- probe (slotOf pb bits (fmix k))+                        if not new+                            then go t bits cnt sawEmpty (i + 1)+                            else+                                if 2 * (cnt + 1) > mask + 1+                                    then do+                                        t' <- growWords pb t bits+                                        go t' (bits + 1) (cnt + 1) sawEmpty (i + 1)+                                    else go t bits (cnt + 1) sawEmpty (i + 1)+    go t0 bits0 0 False s++emptyWord :: Int+emptyWord = minBound++growWords :: Int -> VUM.IOVector Int -> Int -> IO (VUM.IOVector Int)+growWords pb t bits = do+    let !bits' = bits + 1+        !mask' = (1 `unsafeShiftL` bits') - 1+    t' <- VUM.replicate (mask' + 1) emptyWord+    forM_ [0 .. VUM.length t - 1] $ \j -> do+        k <- VUM.unsafeRead t j+        when (k /= emptyWord) $ do+            let place !sl = do+                    x <- VUM.unsafeRead t' sl+                    if x == emptyWord+                        then VUM.unsafeWrite t' sl k+                        else place ((sl + 1) .&. mask')+            place (slotOf pb bits' (fmix k))+    pure t'++-- | Distinct @(w0, w1)@ entries in @[s, e)@ under @eq@, slotted by @hashOf@.+dedupPairs ::+    Int ->+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    Int ->+    (Int -> Int -> Int) ->+    (Int -> Int -> Int -> Int -> Bool) ->+    (Int -> Int -> IO ()) ->+    IO Int+dedupPairs pb w0 w1 s e hashOf eq onNew = do+    let !bits0 = initialBits (e - s)+    a0 <- VUM.unsafeNew (1 `unsafeShiftL` bits0)+    a1 <- VUM.replicate (1 `unsafeShiftL` bits0) (-1)+    let go !t0 !t1 !bits !cnt !i+            | i >= e = pure cnt+            | otherwise = do+                let !k0 = VU.unsafeIndex w0 i+                    !k1 = VU.unsafeIndex w1 i+                    !mask = (1 `unsafeShiftL` bits) - 1+                    probe !sl = do+                        y <- VUM.unsafeRead t1 sl+                        if y < 0+                            then do+                                VUM.unsafeWrite t0 sl k0+                                VUM.unsafeWrite t1 sl k1+                                pure True+                            else do+                                x <- VUM.unsafeRead t0 sl+                                if eq x y k0 k1+                                    then pure False+                                    else probe ((sl + 1) .&. mask)+                new <- probe (slotOf pb bits (hashOf k0 k1))+                if not new+                    then go t0 t1 bits cnt (i + 1)+                    else do+                        onNew k0 k1+                        if 2 * (cnt + 1) > mask + 1+                            then do+                                (t0', t1') <- growPairs pb hashOf t0 t1 bits+                                go t0' t1' (bits + 1) (cnt + 1) (i + 1)+                            else go t0 t1 bits (cnt + 1) (i + 1)+    go a0 a1 bits0 0 s+{-# INLINE dedupPairs #-}++growPairs ::+    Int ->+    (Int -> Int -> Int) ->+    VUM.IOVector Int ->+    VUM.IOVector Int ->+    Int ->+    IO (VUM.IOVector Int, VUM.IOVector Int)+growPairs pb hashOf t0 t1 bits = do+    let !bits' = bits + 1+        !mask' = (1 `unsafeShiftL` bits') - 1+    n0 <- VUM.unsafeNew (mask' + 1)+    n1 <- VUM.replicate (mask' + 1) (-1)+    forM_ [0 .. VUM.length t1 - 1] $ \j -> do+        y <- VUM.unsafeRead t1 j+        when (y >= 0) $ do+            x <- VUM.unsafeRead t0 j+            let place !sl = do+                    y' <- VUM.unsafeRead n1 sl+                    if y' < 0+                        then VUM.unsafeWrite n0 sl x >> VUM.unsafeWrite n1 sl y+                        else place ((sl + 1) .&. mask')+            place (slotOf pb bits' (hashOf x y))+    pure (n0, n1)+{-# INLINE growPairs #-}

@@ -4,22 +4,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | Fused multi-reduction aggregation kernels: several reductions over the-SAME grouping, evaluated in one pass over the rows.--Two shapes, picked by group-domain width:--* 'mkFusedAgg' \/ 'runFusedAggs' — the STREAMING pass, driven by @rowToGroup@.-  No @valueIndices@ and no placement pass; each worker keeps a private-  accumulator per reduction and the partials merge over group slices. Capped at-  'streamGroupCap', above which the per-worker arrays overflow cache.--* 'mkGatherAgg' \/ 'runGatherAggs' — the GATHER pass, driven by the grouped-  @(valueIndices, offsets)@ layout, for group domains too wide to stream.--Both amortize the memory traffic of the grouping across every reduction in the-batch instead of re-reading it once per reduction.--}+-- | Fused multi-reduction kernels: several reductions in one pass. module DataFrame.Internal.Aggregation.Kernel.Fused (     FusedAgg,     mkFusedAgg,@@ -30,7 +15,7 @@ ) where  import Control.Exception (evaluate)-import Control.Monad (replicateM, when)+import Control.Monad (replicateM) import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM@@ -57,15 +42,7 @@ shouldPar :: Int -> Bool shouldPar = shouldParallelize parThreshold ----------------------------------------------------------------------------------- Fused multi-reduction streaming pass----------------------------------------------------------------------------------{- | One fused reduction of a multi-reduction streaming pass: how to allocate a-per-worker accumulator, fold a row range into it, merge another worker's-accumulator into it over a group range (callers merge in worker order), and-finalize the fully merged accumulator into the output column.--}+-- | One fused reduction of a multi-reduction streaming pass. data FusedAgg     = forall s.         FusedAgg@@ -79,28 +56,11 @@  -- \^ finalize the merged accumulator -{- | Below this many groups a Double sum/mean does NOT join the fused streaming-pass: the per-expression kernel it would replace-('DataFrame.Internal.AggKernelDirect.sumDblDirect' under its @seqFloatGroups@-policy, value mirrored here) runs those as ONE sequential row-order pass that-is byte-identical to the gather kernels and the interpreter, and the vectorized-parity gate asserts exactly that. Above it the per-expression kernel already-merges chunk partials in worker order, so fusing changes nothing semantically-new. Int reductions and count/min/max are exact under any merge and always-fuse.--}+-- | Below this many groups a Double sum/mean skips the fused streaming pass. fusedSeqFloatGroups :: Int fusedSeqFloatGroups = 65536 -{- | Build the fused-pass reduction for one @(reduction, column)@ pair, or-'Nothing' when the pair cannot stream (nullable/boxed columns; the-order-sensitive var/std/top2 reductions, which keep their per-expression-kernels; or a Double sum/mean below 'fusedSeqFloatGroups', which keeps its-byte-identical sequential pass). Sum/min/max/count/mean over Int are exact-under any chunk merge; the admitted Double sum and mean merge their per-worker-partials in worker order (deterministic at a fixed @-N@, chunk-major float-summation order).--}+-- | Build the fused-pass reduction for a pair, or 'Nothing' if it can't stream. mkFusedAgg :: Int -> VU.Vector Int -> Reduction -> Column -> Maybe FusedAgg mkFusedAgg nGroups rtg red col     | nGroups <= 0 || nGroups > streamGroupCap = Nothing@@ -125,27 +85,10 @@                             Just (meanDblFusedAgg nGroups rtg v)                     RMin -> Just (extremaDblFusedAgg True nGroups rtg v)                     RMax -> Just (extremaDblFusedAgg False nGroups rtg v)-                    -- Top-2 selection is an exact multiset selection (no float-                    -- arithmetic before finalize), so its per-worker merge is-                    -- byte-identical to the per-expression kernels at any -N.-                    RTop2Snd -> Just (top2SndDblFusedAgg nGroups rtg v)                     _ -> Nothing         _ -> Nothing -{- | Execute all fused reductions in ONE pass over @rowToGroup@ and the value-columns: the rows are split into one contiguous chunk per capability, each-worker walks its chunk in 4096-row blocks running every reduction's-accumulate-step on the block (the block's @rowToGroup@ slice stays in L1-across the k steps), and each reduction then merges its per-worker partials in-fixed worker order (parallel over group slices) and finalizes. Per-group update-order within each worker is ascending original row order, and chunk boundaries-are a fixed function of the row and capability counts, so the result is-deterministic at a fixed @-N@; on a single capability it is bit-identical to-running each unfused sequential scatter kernel separately.--Pure w.r.t. its immutable inputs (deterministic fan-out and merge order), so-the 'unsafePerformIO' is safe.--}+-- | Execute all fused reductions in one pass over @rowToGroup@ and values. runFusedAggs :: Int -> Int -> [FusedAgg] -> [Column] runFusedAggs n nGroups aggs = unsafePerformIO $ do     let !caps' = if shouldPar n then capabilities else 1@@ -163,9 +106,7 @@     mapM snd opened {-# NOINLINE runFusedAggs #-} -{- | Open one fused reduction for @caps'@ workers: its per-worker step-functions (worker order) and the merge+finalize action.--}+-- | Open one fused reduction for @caps'@ workers: step functions and merge. openFusedAgg :: Int -> Int -> FusedAgg -> IO ([Int -> Int -> IO ()], IO Column) openFusedAgg caps' nGroups (FusedAgg new step mergeR fin) = do     ss <- replicateM caps' new@@ -255,14 +196,7 @@             VUM.unsafeWrite acc k (c + VU.unsafeIndex v i)             go (i + 1) -{- | The mean aggs hold their (sum, count) pair INTERLEAVED in one array —-slots @2g@/@2g+1@ share a cache line (pairs are 16-byte aligned, so they never-straddle one), halving the accumulator misses of the random per-row update-against two separate arrays (measured ~13% off a fused sum+mean pass at 1e6-groups / 1e8 rows on -N16). The count is exact in both layouts (an integer, or-integer-valued Double additions well below 2^53), so sums, merges and the-finalize divide are bit-identical to the two-array layout.--}+-- | Mean aggs keep their (sum, count) pair interleaved in one array. meanIntFusedAgg :: Int -> VU.Vector Int -> VU.Vector Int -> FusedAgg meanIntFusedAgg nGroups rtg v =     FusedAgg@@ -401,78 +335,6 @@             VUM.unsafeWrite acc k (if isMin then min c x else max c x)             go (i + 1) -{- | Fused second-largest over a Double column. The per-group-(largest, second-largest) pair is INTERLEAVED at slots @2g@/@2g+1@ (one cache-line per group, as 'meanDblFusedAgg'); the update is the same top-2 selection-as every other top2 kernel, the merge keeps the top two of the four candidates-per group (exact — no float arithmetic), and the finalize returns the second-max, NaN for a group of size < 2 (the @-inf@ seed; see-'DataFrame.Internal.AggKernel.top2SndScatter').--}-top2SndDblFusedAgg :: Int -> VU.Vector Int -> VU.Vector Double -> FusedAgg-top2SndDblFusedAgg nGroups rtg v =-    FusedAgg-        (VUM.replicate (2 * nGroups) (negate (1 / 0) :: Double))-        (top2SndStepDbl rtg v)-        mergeTop2Range-        ( \s -> do-            sv <- VU.unsafeFreeze s-            pure-                ( fromUnboxedVector-                    ( VU.generate-                        nGroups-                        ( \g ->-                            let !a2 = VU.unsafeIndex sv (2 * g + 1)-                             in if isInfinite a2 then 0 / 0 else a2-                        )-                    )-                )-        )--top2SndStepDbl ::-    VU.Vector Int ->-    VU.Vector Double ->-    VUM.IOVector Double ->-    Int ->-    Int ->-    IO ()-top2SndStepDbl rtg v s lo hi = go lo-  where-    go !i-        | i >= hi = pure ()-        | otherwise = do-            let !k2 = 2 * VU.unsafeIndex rtg i-                !x = VU.unsafeIndex v i-            a1 <- VUM.unsafeRead s k2-            if x > a1-                then do-                    VUM.unsafeWrite s k2 x-                    VUM.unsafeWrite s (k2 + 1) a1-                else do-                    a2 <- VUM.unsafeRead s (k2 + 1)-                    when (x > a2) $ VUM.unsafeWrite s (k2 + 1) x-            go (i + 1)---- | Top two of the four candidates per group (pairs already ordered m1 >= m2).-mergeTop2Range ::-    VUM.IOVector Double -> VUM.IOVector Double -> Int -> Int -> IO ()-mergeTop2Range a b lo hi = go lo-  where-    go !g-        | g >= hi = pure ()-        | otherwise = do-            let !g2 = 2 * g-            a1 <- VUM.unsafeRead a g2-            a2 <- VUM.unsafeRead a (g2 + 1)-            b1 <- VUM.unsafeRead b g2-            b2 <- VUM.unsafeRead b (g2 + 1)-            if b1 > a1-                then do-                    VUM.unsafeWrite a g2 b1-                    VUM.unsafeWrite a (g2 + 1) (max a1 b2)-                else VUM.unsafeWrite a (g2 + 1) (max a2 b1)-            go (g + 1)- addIntRange :: VUM.IOVector Int -> VUM.IOVector Int -> Int -> Int -> IO () addIntRange a b lo hi = go lo   where@@ -519,18 +381,7 @@             VUM.unsafeWrite a g (if isMin then min x y else max x y)             go (g + 1) ----------------------------------------------------------------------------------- Fused multi-reduction gather pass (nGroups above 'streamGroupCap')----------------------------------------------------------------------------------{- | One fused reduction of a multi-reduction GATHER pass: allocate the output-array, fold a contiguous group range (each group's rows via the shared-@valueIndices@ slice, accumulator in registers, one write per group), finalize.-Group ranges are disjoint across workers, so there is no merge and every-reduction reproduces the exact per-group fold order and formula of its unfused-gather kernel ('reduceParTyped') — results are bit-identical to running the-kernels separately, at any @-N@.--}+-- | One fused reduction of a multi-reduction gather pass. data GatherAgg     = forall s.         GatherAgg@@ -538,11 +389,7 @@         (s -> Int -> Int -> IO ())         (s -> IO Column) -{- | Build the fused gather reduction for one @(reduction, column)@ pair, or-'Nothing' when the pair cannot fuse (nullable/boxed columns, or var/std/top2,-which keep their per-expression gather kernels). @vis@/@offs@ are captured-lazily: nothing is forced until the pass actually runs.--}+-- | Build the fused gather reduction for a pair, or 'Nothing' if it can't fuse. mkGatherAgg ::     Int ->     VU.Vector Int ->@@ -586,21 +433,11 @@             step             (fmap fromUnboxedVector . VU.unsafeFreeze) -{- | Number of groups each fused-gather block hands to every reduction before-moving on: the block's @valueIndices@ slice stays hot in cache across the k-per-reduction loops.--}+-- | Groups each fused-gather block hands to every reduction before moving on. gatherBlock :: Int gatherBlock = 32 -{- | Execute all fused gather reductions in one traversal: workers own disjoint-contiguous group ranges (row-balanced, same policy as every gather kernel), and-walk them in 'gatherBlock'-group blocks running each reduction's fold on the-block. Deterministic and bit-identical to the unfused kernels (see-'GatherAgg'). Forces @valueIndices@ once, before the fan-out.--Pure w.r.t. its immutable inputs, so the 'unsafePerformIO' is safe.--}+-- | Execute all fused gather reductions in one traversal over group ranges. runGatherAggs ::     VU.Vector Int -> VU.Vector Int -> Int -> [GatherAgg] -> [Column] runGatherAggs vis offs nGroups aggs = unsafePerformIO $ do@@ -620,9 +457,7 @@     mapM snd opened {-# NOINLINE runGatherAggs #-} --- Monomorphic per-reduction gather folds; each replicates the exact per-group--- recurrence of its unfused kernel above ('countPar'/'sumPar'/'extremaPar'/--- 'meanPar'), so fused results are bit-identical.+-- Gather folds replicate their unfused kernels' per-group recurrence exactly.  gatherCount :: VU.Vector Int -> VUM.IOVector Int -> Int -> Int -> IO () gatherCount offs out gs ge = go gs

@@ -5,18 +5,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | The fused two-column moment kernel: one pass over @x@ and @y@ producing the-six sufficient statistics @(n, Sx, Sy, Sxx, Syy, Sxy)@ per group, from which the-whole correlation\/regression family (mean, variance, covariance, correlation,-OLS slope) is algebra requiring no further look at the rows.--The sequential and parallel passes live together because they must agree on-floating-point accumulation order. Moments are additive, so a row-range split-with a merge would be correct in exact arithmetic — but float addition is not-associative, so 'momentScatterPar' partitions by GROUP range instead: every-group accumulates start-to-finish inside a single worker, in the same order as-'momentScatter'. That is what makes the two byte-identical at any @-N@.--}+-- | The fused two-column moment kernel: six sufficient statistics per group. module DataFrame.Internal.Aggregation.Kernel.Moments (     Moments (..),     momentScatter,@@ -51,9 +40,7 @@ shouldPar :: Int -> Bool shouldPar = shouldParallelize parThreshold -{- | The additive moment sums of two columns, each an @nGroups@-length column:-@(n, Sx, Sy, Sxx, Syy, Sxy)@.--}+-- | The additive moment sums of two columns: @(n, Sx, Sy, Sxx, Syy, Sxy)@. data Moments = Moments     { mN :: Column     , mSx :: Column@@ -63,10 +50,7 @@     , mSxy :: Column     } -{- | One pass over two Double-coercible columns @x@ and @y@ filling the count and-five sums, collapsing the Q9 regression family's six folds into a single pass.-'Nothing' unless both columns are non-null unboxed Int/Double.--}+-- | One pass over two Double-coercible columns filling the count and sums. momentScatter :: VU.Vector Int -> Int -> Column -> Column -> Maybe Moments momentScatter g nGroups colX colY = do     xs <- cleanDoubleVector colX@@ -125,10 +109,7 @@         <*> VU.unsafeFreeze syy         <*> VU.unsafeFreeze sxy -{- | Parallel counterpart of 'momentScatter': one fused pass over both columns,-each group's six sums accumulated within one worker's range. Byte-identical to-'momentScatter'. 'Nothing' unless both columns are non-null unboxed Int/Double.--}+-- | Parallel counterpart of 'momentScatter'. momentScatterPar ::     VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -> Maybe Moments momentScatterPar vis offs nGroups colX colY@@ -142,8 +123,6 @@         pure (unsafePerformIO (momentPar vis offs nGroups xs ys caps bounds)) {-# NOINLINE momentScatterPar #-} --------------------------------------------------------------------------------- momentPar ::     VU.Vector Int ->     VU.Vector Int ->@@ -197,21 +176,7 @@         <*> (fromUnboxedVector <$> VU.unsafeFreeze syy)         <*> (fromUnboxedVector <$> VU.unsafeFreeze sxy) ----------------------------------------------------------------------------------- Streaming fused two-column moments (Q9)----------------------------------------------------------------------------------{- | Streaming counterpart of 'momentScatterPar': one fused pass over-@rowToGroup@ and the two TYPED value columns (Int values convert to Double-in-register — bit-identical to the @VU.map fromIntegral@ materialization it-replaces, with no 800MB intermediate column and no sequential conversion pass).-Each worker accumulates the six per-group sums over its contiguous row chunk in-original row order; partials merge in fixed worker order (counts exactly, the-five Double sums in chunk-major float order — deterministic at a fixed @-N@,-but a different summation order than the per-group gather kernel).-'Nothing' above 'streamGroupCap' or unless both columns are clean unboxed-Int/Double; the caller then keeps the gather path.--}+-- | Streaming counterpart of 'momentScatterPar' over typed value columns. momentStreamPar :: VU.Vector Int -> Int -> Column -> Column -> Maybe Moments momentStreamPar rtg nGroups colX colY     | nGroups <= 0 || nGroups > streamGroupCap = Nothing@@ -234,9 +199,7 @@         _ -> Nothing {-# NOINLINE momentStreamPar #-} -{- | Monomorphic entry points (see 'reduceParInt' for why the 'testEquality'-dispatch needs them).--}+-- | Monomorphic entry points for the 'testEquality' dispatch. momentStreamII ::     VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int -> Moments momentStreamII = momentStreamTyped@@ -285,9 +248,7 @@     VU.Vector a ->     VU.Vector b ->     Moments-{- The SPECIALIZE pragmas matter for the same reason as 'reduceParTyped': the-per-element @realToFrac@ must rewrite to @int2Double@/@id@ at a concrete type-or it goes through 'Rational' at runtime. -}+-- SPECIALIZE keeps per-element @realToFrac@ from going through 'Rational'. {-# SPECIALIZE momentStreamTyped ::     VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int -> Moments     #-}

@@ -5,24 +5,11 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | The group-range scatter-accumulate reduction kernel: reduces a value column-over the grouped layout @(valueIndices, offsets)@.--Sequential and parallel are the same algorithm at different row counts, so they-live together. 'scatterReducePar' cuts the GROUP axis into @caps@ ranges of-roughly equal row count and lets workers write disjoint slots of one shared-output — no per-worker accumulator, no merge — which keeps each group's-accumulation order identical to 'scatterReduce' and the results byte-identical-at any @-N@. Below 'parThreshold' it delegates to 'scatterReduce' directly.--Contrast "DataFrame.Internal.Aggregation.Kernel.Dense", which scatters off-@rowToGroup@ with no gather but needs a small dense group domain.--}+-- | The group-range scatter-accumulate reduction kernel over grouped layout. module DataFrame.Internal.Aggregation.Kernel.Scatter (     scatterReduce,     scatterReducePar,     maxMinusMinScatterPar,-    top2SndScatter,      -- * Group-range helpers     -- $shared@@ -33,7 +20,6 @@     streamGroupCap, ) where -import Control.Monad (when) import Control.Monad.ST (ST, runST) import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) import qualified Data.Vector.Unboxed as VU@@ -56,10 +42,7 @@     shouldParallelize,  ) -{- $shared-Also used by "DataFrame.Internal.Aggregation.Kernel.Moments", which partitions-the group axis the same way.--}+-- $shared Also used by the Moments kernel, which partitions groups alike.  -- | Whether to fan out at this row count. shouldPar :: Int -> Bool@@ -99,8 +82,6 @@     RMean -> fromUnboxedVector (meanScatter g nGroups v)     RVar -> fromUnboxedVector (varScatter False g nGroups v)     RStd -> fromUnboxedVector (varScatter True g nGroups v)-    RTop2Sum -> fromUnboxedVector (top2Scatter g nGroups v)-    RTop2Snd -> fromUnboxedVector (top2SndScatter g nGroups v) {-# INLINE reduceTyped #-}  countScatter :: VU.Vector Int -> Int -> VU.Vector Int@@ -233,41 +214,6 @@     VU.unsafeFreeze out {-# INLINE varScatter #-} -top2Scatter ::-    (VU.Unbox a, Real a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double-top2Scatter g nGroups v = runST $ do-    let ninf = negate (1 / 0) :: Double-    m1 <- VUM.replicate nGroups ninf-    m2 <- VUM.replicate nGroups ninf-    let n = VU.length v-        go !i-            | i >= n = pure ()-            | otherwise = do-                let !k = VU.unsafeIndex g i-                    !x = realToFrac (VU.unsafeIndex v i)-                a1 <- VUM.unsafeRead m1 k-                if x > a1-                    then do-                        VUM.unsafeWrite m1 k x-                        VUM.unsafeWrite m2 k a1-                    else do-                        a2 <- VUM.unsafeRead m2 k-                        when (x > a2) (VUM.unsafeWrite m2 k x)-                go (i + 1)-    go 0-    out <- VUM.new nGroups-    let fin !k-            | k >= nGroups = pure ()-            | otherwise = do-                a1 <- VUM.unsafeRead m1 k-                a2 <- VUM.unsafeRead m2 k-                let s = (if isInfinite a1 then 0 else a1) + (if isInfinite a2 then 0 else a2)-                VUM.unsafeWrite out k s-                fin (k + 1)-    fin 0-    VU.unsafeFreeze out-{-# INLINE top2Scatter #-}- groupRangeBounds :: VU.Vector Int -> Int -> Int -> VU.Vector Int groupRangeBounds offs nGroups caps = VU.create $ do     b <- VUM.new (caps + 1)@@ -304,12 +250,7 @@         _ -> Nothing {-# NOINLINE scatterReducePar #-} -{- | Monomorphic entry points: the 'testEquality' dispatch above only yields an-unsafe coercion, so a direct call to the polymorphic 'reduceParTyped' there-would stay at the abstract element type and never meet its SPECIALIZE rules;-calling through these fixed-type wrappers (the coercion lands on the argument)-does.--}+-- | Monomorphic entry points for the 'testEquality' dispatch above. reduceParInt ::     Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Int -> Column reduceParInt red vis offs nGroups v = reduceParTyped red vis offs nGroups v intIdent@@ -346,10 +287,7 @@     VU.Vector a ->     Idents a ->     Column-{- The SPECIALIZE pragmas matter: without them the @realToFrac@ in the-mean/var/top2 kernels survives to runtime as a dictionary call through-'Rational' (the Double->Double/Int->Double rewrite rules only fire once the-type is concrete), costing ~4x on the whole pass. -}+-- SPECIALIZE lets @realToFrac@ resolve at a concrete type (~4x faster). {-# SPECIALIZE reduceParTyped ::     Reduction ->     VU.Vector Int ->@@ -386,18 +324,9 @@                     (unsafePerformIO (varPar False vis offs nGroups v caps bounds))             RStd ->                 fromUnboxedVector (unsafePerformIO (varPar True vis offs nGroups v caps bounds))-            RTop2Sum -> fromUnboxedVector (unsafePerformIO (top2Par vis offs nGroups v caps bounds))-            RTop2Snd ->-                fromUnboxedVector (unsafePerformIO (top2SndPar vis offs nGroups v caps bounds)) {-# INLINEABLE reduceParTyped #-} -{- | For each group in @[gs, ge)@, fold the group's rows (in @valueIndices@-order, i.e. ascending original-row order) into an accumulator held in-registers, then hand the final accumulator to @done@ exactly once. Keeping the-running state out of memory leaves one write per group instead of a-read-modify-write per row; the per-group fold order is unchanged, so results-stay byte-identical to the row-wise variant.--}+-- | Fold each group's rows into a register accumulator, then hand it on. overGroupsAcc ::     VU.Vector Int ->     VU.Vector Int ->@@ -526,8 +455,7 @@ varPar takeSqrt vis offs nGroups v caps bounds = do     out <- VUM.replicate nGroups (0 :: Double)     parallelBounds_ caps bounds $ \gs ge ->-        -- Per-group Welford state (count, mean, M2) carried in registers; the-        -- update order per group is the same ascending row order as before.+        -- Per-group Welford state (count, mean, M2) is carried in registers.         let grp !g                 | g >= ge = pure ()                 | otherwise = do@@ -550,90 +478,7 @@     VU.unsafeFreeze out {-# INLINE varPar #-} -top2Par ::-    (VU.Unbox a, Real a) =>-    VU.Vector Int ->-    VU.Vector Int ->-    Int ->-    VU.Vector a ->-    Int ->-    VU.Vector Int ->-    IO (VU.Vector Double)-top2Par vis offs nGroups v caps bounds = do-    let ninf = negate (1 / 0) :: Double-    out <- VUM.replicate nGroups (0 :: Double)-    parallelBounds_ caps bounds $ \gs ge ->-        -- The (largest, second-largest) pair carried in registers per group.-        let grp !g-                | g >= ge = pure ()-                | otherwise = do-                    let !e = VU.unsafeIndex offs (g + 1)-                        inner !pos !a1 !a2-                            | pos >= e =-                                (if isInfinite a1 then 0 else a1)-                                    + (if isInfinite a2 then 0 else a2)-                            | otherwise =-                                let !x = realToFrac (VU.unsafeIndex v (VU.unsafeIndex vis pos))-                                 in if x > a1-                                        then inner (pos + 1) x a1-                                        else inner (pos + 1) a1 (max a2 x)-                        !res = inner (VU.unsafeIndex offs g) ninf ninf-                    VUM.unsafeWrite out g res-                    grp (g + 1)-         in grp gs-    VU.unsafeFreeze out-{-# INLINE top2Par #-}--{- | Second-largest value per group: the same (largest, second-largest)-register pair as 'top2Par', finalized to the second max alone. Size-1 groups-finalize the @-inf@ seed to NaN (documented; see-'DataFrame.Internal.AggKernel.top2SndScatter').--}-top2SndPar ::-    (VU.Unbox a, Real a) =>-    VU.Vector Int ->-    VU.Vector Int ->-    Int ->-    VU.Vector a ->-    Int ->-    VU.Vector Int ->-    IO (VU.Vector Double)-top2SndPar vis offs nGroups v caps bounds = do-    let ninf = negate (1 / 0) :: Double-    out <- VUM.replicate nGroups (0 :: Double)-    parallelBounds_ caps bounds $ \gs ge ->-        -- The (largest, second-largest) pair carried in registers per group.-        let grp !g-                | g >= ge = pure ()-                | otherwise = do-                    let !e = VU.unsafeIndex offs (g + 1)-                        inner !pos !a1 !a2-                            | pos >= e = if isInfinite a2 then 0 / 0 else a2-                            | otherwise =-                                let !x = realToFrac (VU.unsafeIndex v (VU.unsafeIndex vis pos))-                                 in if x > a1-                                        then inner (pos + 1) x a1-                                        else inner (pos + 1) a1 (max a2 x)-                        !res = inner (VU.unsafeIndex offs g) ninf ninf-                    VUM.unsafeWrite out g res-                    grp (g + 1)-         in grp gs-    VU.unsafeFreeze out-{-# INLINE top2SndPar #-}------------------------------------------------------------------------------------ Parallel fused max(a) - min(b) (Q7 at wide group domains)----------------------------------------------------------------------------------{- | Fused @max a - min b@ over the group-range layout: ONE traversal of-@valueIndices@ accumulating both extrema, parallel by disjoint group range with-no cross-worker merge. min/max are order-independent, so the result is-byte-identical to running the two gather extrema passes separately; the fusion-halves the index traffic. 'Nothing' below the parallel threshold or unless both-columns are clean unboxed and same-typed (Int/Int keeps the Int result of the-interpreter; Double/Double the Double one) — the caller then keeps its two-pass-fallback.--}+-- | Fused @max a - min b@ over the group-range layout in one traversal. maxMinusMinScatterPar ::     VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -> Maybe Column maxMinusMinScatterPar vis offs nGroups ca cb@@ -654,9 +499,7 @@     !bounds = groupRangeBounds offs nGroups caps {-# NOINLINE maxMinusMinScatterPar #-} -{- | Monomorphic entry points (see 'reduceParInt' for why the 'testEquality'-dispatch needs them).--}+-- | Monomorphic entry points for the 'testEquality' dispatch. maxMinusMinParInt ::     VU.Vector Int ->     VU.Vector Int ->@@ -728,8 +571,7 @@ maxMinusMinPar maxSeed minSeed vis offs nGroups va vb caps bounds = do     out <- VUM.new nGroups     parallelBounds_ caps bounds $ \gs ge ->-        -- Both extrema carried in registers per group; one traversal of the-        -- shared index slice reads both value columns.+        -- Both extrema live in registers; one traversal reads both columns.         let grp !g                 | g >= ge = pure ()                 | otherwise = do@@ -748,71 +590,10 @@          in grp gs     VU.unsafeFreeze out ----------------------------------------------------------------------------------- Streaming (rowToGroup-scatter) kernels: no valueIndices, no placement pass----------------------------------------------------------------------------------{- | Group-count cap for the FUSED streaming rtg-scatter kernels-('DataFrame.Internal.Aggregation.Kernel.Moments.momentStreamPar',-'DataFrame.Internal.Aggregation.Kernel.Fused.runFusedAggs'). Above-'DataFrame.Internal.Grouping.Direct.directThreshold' the per-worker accumulator-arrays overflow cache, so a SINGLE streaming reduction loses to a gather pass —-the per-expression dispatch keeps that threshold. A fused multi-reduction pass-amortizes those misses across all its reductions AND avoids the deferred-@valueIndices@ placement entirely, which flips the comparison (measured at 1e6-groups / 1e8 rows on -N16: rowToGroup 0.4s + fused 3-sum stream 1.4s, against-placement 1.1s + fused gather 1.05s), so the fused cap extends to-'directGroupThreshold' — every direct-grouped frame can stream. Wider-groupings are necessarily hash-path (eager @valueIndices@) and use the fused-GATHER kernel ('DataFrame.Internal.Aggregation.Kernel.Fused.runGatherAggs')-instead. Memory: @capabilities * nGroups@ words per accumulator array, at most-~128MB transient at -N16.--}+-- | Group-count cap for the fused streaming scatter kernels. streamGroupCap :: Int streamGroupCap = 1048576 -{- | Near-equal contiguous slices of the group domain for parallel merges.-Below 4096 groups (or single-capability) the merge stays on one thread.--}+-- | Near-equal contiguous slices of the group domain for parallel merges. groupSlices :: Int -> [(Int, Int)] groupSlices = chunksFor 4096--{- | Second-largest value per group: the same (largest, second-largest)-accumulator pair as 'top2Scatter', but the finalize returns the second max-alone. A group of size 1 (or 0) leaves the @-inf@ seed in the second slot, so-its output is NaN — documented behaviour (the db-benchmark Q8 data has no-size-1 @id6@ groups). Like 'top2Scatter''s @-inf -> 0@ guard, an actual-infinite data value in the second slot is indistinguishable from the seed.--}-top2SndScatter ::-    (VU.Unbox a, Real a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double-top2SndScatter g nGroups v = runST $ do-    let ninf = negate (1 / 0) :: Double-    m1 <- VUM.replicate nGroups ninf-    m2 <- VUM.replicate nGroups ninf-    let n = VU.length v-        go !i-            | i >= n = pure ()-            | otherwise = do-                let !k = VU.unsafeIndex g i-                    !x = realToFrac (VU.unsafeIndex v i)-                a1 <- VUM.unsafeRead m1 k-                if x > a1-                    then do-                        VUM.unsafeWrite m1 k x-                        VUM.unsafeWrite m2 k a1-                    else do-                        a2 <- VUM.unsafeRead m2 k-                        when (x > a2) (VUM.unsafeWrite m2 k x)-                go (i + 1)-    go 0-    out <- VUM.new nGroups-    let fin !k-            | k >= nGroups = pure ()-            | otherwise = do-                a2 <- VUM.unsafeRead m2 k-                VUM.unsafeWrite out k (if isInfinite a2 then 0 / 0 else a2)-                fin (k + 1)-    fin 0-    VU.unsafeFreeze out-{-# INLINE top2SndScatter #-}

@@ -6,11 +6,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | The aggregation fast-path planner. 'planAgg' recognises a supported-aggregate shape over a clean unboxed Int/Double column and returns an 'AggPlan';-'planMoments' recognises the six-fold regression shape and returns a-'MomentPlan'. Planning only — the kernels live under "Kernel".--}+-- | The aggregation fast-path planner: recognises supported aggregate shapes. module DataFrame.Internal.Aggregation.Plan (     AggPlan (..),     planAgg,@@ -23,6 +19,7 @@ import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) import qualified Data.Vector.Unboxed as VU +import DataFrame.Internal.Aggregation.Kernel.Distinct (distinctSupported) import DataFrame.Internal.Aggregation.Reduction (Reduction (..)) import DataFrame.Internal.Column (Column (..)) import DataFrame.Internal.DataFrame (@@ -34,14 +31,11 @@     AggStrategy (..),     BinaryOp (binaryCommutative, binaryName),     Expr (..),-    UExpr (..),+    SomeExpr (..),  ) import Type.Reflection (Typeable, typeRep) -{- | The plan 'planAgg' produces for a recognised output expression. The median-plan carries only the column name (the holistic grouped sort lives in the-operations layer, where @vector-algorithms@ is available).--}+-- | The plan 'planAgg' produces for a recognised output expression. data AggPlan     = -- | A single scatter reduction over one named column.       PlanScatter Reduction T.Text@@ -49,13 +43,12 @@       PlanMaxMinusMin T.Text T.Text     | -- | Holistic median over one named column.       PlanMedian T.Text+    | -- | Exact distinct count over one named column.+      PlanCountDistinct T.Text -{- | Inspect a named output expression; return @Just plan@ on a recognised shape-over a present clean column, else 'Nothing'. Nullable or non-Int/Double columns-are rejected here so the scatter only sees a clean unboxed vector.--}-planAgg :: GroupedDataFrame -> UExpr -> Maybe AggPlan-planAgg gdf (UExpr (expr :: Expr a)) = case expr of+-- | Plan a named output expression, or 'Nothing' if the shape is unsupported.+planAgg :: GroupedDataFrame -> SomeExpr -> Maybe AggPlan+planAgg gdf (SomeExpr (expr :: Expr a)) = case expr of     Agg (FoldAgg tag _ _) (Col name) -> foldPlan tag name     Agg (MergeAgg tag _ _ _ _) (Col name) -> mergePlan tag name     Agg (CollectAgg tag _) (Col name) -> collectPlan tag name@@ -76,6 +69,10 @@     mergePlan tag name = case tag of         "mean" -> outputType @Double >> require name (PlanScatter RMean name)         "count" -> outputType @Int >> require name (PlanScatter RCount name)+        "countDistinct" -> do+            outputType @Int+            c <- getColumn name (fullDataframe gdf)+            if distinctSupported c then Just (PlanCountDistinct name) else Nothing         _ -> Nothing     outputType :: forall t. (Typeable t) => Maybe ()     outputType = case testEquality (typeRep @a) (typeRep @t) of@@ -84,8 +81,6 @@     collectPlan tag name = case tag of         "stddev" -> require name (PlanScatter RStd name)         "variance" -> require name (PlanScatter RVar name)-        "top2Sum" -> require name (PlanScatter RTop2Sum name)-        "top2Snd" -> require name (PlanScatter RTop2Snd name)         "median" -> require name (PlanMedian name)         _ -> Nothing     require name plan = colUnboxedNumeric name >> Just plan@@ -105,10 +100,7 @@                 Nothing -> False     _ -> False -{- | A recognised moment (Q9 regression) aggregate group: six output columns that-form the sufficient statistics of two base columns @x@ and @y@. The caller runs-'momentScatter' once and binds each output name to a field of the result.--}+-- | A recognised moment aggregate group: six sufficient-statistic columns. data MomentPlan = MomentPlan     { mpColX :: T.Text     , mpColY :: T.Text@@ -120,20 +112,14 @@     , 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).--}+-- | The shape of a sum's argument: linear in one base column, or a product. data Term     = Lin T.Text     | Prod T.Text T.Text     deriving (Eq, Ord, Show) -{- | Recognise the moment shape across a whole @aggregate@ list: exactly-@count@, @sum(x)@, @sum(y)@, @sum(x*x)@, @sum(y*y)@, @sum(x*y)@ over two distinct-clean unboxed base columns. 'Nothing' on any other set.--}-planMoments :: GroupedDataFrame -> [(T.Text, UExpr)] -> Maybe MomentPlan+-- | Recognise the moment shape across a whole @aggregate@ list.+planMoments :: GroupedDataFrame -> [(T.Text, SomeExpr)] -> Maybe MomentPlan planMoments gdf aggs     | length aggs /= 6 = Nothing     | otherwise = do@@ -174,33 +160,30 @@     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+classify :: M.Map T.Text SomeExpr -> (T.Text, SomeExpr) -> Maybe (T.Text, Role)+classify exprs (name, SomeExpr expr) = case expr of     Agg (MergeAgg "count" _ _ _ _) _ -> Just (name, RoleN)-    Agg (FoldAgg "sum" _ _) arg -> (\t -> (name, termRole t)) <$> resolveTerm exprs (UExpr arg)+    Agg (FoldAgg "sum" _ _) arg -> (\t -> (name, termRole t)) <$> resolveTerm exprs (SomeExpr arg)     _ -> Nothing  termRole :: Term -> Role termRole (Lin a) = RoleLin a termRole (Prod a b) = RoleProd a b -{- | Resolve a (sum-argument) expression to its 'Term'. 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+-- | Resolve a sum-argument expression to its 'Term'.+resolveTerm :: M.Map T.Text SomeExpr -> SomeExpr -> Maybe Term resolveTerm exprs = go (8 :: Int)   where     go 0 _ = Nothing-    go fuel (UExpr e) = case e of+    go fuel (SomeExpr e) = case e of         Col nm -> case M.lookup nm exprs of             Just ue -> go (fuel - 1) ue             Nothing -> Just (Lin nm)-        Unary _ inner -> go (fuel - 1) (UExpr inner)+        Unary _ inner -> go (fuel - 1) (SomeExpr 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)+                Lin a <- go (fuel - 1) (SomeExpr l)+                Lin b <- go (fuel - 1) (SomeExpr r)                 Just (sortProd a b)         _ -> Nothing @@ -210,9 +193,7 @@     | 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.--}+-- | Find the unordered pair of base columns that the linear sums name. pickBaseColumns :: [(T.Text, Role)] -> Maybe (T.Text, T.Text) pickBaseColumns roles =     case lins of

@@ -1,12 +1,6 @@ {-# LANGUAGE ScopedTypeVariables #-} -{- | The vocabulary shared by the aggregation planner and the aggregation-kernels: the set of recognised reductions, and the admission gate that decides-whether a column may enter a Double-valued fast path.--This module holds no kernel and no policy, so the planner can depend on it-without depending on an implementation.--}+-- | Vocabulary shared by the aggregation planner and kernels. module DataFrame.Internal.Aggregation.Reduction (     Reduction (..),     cleanDoubleVector,@@ -20,10 +14,7 @@  ) import DataFrame.Internal.Column.Conversion (toDoubleVector) -{- | A recognised fast-path reduction over a single value column. The element-type (Int vs Double) is resolved at scatter time; sum/min/max preserve the-column's element type, everything else produces a Double column.--}+-- | A recognised fast-path reduction over a single value column. data Reduction     = RSum     | RCount@@ -32,8 +23,6 @@     | RMean     | RStd     | RVar-    | RTop2Sum-    | RTop2Snd     deriving (Eq, Show)  cleanDoubleVector :: Column -> Maybe (VU.Vector Double)

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

@@ -1,10 +1,7 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE ScopedTypeVariables #-} -{- | Stable rank of a set of group representatives by ascending unsigned hash-order. Shared by the sequential and parallel group-by canonical-ordering steps-so they stay bit-for-bit identical. @O(ng)@ stable LSD radix sort.--}+-- | Stable rank of group representatives by ascending unsigned hash. module DataFrame.Internal.Algorithms.Rank.Radix (     rankByHash,     sortKey,@@ -17,10 +14,7 @@ import qualified Data.Vector.Unboxed.Mutable as VUM import Data.Word (Word64) -{- | Unsigned sort key of a hash: ascending 'Word64' order of @sortKey h@ equals-ascending signed-'Int' order of @h@. Reinterpreted to 'Int' for the byte-wise-radix passes (the byte mask makes the sign extension irrelevant).--}+-- | Unsigned sort key of a hash, ordering like signed 'Int'. sortKey :: Int -> Int sortKey h = fromIntegral (fromIntegral h + 0x8000000000000000 :: Word64) {-# INLINE sortKey #-}

@@ -1,17 +1,14 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE ScopedTypeVariables #-} -{- | Parallel stable sort of row indices by ascending unsigned order of a per-row-'Int' hash, used by the join build side. A counting sort buckets rows into-key-ordered partitions that workers LSD-radix-sort in parallel, with no merge step.--}+-- | Parallel stable sort of row indices by unsigned per-row hash. module DataFrame.Internal.Algorithms.Sort.Radix.Parallel (     parSortByHash,     parSortThreshold, ) where  import Control.Concurrent (getNumCapabilities)-import Control.Monad (forM_, when)+import Control.Monad (when) import Data.Bits (countLeadingZeros, unsafeShiftR, (.&.)) import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM@@ -25,15 +22,11 @@  ) 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.--}+-- | Below this many rows the sequential radix path is used instead. parSortThreshold :: Int parSortThreshold = 500000 -{- | Top-bits partition index of a hash: the high @64 - shift@ bits of its-unsigned 'sortKey'. Ascending partition order equals ascending key order.--}+-- | Top-bits partition index of a hash, in ascending key order. partIx :: Int -> Int -> Int partIx shift h = fromIntegral ((fromIntegral (sortKey h) :: Word64) `unsafeShiftR` shift) {-# INLINE partIx #-}@@ -52,9 +45,7 @@ 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.--}+-- | Parallel stable sort of @[0, n)@ by ascending unsigned hash order. parSortByHash :: Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) parSortByHash n hashes     | n <= 1 =@@ -64,13 +55,7 @@     | 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)@.--}+-- | Stable LSD radix sort of @[0, n)@ by ascending 'sortKey' of their hash. seqSortByHash :: Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) seqSortByHash n hashes = unsafePerformIO $ do     keysA <- VUM.new n@@ -88,10 +73,7 @@     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.--}+-- | Run all eight stable 8-bit LSD passes, ping-ponging between buffers. radixPasses ::     Int ->     VUM.IOVector Int ->@@ -101,12 +83,7 @@     IO () radixPasses = radixPassesN 8 -{- | Run the first @np@ stable 8-bit LSD passes (bits @0 .. 8*np-1@),-ping-ponging between the buffer pairs. For odd @np@ the sorted order lands in-@(keysB, orderB)@, for even in @(keysA, orderA)@. Callers whose rows share-their top bytes (per-partition sorts partitioned on the top byte) can pass-@np = 7@: the eighth pass would be a stable identity copy.--}+-- | Run the first @np@ stable 8-bit LSD passes, ping-ponging between buffers. radixPassesN ::     Int ->     Int ->@@ -159,10 +136,6 @@             | otherwise = pass (8 * k) keysB orderB keysA orderA >> run (k + 1)     run 0 ----------------------------------------------------------------------------------- Parallel path: counting-sort partition, then per-partition sort in parallel--------------------------------------------------------------------------------- parSortByHashIO :: Int -> VU.Vector Int -> IO (VU.Vector Int, VU.Vector Int) parSortByHashIO n hashes = do     caps <- getNumCapabilities@@ -176,17 +149,7 @@     sortedHashes <- VU.unsafeFreeze outKeys     pure (sortedHashes, order) -{- | Bucket every row index into its top-bits partition by a counting sort.-Returns the exclusive prefix sum @partStart@ (length @p+1@, @partStart[p] == n@),-the row indices laid out partition-by-partition in ascending key order, and-each sorted position's hash in the same layout (so downstream passes read-hashes sequentially instead of a random @hashes[row]@ gather per row).--Runs chunked across capabilities: per-chunk partition histograms are prefix-summed (in chunk order) into disjoint per-chunk write cursors, so the scatter-threads never contend and each partition keeps its rows in ascending original-row order — bit-for-bit the sequential counting sort's layout.--}+-- | Bucket every row index into its top-bits partition by counting sort. partitionRows ::     Int ->     VU.Vector Int ->@@ -197,8 +160,7 @@     caps <- getNumCapabilities     let chunks = rowChunks caps n     cursors <- forkJoin [histChunk hashes p shift lo hi | (lo, hi) <- chunks]-    -- Exclusive prefix over partitions (outer) and chunks (inner): partStart-    -- from the totals, and each chunk's histogram rewritten into its cursor.+    -- Exclusive prefix over partitions and chunks; histograms become cursors.     partStartM <- VUM.new (p + 1)     let seed !pp !acc             | pp >= p = VUM.unsafeWrite partStartM p acc@@ -248,9 +210,7 @@                 go (i + 1)     go lo -{- | Scatter one row chunk into the partitioned layout using the chunk's-pre-summed cursor (disjoint write regions per chunk, no contention).--}+-- | Scatter one row chunk into the partitioned layout using its cursor. scatterChunk ::     VU.Vector Int ->     Int ->@@ -273,19 +233,7 @@             VUM.unsafeWrite cursor pp (pos + 1)             go (i + 1) -{- | Stable-sort each partition by full key, writing sorted original indices-into @outOrder@ and their hashes into @outKeys@ at the partition's slot range.-Forks @caps@ workers that pull partition indices off a shared atomic counter.-Within a partition the counting sort already left rows in ascending original-order, so the LSD radix sort's stability reproduces the global @(key, row)@-order. Partitions below two elements are already sorted (counting sort kept-original order) and are copied directly.--@partHashes@ is the partition-layout hash vector from 'partitionRows', so-seeding reads hashes sequentially; only 7 LSD passes run (the top byte is the-partition byte, constant within a partition), and the sorted hash is recovered-from the sort key ('sortKey' is self-inverse) instead of a random gather.--}+-- | Stable-sort each partition by full key into the output slots. sortPartitions ::     Int ->     Int ->

@@ -1,17 +1,4 @@-{- |-Umbrella re-export of the column implementation. The module is split by-concern:--* "DataFrame.Internal.Column.Types" — type-level machinery ('Rep', 'SBool', ...)-* "DataFrame.Internal.Column.Base" — the 'Column' GADT and core definitions-* "DataFrame.Internal.Column.Bitmap" — the validity bitmap and its helpers-* "DataFrame.Internal.Column.Properties" — predicates and introspection-* "DataFrame.Internal.Column.Conversion" — vector\/list conversions-* "DataFrame.Internal.Column.Operations" — bulk transformations--Import this module to get the whole surface; import a submodule directly when-you only need one layer.--}+-- | Umbrella re-export of the column implementation, split by concern. module DataFrame.Internal.Column (     module DataFrame.Internal.Column.Base,     module DataFrame.Internal.Column.Bitmap,

@@ -15,16 +15,7 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} -{- |-Core column definitions: the type-erased 'Column' GADT and its mutable/typed-companions, the 'Columnable' constraint, and the representation-level-primitives ('materializePacked', 'materializeMerged') that the non-orphan-'Show'/'Eq' instances depend on.--Predicates live in "DataFrame.Internal.Column.Properties", vector/list-conversions in "DataFrame.Internal.Column.Conversion", and bulk transformations-in "DataFrame.Internal.Column.Operations".--}+-- | Core column definitions: the type-erased 'Column' GADT and companions. module DataFrame.Internal.Column.Base where  import qualified Data.Vector as VB@@ -62,24 +53,16 @@     , SBoolI (FloatingTypes a)     ) -{- | Type-erased column GADT. Pattern-matching on the constructor recovers the-representation; nullability is an optional bit-packed 'Bitmap' (@Nothing@ = no-nulls, @Just bm@ = bit @i@ set iff row @i@ is valid).--}+-- | Type-erased column GADT with optional bit-packed null bitmap. data Column where     BoxedColumn :: (Columnable a) => Maybe Bitmap -> VB.Vector a -> Column     UnboxedColumn ::         (Columnable a, VU.Unbox a) => Maybe Bitmap -> VU.Vector a -> Column-    -- TODO: mchavinda - investigate splitting this into a separate intermediate-    --                   representation.-    -- Efficient intermediate formats.z-    -- Bit-packed Text: shared UTF-8 byte buffer + row offsets + optional bitmap;-    -- Text is materialized on demand. Only CSV ingest emits this; user-built-    -- Text columns stay 'BoxedColumn'.+    -- TODO: mchavinda - split this into a separate intermediate representation.++    -- Bit-packed Text: shared UTF-8 buffer, offsets, optional bitmap; CSV only.     PackedText :: Maybe Bitmap -> {-# UNPACK #-} !PackedTextData -> Column-    -- A join's same-named non-key column pair ('mkMergedColumns'): both sides-    -- keep their native (packed/dict/unboxed) representation; per-row 'These'-    -- values only materialize on element access ('materializeMerged').+    -- A join's same-named non-key column pair; both sides keep native form.     MergedColumn :: !Column -> !Column -> Column  instance Show Column where@@ -103,9 +86,7 @@          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.--}+-- | Compare two nullable boxed columns element by element, skipping nulls. eqBoxedCols ::     (Eq a) => Maybe Bitmap -> VB.Vector a -> Maybe Bitmap -> VB.Vector a -> Bool eqBoxedCols bm1 a bm2 b@@ -150,9 +131,7 @@     (==) 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'.--}+-- | Byte-slice equality of two packed-text columns, skipping null slots. eqPackedCols ::     Maybe Bitmap -> PackedTextData -> Maybe Bitmap -> PackedTextData -> Bool eqPackedCols bm1 p1 bm2 p2@@ -172,18 +151,12 @@         nullB = maybe False (\bm -> not (bitmapTestBit bm i)) bm2 {-# INLINE eqPackedCols #-} -{- | A mutable companion struct to dataframe columns.--Used mostly as an intermediate structure for I/O.--}+-- | A mutable companion to dataframe columns, used as an I/O intermediate. data MutableColumn where     MBoxedColumn :: (Columnable a) => VBM.IOVector a -> MutableColumn     MUnboxedColumn :: (Columnable a, VU.Unbox a) => VUM.IOVector a -> MutableColumn -{- | A wrapper around the type-erased 'Column' carrying a phantom element type,-used to type-check expressions. The phantom is not guaranteed to match the-underlying vector's type.--}+-- | A type-erased 'Column' with a phantom element type, for expressions. data TypedColumn a where     TColumn :: (Columnable a) => Column -> TypedColumn a @@ -199,9 +172,7 @@ unwrapTypedColumn :: TypedColumn a -> Column unwrapTypedColumn (TColumn value) = value -{- | Decode a 'PackedText' into a @BoxedColumn Text@ (bit-identical to-materializing at freeze). Identity on every other column.--}+-- | Decode a 'PackedText' into a @BoxedColumn Text@; identity otherwise. materializePacked :: Column -> Column materializePacked (PackedText bm p) =     BoxedColumn bm (VB.generate (packedLength p) (packedIndexText p))@@ -215,9 +186,7 @@ columnBitmap (PackedText bm _) = bm columnBitmap (MergedColumn _ _) = Nothing -{- | A class for converting a vector to a column of the appropriate type.-Given each Rep we tell the `toColumnRep` function which Column type to pick.--}+-- | Selects which Column type to build from a vector for each 'Rep'. class ColumnifyRep (r :: Rep) a where     toColumnRep :: VB.Vector a -> Column @@ -276,9 +245,7 @@     PackedText (fmap (bitmapSlice 0 n) bm) (packedTake n p) {-# INLINE takeColumn #-} -{- | Merge two columns using `These`. O(1): the sides are kept in their-native representation and 'These' values materialize on element access.--}+-- | Merge two columns using 'These' in O(1), keeping native representations. mkMergedColumns :: Column -> Column -> Column mkMergedColumns = MergedColumn {-# INLINE mkMergedColumns #-}@@ -294,10 +261,7 @@     materializeMerged (MergedColumn (takeColumn 1 a) (takeColumn 1 b)) mergedHead c = c -{- | The eager element-wise merge ('These' per row, boxed). Bitmaps are-honored for every representation pair: a null side yields 'This'/'That',-both-null is an error (the join kernels never produce such a row).--}+-- | The eager element-wise merge ('These' per row, boxed). mergeEager :: Column -> Column -> Column mergeEager colA colB = case (colA, colB) of     (MergedColumn a b, _) -> mergeEager (mergeEager a b) colB

@@ -14,13 +14,13 @@     (.&.),     (.|.),  )-import Data.List (foldl')+import qualified Data.List as L import Data.Maybe (fromMaybe, isNothing) import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM import Data.Word (Word8) --- | A bit-packed validity bitmap. Bit @i@ = 1 means row @i@ is valid (not null).+-- | A bit-packed validity bitmap; bit @i@ = 1 means row @i@ is valid. type Bitmap = VU.Vector Word8  -- | A bitmap attached to its row counts so we can splice it.@@ -32,11 +32,7 @@ vRowCount :: Validity -> Int vRowCount (Validity _ n) = n -{- | Test whether row @i@ is valid (not null) in a bitmap.--The bit-level arithmetric is dense but read this as:-`shiftR` 3 is equivalent to `div` 8, and `.&. 7` is equivalent to `mod` 8.--}+-- | Test whether row @i@ is valid (not null) in a bitmap. bitmapTestBit :: Bitmap -> Int -> Bool bitmapTestBit bm i = testBit (VU.unsafeIndex bm (i `shiftR` 3)) (i .&. 7) {-# INLINE bitmapTestBit #-}@@ -61,9 +57,7 @@                 pure mv {-# INLINE allValidBitmap' #-} -{- | Build a bitmap from a @VU.Vector Word8@ validity vector-(1 = valid, 0 = null), as produced by Arrow / Parquet decoders.--}+-- | Build a bitmap from a validity vector (1 = valid, 0 = null). buildBitmapFromValid :: VU.Vector Word8 -> Bitmap buildBitmapFromValid valid =     let n = VU.length valid@@ -75,11 +69,9 @@                      in if idx < n && VU.unsafeIndex valid idx /= 0                             then setBit acc bit                             else acc-             in foldl' setBitIf (0 :: Word8) [0 .. 7]+             in L.foldl' setBitIf (0 :: Word8) [0 .. 7] -{- | Build a bitmap from a list of null-row indices.-@nullIdxs@ are the positions that are NULL.--}+-- | Build a bitmap from a list of null-row indices. buildBitmapFromNulls :: Int -> [Int] -> Bitmap buildBitmapFromNulls n idxs = buildBitmapFromNulls' n (VU.fromList idxs) @@ -109,11 +101,7 @@                 VU.generate n $                     \i -> if bitmapTestBit bm (start + i) then 1 else 0 -{- | Count the set bits among the first @n@ bits of a bitmap. A bitmap does-not know the length of the column it describes, and 'bitmapSlice' keeps whole-bytes on its aligned path, so the bits past @n@ may still describe rows-outside the slice.--}+-- | Count the set bits among the first @n@ bits of a bitmap. popCountUpTo :: Int -> Bitmap -> Int popCountUpTo n bm = whole + partial   where@@ -126,6 +114,11 @@             popCount (VU.unsafeIndex bm fullBytes .&. ((1 `shiftL` rest) - 1)) {-# INLINE popCountUpTo #-} +-- | Whether any of the first @n@ rows is null.+bitmapHasNulls :: Int -> Bitmap -> Bool+bitmapHasNulls n bm = popCountUpTo n bm < n+{-# INLINE bitmapHasNulls #-}+ -- | Concatenate two bitmaps covering @n1@ and @n2@ rows respectively. bitmapConcat :: Int -> Bitmap -> Int -> Bitmap -> Bitmap bitmapConcat n1 bm1 n2 bm2 =@@ -139,9 +132,7 @@ andBitmaps :: Bitmap -> Bitmap -> Bitmap andBitmaps = VU.zipWith (.&.) -{- | Splice chunk bitmaps end to end at the bit level. 'Nothing' if no chunk-carries a bitmap; chunks without one count as all-valid otherwise.--}+-- | Splice chunk bitmaps end to end at the bit level. concatValidity :: [Validity] -> Maybe Bitmap concatValidity parts     | all (isNothing . vBitmap) parts = Nothing

@@ -3,10 +3,7 @@ {-# 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.--}+-- | Mutable, growable column builders for high-throughput ingest. module DataFrame.Internal.Column.Builder (     ColumnBuilder (..),     NumBuilder,@@ -26,7 +23,9 @@     appendTextSliceFromPtr,     freezeTextChunk,     concatColumns,+    concatColumnsIO,     mergeTextChunks,+    mergeTextChunksIO, ) where  import qualified Data.Text as T@@ -48,14 +47,13 @@ import DataFrame.Internal.Column.Merge (     TextChunk (..),     concatColumns,+    concatColumnsIO,     mergeTextChunks,+    mergeTextChunksIO,  ) import Foreign.Ptr (Ptr) -{- | Operations shared by all column builders.--NB: Do not use after freezing--}+-- | Operations shared by all column builders; do not use after freezing. class ColumnBuilder b where     -- | Append a null row (sentinel payload + invalid bit).     appendNull :: b s -> ST s ()@@ -72,9 +70,7 @@ cAnyNull = 1 cBytes = 2 -{- | Builder for unboxed numeric payloads ('Int', 'Double', ...). 'nbNull'-is the sentinel written into null slots.--}+-- | Builder for unboxed numeric payloads, with a null-slot sentinel. data NumBuilder a s = NumBuilder     { nbNull :: !a     , nbCounters :: !(VUM.MVector s Int)@@ -174,11 +170,7 @@     | VUM.length mv - n <= n `shiftR` 2 = VU.unsafeFreeze (VUM.slice 0 n mv)     | otherwise = VU.freeze (VUM.slice 0 n mv) -{- | Builder for 'Text' columns.--Representation is packed. I.e all field bytes go into one exponentially-grown byte array with rows recorded as offsets.--}+-- | Builder for 'Text' columns using one packed, growable byte array. data TextBuilder s = TextBuilder     { tbCounters :: !(VUM.MVector s Int)     , tbArrays :: !(STRef s (TextArrays s))@@ -268,9 +260,7 @@     writeSTRef (tbArrays b) arrs     pure arrs -{- | Freeze a 'TextBuilder' into a raw 'TextChunk' for later merging-('mergeTextChunks').--}+-- | Freeze a 'TextBuilder' into a raw 'TextChunk' for later merging. freezeTextChunk :: TextBuilder s -> ST s TextChunk freezeTextChunk b = do     n <- VUM.unsafeRead (tbCounters b) cRows

@@ -14,11 +14,7 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} -{- |-Conversions between 'Column' and ordinary vectors\/lists, plus the typed-extraction functions ('toVector', 'toDoubleVector', ...) that recover a-column's element type.--}+-- | Conversions between 'Column' and vectors or lists, plus typed extraction. module DataFrame.Internal.Column.Conversion where  import qualified Data.Text as T@@ -29,9 +25,10 @@  import Control.Exception (throw) import Control.Monad.ST (ST)+import Data.Int (Int16, Int32, Int64, Int8) import Data.Kind (Type) import Data.Type.Equality (TestEquality (..))-import Data.Word (Word8)+import Data.Word (Word16, Word32, Word64, Word8) import DataFrame.Errors (     DataFrameException (ExpectedNonNullableException, TypeMismatchException),     TypeErrorContext (@@ -50,63 +47,26 @@ import System.Random (RandomGen, UniformRange, uniformR) import Type.Reflection (TypeRep, Typeable, typeRep, type (:~:) (Refl)) -{- | O(n) Convert a vector to a column. Automatically picks the best representation of a vector to store the underlying data in.--__Examples:__--@-> import qualified Data.Vector as V-> fromVector (VB.fromList [(1 :: Int), 2, 3, 4])-[1,2,3,4]-@--}+-- | O(n) Convert a vector to a column, picking the best representation. fromVector ::     forall a.     (Columnable a, ColumnifyRep (KindOf a) a) =>     VB.Vector a -> Column fromVector = toColumnRep @(KindOf a) -{- | O(n) Convert an unboxed vector to a column. This avoids the extra conversion if you already have the data in an unboxed vector.--__Examples:__--@-> import qualified Data.Vector.Unboxed as V-> fromUnboxedVector (VB.fromList [(1 :: Int), 2, 3, 4])-[1,2,3,4]-@--}+-- | O(n) Convert an unboxed vector to a column. fromUnboxedVector ::     forall a. (Columnable a, VU.Unbox a) => VU.Vector a -> Column fromUnboxedVector = UnboxedColumn Nothing -{- | O(n) Convert a list to a column. Automatically picks the best representation of a vector to store the underlying data in.--__Examples:__--@-> fromList [(1 :: Int), 2, 3, 4]-[1,2,3,4]-@--}+-- | O(n) Convert a list to a column, picking the best representation. fromList ::     forall a.     (Columnable a, ColumnifyRep (KindOf a) a) =>     [a] -> Column fromList = toColumnRep @(KindOf a) . VB.fromList -{- | O(n) Create a column of random elements within a range.--Takes a random number generator, a length, and a lower and upper bound for the random values.--__Examples:__--@-> import System.Random (mkStdGen)-> mkRandom (mkStdGen 42) 4 0 10-[4,2,6,5]-@--}+-- | O(n) Create a column of random elements within a range. mkRandom ::     (RandomGen g, Columnable a, ColumnifyRep (KindOf a) a, UniformRange a) =>     g -> Int -> a -> a -> Column@@ -119,32 +79,13 @@          in             v : go g' (n - 1) -{- | O(n) Converts a column to a list. Throws an exception if the wrong type is specified.--__Examples:__--@-> column = fromList [(1 :: Int), 2, 3, 4]-> toList @Int column-[1,2,3,4]-> toList @Double column-exception: ...-@--}+-- | O(n) Convert a column to a list; throws if the type is wrong. toList :: forall a. (Columnable a) => Column -> [a] toList xs = case toVector @a xs of     Left err -> throw err     Right val -> VB.toList val -{- | Type-safe conversion of a column to a vector of element type @a@ (specify via-type application); 'Left' 'TypeMismatchException' when the column's type differs.-->>> toVector @Int @VU.Vector column-Right (unboxed vector of Ints)-->>> toVector @Text @VB.Vector column-Right (boxed vector of Text)--}+-- | Type-safe conversion of a column to a vector of element type @a@. toVector ::     forall a v.     (VG.Vector v a, Columnable a) => Column -> Either DataFrameException (v a)@@ -154,7 +95,9 @@     BoxedColumn bm (inner :: VB.Vector c) ->         -- Check if user wants Maybe c (nullable) or c directly         case testEquality (typeRep @a) (typeRep @c) of-            Just Refl -> Right $ VG.convert inner+            Just Refl+                | hasNulls (VB.length inner) bm -> Left (nullsMismatch @c)+                | otherwise -> Right $ VG.convert inner             Nothing ->                 -- Try: a = Maybe c                 case testEquality (typeRep @a) (typeRep @(Maybe c)) of@@ -178,7 +121,9 @@                                 )     UnboxedColumn bm (inner :: VU.Vector c) ->         case testEquality (typeRep @a) (typeRep @c) of-            Just Refl -> Right $ VG.convert inner+            Just Refl+                | hasNulls (VU.length inner) bm -> Left (nullsMismatch @c)+                | otherwise -> Right $ VG.convert inner             Nothing ->                 case testEquality (typeRep @a) (typeRep @(Maybe c)) of                     Just Refl ->@@ -199,12 +144,24 @@                                 ) {-# INLINEABLE toVector #-} +hasNulls :: Int -> Maybe Bitmap -> Bool+hasNulls n = maybe False (bitmapHasNulls n)++-- | The column has nulls, so it must be read as @Maybe c@.+nullsMismatch :: forall c. (Columnable c) => DataFrameException+nullsMismatch =+    TypeMismatchException+        ( MkTypeErrorContext+            { userType = Right (typeRep @c)+            , expectedType = Right (typeRep @(Maybe c))+            , callingFunctionName = Just "toVector"+            , errorColumnName = Nothing+            }+        )+ -- Some common types we will use for numerical computing. -{- | Convert a column to an unboxed 'Double' vector, coercing numeric types-('realToFrac' for floats, 'fromIntegral' for integrals; nulls become @NaN@).-'Left' 'TypeMismatchException' when the column is not numeric.--}+-- | Convert a column to an unboxed 'Double' vector, coercing numeric types. toDoubleVector :: Column -> Either DataFrameException (VU.Vector Double) toDoubleVector column =     case column of@@ -252,10 +209,7 @@                             }                         ) -{- | Convert a column to an unboxed 'Float' vector, coercing numeric types (nulls-become @NaN@); 'Left' 'TypeMismatchException' when not numeric. Converting from-'Double' may lose precision.--}+-- | Convert a column to an unboxed 'Float' vector, coercing numeric types. toFloatVector :: Column -> Either DataFrameException (VU.Vector Float) toFloatVector column =     case column of@@ -317,10 +271,7 @@                             }                         ) -{- | Convert a column to an unboxed 'Int' vector, coercing numeric types-(floats are 'round'ed via banker's rounding); 'Left' 'TypeMismatchException'-when the column is not numeric. Does not support nullable columns.--}+-- | Convert a column to an unboxed 'Int' vector, coercing numeric types. toIntVector :: Column -> Either DataFrameException (VU.Vector Int) toIntVector column =     case column of@@ -330,9 +281,9 @@             Just Refl -> Right f             Nothing -> case sFloating @a of                 STrue -> Right (VU.map (round . (realToFrac :: a -> Double)) f)-                SFalse -> case sIntegral @a of-                    STrue -> Right (VU.map fromIntegral f)-                    SFalse ->+                SFalse -> case widenToInt f of+                    Just v -> Right v+                    Nothing ->                         Left $                             TypeMismatchException                                 ( MkTypeErrorContext@@ -355,6 +306,33 @@                             }                         ) +-- | Widen an integral vector to 'Int', or 'Nothing' if not integral.+widenToInt ::+    forall a. (Columnable a, VU.Unbox a) => VU.Vector a -> Maybe (VU.Vector Int)+widenToInt v+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = Just v+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int64) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int32) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int16) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int8) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Word) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Word64) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Word32) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Word16) =+        Just (VU.map fromIntegral v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Word8) =+        Just (VU.map fromIntegral v)+    | otherwise = case sIntegral @a of+        STrue -> Just (VU.map fromIntegral v)+        SFalse -> Nothing+ toUnboxedVector ::     forall a.     (Columnable a, VU.Unbox a) => Column -> Either DataFrameException (VU.Vector a)@@ -384,9 +362,7 @@                     ) {-# INLINE toUnboxedVector #-} --- Shared finaliser for the two parseUnboxedColumn* helpers.  Freezes--- the mutable data vector, and only materialises the bitmap when the--- column actually had nulls.+-- Freezes the data vector; builds the bitmap only if the column had nulls. {-# INLINE finalizeParseResult #-} finalizeParseResult ::     (VU.Unbox a) =>

@@ -4,11 +4,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | Dictionary-encode a text (or factor) group key to dense @Int@ codes: each row-gets a first-appearance code @0..card-1@ (NULL reserved) plus the cardinality.--TODO: mchavinda - revise if this module is still necessary.--}+-- | Dictionary-encode a text group key to dense first-appearance Int codes. module DataFrame.Internal.Column.Encode (     dictEncodeColumn,     dictEncodeColumnUpTo,@@ -44,24 +40,15 @@     sliceEqBytes,  ) -{- | Largest distinct-value count we will dictionary-encode. Above this the codes-no longer index a reasonable direct accumulator and the encode pass is pure-overhead, so the caller keeps the plain hash group-by.--}+-- | Largest distinct-value count we will dictionary-encode. dictMaxCardinality :: Int dictMaxCardinality = 1048576 -{- | Dictionary-encode a text-like column to dense first-appearance @Int@ codes,-returning @Just (codes, cardinality)@ (a NULL row gets its own reserved code).-'Nothing' for non-text columns or cardinality above 'dictMaxCardinality'.--}+-- | Dictionary-encode a text-like column to dense codes plus cardinality. dictEncodeColumn :: Column -> Maybe (VU.Vector Int, Int) dictEncodeColumn = dictEncodeColumnUpTo dictMaxCardinality -{- | Dictionary-encode like 'dictEncodeColumn' but bail to 'Nothing' as soon as-the distinct count would exceed @maxCard@, letting a low-cardinality probe avoid-a full high-cardinality pass.--}+-- | Like 'dictEncodeColumn' but bails to 'Nothing' past @maxCard@ values. dictEncodeColumnUpTo :: Int -> Column -> Maybe (VU.Vector Int, Int) dictEncodeColumnUpTo maxCard (PackedText bm p) = encodePacked maxCard bm p dictEncodeColumnUpTo maxCard (BoxedColumn bm (v :: V.Vector a)) =@@ -70,10 +57,7 @@         Nothing -> Nothing dictEncodeColumnUpTo _ _ = Nothing -{- | Encode a packed-text column: hash each row's raw UTF-8 bytes (the grouping-'mixBytes'), re-verify byte equality on collisions, assign dense codes in-first-appearance order. A null row hashes 'nullSalt'.--}+-- | Encode a packed-text column by hashing raw UTF-8 bytes. encodePacked ::     Int -> Maybe Bitmap -> PackedTextData -> Maybe (VU.Vector Int, Int) encodePacked maxCard bm p =@@ -95,9 +79,7 @@                 _ -> 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).--}+-- | Encode a boxed 'Data.Text.Text' column, mirroring 'encodePacked'. encodeBoxedText ::     Int -> Maybe Bitmap -> V.Vector T.Text -> Maybe (VU.Vector Int, Int) encodeBoxedText maxCard bm v =@@ -114,10 +96,7 @@                 _ -> False      in buildCodes maxCard n hashAt eqAt -{- | The shared code-assignment loop: bucket every row through an open-addressing-table on its precomputed hash, re-verify with @eqAt@ on a hit, assign dense-first-appearance codes. Bails to 'Nothing' once the distinct count exceeds @maxCard@.--}+-- | The shared code-assignment loop over an open-addressing hash table. buildCodes ::     Int -> Int -> (Int -> Int) -> (Int -> Int -> Bool) -> Maybe (VU.Vector Int, Int) buildCodes maxCard n hashAt eqAt

@@ -4,14 +4,13 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | Concatenation of per-chunk 'Column's (e.g. from parallel CSV chunks). Text-columns merge at the byte level via 'TextChunk' \/ 'mergeTextChunks', so no-per-chunk 'Data.Text.Text' values are ever materialized.--}+-- | Concatenation of per-chunk 'Column's, merging text at the byte level. module DataFrame.Internal.Column.Merge (     TextChunk (..),     concatColumns,+    concatColumnsIO,     mergeTextChunks,+    mergeTextChunksIO,     packedFromTextChunk,     concatValidity,     tcRows,@@ -22,8 +21,8 @@ import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM -import Control.Monad (foldM_, forM_)-import Control.Monad.ST (runST)+import Control.Monad (foldM_, forM_, when)+import Control.Monad.ST (runST, stToIO) import Data.Type.Equality (testEquality, (:~:) (Refl)) import DataFrame.Internal.Column (     Column (..),@@ -38,13 +37,21 @@     Validity (Validity),     concatValidity,  )-import DataFrame.Internal.Data.PackedText (mkPackedContiguous)+import DataFrame.Internal.Control.Concurrent (+    forkJoin_,+    parThreshold,+    shouldParallelize,+ )+import DataFrame.Internal.Data.PackedText (+    PackedOffsets (..),+    mkPackedContiguous,+    offAt,+    offCount,+    packedRowOffsets,+ ) import Type.Reflection (typeRep) -{- | A frozen text-builder chunk: 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'.--}+-- | A frozen text-builder chunk: UTF-8 bytes, row offsets and validity. data TextChunk = TextChunk     { tcBytes :: !A.Array     , tcUsed :: !Int@@ -55,18 +62,12 @@ tcRows :: TextChunk -> Int tcRows c = VU.length (tcOffsets c) - 1 -{- | Freeze a builder chunk directly into a packed-text column: no-'Data.Text.Text' materialization, no UTF-8 validation pass (deferred to decode).-Not yet called by any reader.--}+-- | Freeze a builder chunk directly into a packed-text column. packedFromTextChunk :: TextChunk -> Column packedFromTextChunk (TextChunk arr _used offs bm) =     PackedText bm (mkPackedContiguous arr offs) -{- | Merge text chunks into one packed-text 'Column': one byte-array copy per-chunk, one offset rebase, then wrap the shared buffer + offsets as 'PackedText'-(no per-row header, decode deferred).--}+-- | Merge text chunks into one packed-text 'Column' with one copy per chunk. mergeTextChunks :: [TextChunk] -> Column mergeTextChunks [] = error "DataFrame.Internal.ColumnMerge.mergeTextChunks: empty list" mergeTextChunks [c] = packedFromTextChunk c@@ -89,16 +90,48 @@     let !bm = concatValidity [Validity (tcBitmap c) (tcRows c) | c <- cs]     pure (PackedText bm (mkPackedContiguous farr foffs)) -{- | Merge per-chunk columns into one column.+-- | 'mergeTextChunks' with copies and offset rebases spread over workers.+mergeTextChunksIO :: [TextChunk] -> IO Column+mergeTextChunksIO [] = error "DataFrame.Internal.ColumnMerge.mergeTextChunksIO: empty list"+mergeTextChunksIO [c] = pure (packedFromTextChunk c)+mergeTextChunksIO cs+    | not (shouldParallelize parThreshold totalRows) = pure (mergeTextChunks cs)+    | otherwise = do+        arr <- stToIO (A.new (max 1 totalBytes))+        offs <- VUM.unsafeNew (totalRows + 1)+        VUM.unsafeWrite offs 0 0+        let byteBases = scanl (+) 0 (map tcUsed cs)+            rowBases = scanl (+) 0 (map tcRows cs)+            splice byteBase rowBase c = do+                let n = tcRows c+                    co = tcOffsets c+                    go !i = when (i <= n) $ do+                        VUM.unsafeWrite offs (rowBase + i) (byteBase + VU.unsafeIndex co i)+                        go (i + 1)+                stToIO (A.copyI (tcUsed c) arr byteBase (tcBytes c) 0)+                go 1+        forkJoin_ (zipWith3 splice byteBases rowBases cs)+        farr <- stToIO (A.unsafeFreeze arr)+        foffs <- VU.unsafeFreeze offs+        let !bm = concatValidity [Validity (tcBitmap c) (tcRows c) | c <- cs]+        pure (PackedText bm (mkPackedContiguous farr foffs))+  where+    totalBytes = sum (map tcUsed cs)+    totalRows = sum (map tcRows cs) -TODO: mchavinda - this is very similar to mappendColumns can could possibly-be defined in terms of it but I'll have to ivnestigate further.--}+-- | 'concatColumns', merging packed text in parallel.+concatColumnsIO :: [Column] -> IO Column+concatColumnsIO cols+    | Just chunks <- mapM packedChunk cols = mergeTextChunksIO chunks+    | otherwise = pure (concatColumns cols)++-- | Merge per-chunk columns into one column. concatColumns :: [Column] -> Column concatColumns [] = error "DataFrame.Internal.Column.Builder.concatColumns: empty list" concatColumns [c] = c--- Normalize on the whole list, not the head: a packed or merged chunk in any--- position must demote every chunk to the common boxed form.+-- Normalize the whole list: a packed or merged chunk demotes all to boxed.+concatColumns cols+    | Just chunks <- mapM packedChunk cols = mergeTextChunks chunks concatColumns cols@(c0 : _)     | any isMergedColumn cols = concatColumns (map materializeMerged cols)     | any isPackedText cols = concatColumns (map materializePacked cols)@@ -136,3 +169,15 @@ mergeMismatch =     error         "DataFrame.Internal.Column.Builder.concatColumns: chunk column types differ"++-- | A packed-text column without a selection, viewed as a builder chunk.+packedChunk :: Column -> Maybe TextChunk+packedChunk (PackedText bm p) = case packedRowOffsets p of+    Just (arr, offs) ->+        let n = offCount offs - 1+            offsV = case offs of+                Offs64 v -> v+                Offs32 v -> VU.map fromIntegral v+         in Just (TextChunk arr (offAt offs n) offsV bm)+    Nothing -> Nothing+packedChunk _ = Nothing

@@ -8,27 +8,28 @@ {-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE PolyKinds #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} -{- |-Bulk operations over columns: mapping, folding, slicing, gathering, zipping,-appending, and the mutable-column IO helpers.--}+-- | Bulk operations over columns: map, fold, slice, gather, zip, append. module DataFrame.Internal.Column.Operations where  import qualified Data.Text as T+import qualified Data.Text.Array as A+import Data.Text.Internal (Text (Text)) import qualified Data.Vector as VB import qualified Data.Vector.Generic as VG import qualified Data.Vector.Mutable as VBM import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM -import Control.Monad (when)+import Control.Monad (forM_, join, when) import Control.Monad.ST (runST) import Data.Bits (setBit, shiftL, shiftR) import Data.Int (Int32)@@ -52,31 +53,38 @@ import DataFrame.Internal.Column.Properties import DataFrame.Internal.Column.Types import DataFrame.Internal.Control.Concurrent (+    capabilities,+    forkJoin,+    forkJoin_,     parThreshold,+    parallelChunks,     parallelChunks_,     shouldParallelize,+    splitChunkRange,  ) import DataFrame.Internal.Data.PackedText (-    PackedOffsets,     PackedSel (..),     PackedTextData (..),+    mkPackedContiguous,     offCount,     packedGather,+    packedIndexText,     packedLength,  ) import System.IO.Unsafe (unsafePerformIO) import Type.Reflection (     TypeRep,     Typeable,+    eqTypeRep,     typeOf,     typeRep,     withTypeable,+    pattern App,     type (:~:) (Refl),+    type (:~~:) (HRefl),  ) -{- | Force evaluation of all elements in a column. Replacement for the removed-@instance NFData Column@; used by the IO and lazy-executor strict paths.--}+-- | Force evaluation of all elements in a column. forceColumn :: Column -> () forceColumn (BoxedColumn Nothing (v :: VB.Vector a)) = VB.foldl' (const (`seq` ())) () v forceColumn (BoxedColumn (Just bm) (v :: VB.Vector a)) =@@ -91,11 +99,7 @@ forceColumn (MergedColumn a b) =     forceColumn a `seq` forceColumn b `seq` checkMergedNoBothNull a b -{- | 'MergedColumn' defers element construction, so forcing must still surface-the one deferred error — a row null on both sides — inside strict IO/executor-boundaries. O(rows) bitmap walk, no allocation; both-null needs a bitmap on-each side, so anything else passes immediately.--}+-- | Surface the deferred both-sides-null error of a 'MergedColumn'. checkMergedNoBothNull :: Column -> Column -> () checkMergedNoBothNull a b = case (columnBitmap a, columnBitmap b) of     (Just ba, Just bb) ->@@ -107,7 +111,7 @@          in go 0     _ -> () --- | Allocate a mutable column of size @n@ matching the constructor/type of the given column.+-- | Allocate a mutable column of size @n@ matching the given column. newMutableColumn :: Int -> Column -> IO MutableColumn newMutableColumn n (BoxedColumn _ (_ :: VB.Vector a)) =     MBoxedColumn <$> (VBM.new n :: IO (VBM.IOVector a))@@ -146,9 +150,21 @@ mapColumn f = \case     BoxedColumn bm (col :: VB.Vector a) -> runBoxed bm col     UnboxedColumn bm (col :: VU.Vector a) -> runUnboxed bm col+    PackedText Nothing p+        | Just Refl <- testEquality (typeRep @b) (typeRep @T.Text) ->+            Right (mapPacked p)     c@(PackedText _ _) -> mapColumn f (materializePacked c)     c@(MergedColumn _ _) -> mapColumn f (materializeMerged c)   where+    -- Rows decode one at a time from packed bytes; text results are re-packed.+    mapPacked :: (b ~ T.Text) => PackedTextData -> Column+    mapPacked p =+        let !n = packedLength p+         in case testEquality (typeRep @c) (typeRep @T.Text) of+                Just Refl -> packedMapText n (f . packedIndexText p)+                Nothing -> case sUnbox @c of+                    STrue -> UnboxedColumn Nothing (parGenerateUnboxed n (f . packedIndexText p))+                    SFalse -> fromVector @c (VB.generate n (f . packedIndexText p))     runBoxed ::         forall a.         (Columnable a) =>@@ -173,14 +189,14 @@                                 )         Nothing -> case testEquality (typeRep @a) (typeRep @b) of             Just Refl ->-                Right $ case sUnbox @c of-                    STrue ->-                        UnboxedColumn-                            bm-                            (parGenerateUnboxed (VB.length col) (f . VB.unsafeIndex col))-                    SFalse -> case bm of-                        Nothing -> fromVector @c (VB.map f col)-                        Just _ -> BoxedColumn bm (VB.map f col)+                let !n = VB.length col+                 in Right $ case sUnbox @c of+                        STrue -> UnboxedColumn bm $ case nullsIn n bm of+                            Nothing -> parGenerateUnboxed n (f . VB.unsafeIndex col)+                            Just nulls -> generateValidUnboxed nulls n (f . VB.unsafeIndex col)+                        SFalse -> case bm of+                            Nothing -> fromVector @c (VB.map f col)+                            Just b' -> generateValidBoxed b' n (f . VB.unsafeIndex col)             Nothing -> throwTypeMismatch @a @b      runUnboxed ::@@ -206,18 +222,19 @@                                     else Nothing                                 )         Nothing -> case testEquality (typeRep @a) (typeRep @b) of-            Just Refl -> Right $ case sUnbox @c of-                STrue ->-                    UnboxedColumn-                        bm-                        (parGenerateUnboxed (VU.length col) (f . VU.unsafeIndex col))-                SFalse -> case bm of-                    Nothing -> fromVector @c (VB.generate (VU.length col) (f . VU.unsafeIndex col))-                    Just _ -> BoxedColumn bm (VB.generate (VU.length col) (f . VU.unsafeIndex col))+            Just Refl ->+                let !n = VU.length col+                 in Right $ case sUnbox @c of+                        STrue -> UnboxedColumn bm $ case nullsIn n bm of+                            Nothing -> fromMaybe mapLoop compiledMap f col+                            Just nulls -> generateValidUnboxed nulls n (f . VU.unsafeIndex col)+                        SFalse -> case bm of+                            Nothing -> fromVector @c (VB.generate n (f . VU.unsafeIndex col))+                            Just b' -> generateValidBoxed b' n (f . VU.unsafeIndex col)             Nothing -> throwTypeMismatch @a @b {-# INLINEABLE mapColumn #-} --- | Applies a function that returns an unboxed result to an unboxed vector, storing the result in a column.+-- | Apply a function returning an unboxed result to an unboxed vector. imapColumn ::     forall b c.     (Columnable b, Columnable c) =>@@ -233,12 +250,15 @@         (Columnable a) =>         Maybe Bitmap -> VB.Vector a -> Either DataFrameException Column     runBoxed bm col = case testEquality (typeRep @a) (typeRep @b) of-        Just Refl -> Right $ case sUnbox @c of-            STrue ->-                UnboxedColumn-                    bm-                    (VU.generate (VB.length col) (\i -> f i (VB.unsafeIndex col i)))-            SFalse -> BoxedColumn bm (VB.imap f col)+        Just Refl ->+            let !n = VB.length col+             in Right $ case sUnbox @c of+                    STrue -> UnboxedColumn bm $ case nullsIn n bm of+                        Nothing -> VU.generate n (\i -> f i (VB.unsafeIndex col i))+                        Just nulls -> generateValidUnboxed nulls n (\i -> f i (VB.unsafeIndex col i))+                    SFalse -> case bm of+                        Nothing -> fromVector @c (VB.imap f col)+                        Just b' -> generateValidBoxed b' n (\i -> f i (VB.unsafeIndex col i))         Nothing -> throwTypeMismatch @a @b      runUnboxed ::@@ -246,9 +266,15 @@         (Columnable a, VU.Unbox a) =>         Maybe Bitmap -> VU.Vector a -> Either DataFrameException Column     runUnboxed bm col = case testEquality (typeRep @a) (typeRep @b) of-        Just Refl -> Right $ case sUnbox @c of-            STrue -> UnboxedColumn bm (VU.imap f col)-            SFalse -> BoxedColumn bm (VB.imap f (VG.convert col))+        Just Refl ->+            let !n = VU.length col+             in Right $ case sUnbox @c of+                    STrue -> UnboxedColumn bm $ case nullsIn n bm of+                        Nothing -> VU.imap f col+                        Just nulls -> generateValidUnboxed nulls n (\i -> f i (VU.unsafeIndex col i))+                    SFalse -> case bm of+                        Nothing -> fromVector @c (VB.imap f (VG.convert col))+                        Just b' -> generateValidBoxed b' n (\i -> f i (VU.unsafeIndex col i))         Nothing -> throwTypeMismatch @a @b  -- | O(n) Takes the last n values of a column.@@ -265,8 +291,7 @@ sliceColumn start n (UnboxedColumn bm xs) =     UnboxedColumn (fmap (bitmapSlice start n) bm) (VG.slice start n xs) sliceColumn start n (PackedText bm p)-    -- packedGather decodes an out-of-range index as the empty string, where-    -- the boxed and unboxed arms reject the slice, so check first.+    -- packedGather reads out-of-range as "", so check the slice bounds first.     | start < 0 || n < 0 || start + n > packedLength p =         errorWithoutStackTrace             ( "sliceColumn: invalid slice ("@@ -283,26 +308,9 @@             (packedGather (VU.enumFromN start n) p) {-# INLINE sliceColumn #-} --- | O(n) Selects the elements at a given set of indices. Does not change the order.------------------------------------------------------------------------------------ Parallel element-wise kernels--------------------------------------------------------------------------------+-- (n) Selects the elements at a given set of indices, keeping order. -{- | Parallel unboxed 'VU.generate': splits the index space into one contiguous-chunk per capability, evaluates each chunk into its disjoint slice of a single-pre-allocated mutable vector, then freezes. Element @i@ depends only on @f i@,-so the result is bit-identical to the sequential 'VU.generate' regardless of-capability count. Falls back to 'VU.generate' below 'parThreshold'.--}-{-# SPECIALIZE parGenerateUnboxed ::-    Int -> (Int -> Double) -> VU.Vector Double-    #-}-{-# SPECIALIZE parGenerateUnboxed ::-    Int -> (Int -> Float) -> VU.Vector Float-    #-}-{-# SPECIALIZE parGenerateUnboxed :: Int -> (Int -> Int) -> VU.Vector Int #-}-{-# SPECIALIZE parGenerateUnboxed :: Int -> (Int -> Bool) -> VU.Vector Bool #-}+-- | Parallel unboxed 'VU.generate', one contiguous chunk per capability. parGenerateUnboxed :: (VU.Unbox c) => Int -> (Int -> c) -> VU.Vector c parGenerateUnboxed n f     | not (shouldParallelize parThreshold n) = VU.generate n f@@ -312,10 +320,7 @@         VU.unsafeFreeze mv {-# NOINLINE parGenerateUnboxed #-} -{- | The chunk body shared by 'parGenerateUnboxed' and-'parGenerateUnboxedInline'. INLINE so the latter's monomorphic wrappers each-get their own copy with @f@ inlined.--}+-- | The chunk body shared by 'parGenerateUnboxed' and its inline variant. fillGenerate ::     (VU.Unbox c) => VUM.IOVector c -> (Int -> c) -> Int -> Int -> IO () fillGenerate mv f !lo !hi =@@ -325,26 +330,14 @@      in go lo {-# INLINE fillGenerate #-} -{- | Parallel unboxed gather: element @i@ of the result is-@v ! (ix ! i)@ (unsafe indexing — callers pass in-bounds index vectors, e.g.-grouping-produced representative rows). Same chunking as 'parGenerateUnboxed'-(one contiguous chunk per capability into disjoint slices of one buffer), so-the result is bit-identical to the sequential backpermute at any capability-count; falls back to a sequential loop below 'parThreshold'.--}+-- | Parallel unboxed gather: element @i@ is @v ! (ix ! i)@. parBackpermuteUnboxed ::     (VU.Unbox a) => VU.Vector a -> VU.Vector Int -> VU.Vector a parBackpermuteUnboxed v ix =     parGenerateUnboxed (VU.length ix) (VU.unsafeIndex v . VU.unsafeIndex ix) {-# INLINE parBackpermuteUnboxed #-} -{- | 'parGenerateUnboxed' with an INLINE body: each monomorphic NOINLINE-wrapper below gets its own copy of the fill loop with @f@ inlined, so the-per-element unknown closure call (and the boxed result it returns) disappears-— on a 1e8-row gather that call+alloc dominated the whole pass. Same chunking,-bit-identical results; wrappers must stay NOINLINE so the 'unsafePerformIO'-runs once per call.--}+-- | 'parGenerateUnboxed' with an INLINE body for monomorphic wrappers. parGenerateUnboxedInline :: (VU.Unbox c) => Int -> (Int -> c) -> VU.Vector c parGenerateUnboxedInline n f     | not (shouldParallelize parThreshold n) = VU.generate n f@@ -354,6 +347,95 @@         VU.unsafeFreeze mv {-# INLINE parGenerateUnboxedInline #-} +-- | Run @k@ with the element type known, or 'Nothing' for other types.+withKnownElem ::+    forall a r. (Typeable a) => (forall x. (VU.Unbox x) => a :~: x -> r) -> Maybe r+withKnownElem k+    | Just Refl <- testEquality (typeRep @a) (typeRep @Double) =+        Just (k @Double Refl)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = Just (k @Int Refl)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Bool) = Just (k @Bool Refl)+    | otherwise = Nothing+{-# INLINE withKnownElem #-}++-- | 'Nothing' for element types without a compiled loop.+compiledMap ::+    forall a c.+    (Typeable a, Typeable c) => Maybe ((a -> c) -> VU.Vector a -> VU.Vector c)+compiledMap = join (withKnownElem @a (\Refl -> withKnownElem @c (\Refl -> mapLoop)))++compiledZip ::+    forall a c.+    (Typeable a, Typeable c) =>+    Maybe ((a -> a -> c) -> VU.Vector a -> VU.Vector a -> VU.Vector c)+compiledZip = join (withKnownElem @a (\Refl -> withKnownElem @c (\Refl -> zipLoop)))++compiledFoldl ::+    forall a b.+    (Typeable a, Typeable b) => Maybe ((b -> a -> b) -> b -> VU.Vector a -> b)+compiledFoldl = join (withKnownElem @a (\Refl -> withKnownElem @b (\Refl -> VU.foldl')))++compiledFoldl1 ::+    forall a. (Typeable a) => Maybe ((a -> a -> a) -> VU.Vector a -> a)+compiledFoldl1 = withKnownElem @a (\Refl -> VU.foldl1')++mapLoop :: (VU.Unbox a, VU.Unbox c) => (a -> c) -> VU.Vector a -> VU.Vector c+mapLoop f v = parGenerateUnboxedInline (VU.length v) (f . VU.unsafeIndex v)+{-# INLINE mapLoop #-}++-- | Element-wise over two vectors, to the shorter length.+zipLoop ::+    (VU.Unbox a, VU.Unbox b, VU.Unbox c) =>+    (a -> b -> c) -> VU.Vector a -> VU.Vector b -> VU.Vector c+zipLoop f l r =+    parGenerateUnboxedInline+        (min (VU.length l) (VU.length r))+        (\i -> f (VU.unsafeIndex l i) (VU.unsafeIndex r i))+{-# INLINE zipLoop #-}++-- | The bitmap, if it marks any of the first @n@ rows null.+nullsIn :: Int -> Maybe Bitmap -> Maybe Bitmap+nullsIn n (Just bm) | bitmapHasNulls n bm = Just bm+nullsIn _ _ = Nothing+{-# INLINE nullsIn #-}++validAt :: Maybe Bitmap -> Int -> Bool+validAt bm i = maybe True (`bitmapTestBit` i) bm+{-# INLINE validAt #-}++-- | @g i@ for the valid rows only; null rows hold zero.+generateValidUnboxed ::+    (VU.Unbox c) => Bitmap -> Int -> (Int -> c) -> VU.Vector c+generateValidUnboxed bm n g = unsafePerformIO $ do+    mv <- VUM.new n+    let fill !lo !hi =+            let go !i+                    | i >= hi = pure ()+                    | bitmapTestBit bm i = VUM.unsafeWrite mv i (g i) >> go (i + 1)+                    | otherwise = go (i + 1)+             in go lo+    if shouldParallelize parThreshold n+        then parallelChunks_ parThreshold n fill+        else fill 0 n+    VU.unsafeFreeze mv+{-# NOINLINE generateValidUnboxed #-}++-- | A boxed column of @g i@ for the valid rows only; null rows stay null.+generateValidBoxed ::+    forall c. (Columnable c) => Bitmap -> Int -> (Int -> c) -> Column+generateValidBoxed bm n g = case typeRep @c of+    App tMaybe _+        | Just HRefl <- eqTypeRep tMaybe (typeRep @Maybe) ->+            fromVector @c+                (VB.generate n (\i -> if bitmapTestBit bm i then g i else Nothing))+    _ ->+        BoxedColumn+            (Just bm)+            ( VB.generate+                n+                (\i -> if bitmapTestBit bm i then g i else errorWithoutStackTrace "null slot")+            )+ -- | Closure-free parallel 'Int' gather: @out!i = v ! (ix!i)@. parBackpermuteInt :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int parBackpermuteInt v ix =@@ -370,10 +452,7 @@         (VU.unsafeIndex v . VU.unsafeIndex ix) {-# NOINLINE parBackpermuteDouble #-} -{- | Closure-free double-indirection gather: @out!g = vis ! (offs!g)@ over-@length offs - 1@ groups (the representative-row build of the 'Grouped'-pattern).--}+-- | Closure-free double-indirection gather: @out!g = vis ! (offs!g)@. parBackpermute2Int :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int parBackpermute2Int vis offs =     parGenerateUnboxedInline@@ -381,11 +460,7 @@         (VU.unsafeIndex vis . VU.unsafeIndex offs) {-# NOINLINE parBackpermute2Int #-} -{- | Parallel boxed gather. The read side uses 'VB.unsafeIndexM' so the array-slot is fetched eagerly (element pointers are shared, elements themselves stay-un-forced, exactly as the sequential 'VB.generate' gather). Bit-identical-element values; sequential below 'parThreshold'.--}+-- | Parallel boxed gather that fetches array slots eagerly. parBackpermuteBoxed :: VB.Vector a -> VU.Vector Int -> VB.Vector a parBackpermuteBoxed v ix     | not (shouldParallelize parThreshold n) =@@ -405,9 +480,7 @@     !n = VU.length ix {-# NOINLINE parBackpermuteBoxed #-} -{- | Clamp sentinel (negative) indices to 0 in one closure-free parallel pass.-The clamped rows read row 0's value; callers mask them via the sentinel bitmap.--}+-- | Clamp sentinel (negative) indices to 0 in one parallel pass. parClampNonNeg :: VU.Vector Int -> VU.Vector Int parClampNonNeg ix =     parGenerateUnboxedInline@@ -415,9 +488,7 @@         (\i -> let !x = VU.unsafeIndex ix i in max x 0) {-# NOINLINE parClampNonNeg #-} -{- | Validity bitmap from sentinel indices (bit @i@ valid iff @ix!i >= 0@),-built one byte (8 rows) per element in parallel.--}+-- | Validity bitmap from sentinel indices, built in parallel. parBitmapNonNeg :: VU.Vector Int -> Bitmap parBitmapNonNeg ix =     let !n = VU.length ix@@ -456,8 +527,7 @@             )             bm         )-        -- Int/Double hit the closure-free monomorphic kernels; anything else-        -- keeps the generic (per-element closure call) path.+        -- Int/Double use closure-free kernels; others take the generic path.         ( case testEquality (typeRep @a) (typeRep @Int) of             Just Refl -> parBackpermuteInt column indexes             Nothing -> case testEquality (typeRep @a) (typeRep @Double) of@@ -478,9 +548,7 @@         (packedGather indexes p) {-# INLINE atIndicesStable #-} -{- | Like 'atIndicesStable' but treats negative indices as null.-Keeps the index vector fully unboxed (no @VB.Vector (Maybe Int)@).--}+-- | Like 'atIndicesStable' but treats negative indices as null. gatherWithSentinel :: VU.Vector Int -> Column -> Column gatherWithSentinel indices c@(MergedColumn _ _) =     gatherWithSentinel indices (materializeMerged c)@@ -490,20 +558,16 @@         withBm srcBm = case srcBm of             Nothing -> Just newBm             Just sb -> Just (andBitmaps newBm (gatherSrcBm sb))-        -- Sequential fallback: gather an existing source bitmap through the-        -- raw indices (negative-guarded). Only runs when the source side is-        -- itself nullable, which join build sides normally are not.+        -- Sequential fallback: gather a nullable source bitmap by raw indices.         gatherSrcBm sb =             buildBitmapFromValid $ VU.generate n $ \i ->                 let idx = VU.unsafeIndex indices i                  in if idx >= 0 && bitmapTestBit sb idx then 1 else 0      in case col of-            -- packedGather composes -1 sentinels into the selector natively-            -- (and shares the byte buffers), so text takes the raw indices.+            -- packedGather handles -1 sentinels itself; pass raw indices.             PackedText srcBm p -> PackedText (withBm srcBm) (packedGather indices p)             BoxedColumn srcBm v-                -- An empty source means every index is a sentinel; clamping-                -- would read row 0 of an empty vector.+                -- An empty source means every index is a sentinel.                 | VB.null v -> BoxedColumn (withBm srcBm) (allNullBoxed n v)                 | otherwise ->                     BoxedColumn@@ -512,17 +576,13 @@             UnboxedColumn srcBm v                 | VU.null v -> UnboxedColumn (withBm srcBm) (allNullUnboxed n v)                 | otherwise ->-                    -- Reuse atIndicesStable's Int/Double monomorphic kernel-                    -- dispatch for the payload; the bitmap is replaced below.+                    -- Reuse atIndicesStable's monomorphic Int/Double kernels.                     case atIndicesStable (parClampNonNeg indices) (UnboxedColumn Nothing v) of                         UnboxedColumn _ dat -> UnboxedColumn (withBm srcBm) dat                         other -> other {-# INLINE gatherWithSentinel #-} -{- | An @n@-row payload for a gather whose source is empty: every index is a-sentinel, so the sentinel bitmap masks every row and no element is ever read.-The source vector is passed only to fix the element type.--}+-- | An @n@-row payload for a gather whose source is empty. allNullBoxed :: Int -> VB.Vector a -> VB.Vector a allNullBoxed n _ = VB.replicate n (error "gatherWithSentinel: null row forced") @@ -540,6 +600,30 @@ getIndicesUnboxed indices xs = VU.generate (VU.length indices) (\i -> xs VU.! (indices VU.! i)) {-# INLINE getIndicesUnboxed #-} +-- | Positions of the 'True' entries, in order, computed in parallel.+maskIndices :: VU.Vector Bool -> VU.Vector Int+maskIndices mask+    | not (shouldParallelize parThreshold n) = VU.findIndices id mask+    | otherwise = unsafePerformIO $ do+        let chunks = splitChunkRange capabilities n+            countChunk (lo, hi) = pure $! countFrom lo hi 0+            countFrom !i !hi !c+                | i >= hi = c+                | otherwise = countFrom (i + 1) hi (c + fromEnum (VU.unsafeIndex mask i))+        counts <- forkJoin (map countChunk chunks)+        let starts = scanl (+) 0 counts+        out <- VUM.unsafeNew (last starts)+        let fill (lo, hi) start =+                let go !i !k+                        | i >= hi = pure ()+                        | VU.unsafeIndex mask i = VUM.unsafeWrite out k i >> go (i + 1) (k + 1)+                        | otherwise = go (i + 1) k+                 in go lo start+        forkJoin_ (zipWith fill chunks starts)+        VU.unsafeFreeze out+  where+    !n = VU.length mask+ findIndices ::     forall a.     (Columnable a) =>@@ -547,19 +631,25 @@     Column ->     Either DataFrameException (VU.Vector Int) findIndices predicate = \case-    BoxedColumn _ (v :: VB.Vector b) -> run v VG.convert-    UnboxedColumn _ (v :: VU.Vector b) -> run v id+    BoxedColumn bm (v :: VB.Vector b) -> run bm v+    UnboxedColumn bm (v :: VU.Vector b) -> run bm v     c@(PackedText _ _) -> findIndices predicate (materializePacked c)     c@(MergedColumn _ _) -> findIndices predicate (materializeMerged c)   where+    -- Null rows never match.     run ::         forall b v.         (Typeable b, VG.Vector v b, VG.Vector v Int) =>+        Maybe Bitmap ->         v b ->-        (v Int -> VU.Vector Int) ->         Either DataFrameException (VU.Vector Int)-    run column finalize = case testEquality (typeRep @a) (typeRep @b) of-        Just Refl -> Right . finalize $ VG.findIndices predicate column+    run bm column = case testEquality (typeRep @a) (typeRep @b) of+        Just Refl -> Right $ case nullsIn (VG.length column) bm of+            Nothing -> VG.convert (VG.findIndices predicate column)+            Just nulls ->+                VU.filter+                    (\i -> bitmapTestBit nulls i && predicate (VG.unsafeIndex column i))+                    (VU.enumFromN 0 (VG.length column))         Nothing ->             Left $                 TypeMismatchException@@ -570,94 +660,115 @@                         , errorColumnName = Nothing                         } --- | Fold (right) column with index.+-- | Fold (right) column with index, skipping null rows. ifoldrColumn ::     forall a b.     (Columnable a, Columnable b) =>     (Int -> a -> b -> b) -> b -> Column -> Either DataFrameException b ifoldrColumn f acc = \case-    BoxedColumn _ column -> foldrWorker column-    UnboxedColumn _ column -> foldrWorker column+    BoxedColumn bm column -> foldrWorker bm column+    UnboxedColumn bm column -> foldrWorker bm column     c@(PackedText _ _) -> ifoldrColumn f acc (materializePacked c)     c@(MergedColumn _ _) -> ifoldrColumn f acc (materializeMerged c)   where     foldrWorker ::         forall c v.         (Typeable c, VG.Vector v c) =>+        Maybe Bitmap ->         v c ->         Either DataFrameException b-    foldrWorker vec = case testEquality (typeRep @a) (typeRep @c) of-        Just Refl -> pure $ VG.ifoldr f acc vec+    foldrWorker bm vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl -> pure $ case nullsIn (VG.length vec) bm of+            Nothing -> VG.ifoldr f acc vec+            Just nulls -> VG.ifoldr (\i x r -> if bitmapTestBit nulls i then f i x r else r) acc vec         Nothing ->             Left $                 TypeMismatchException-                    ( MkTypeErrorContext+                    MkTypeErrorContext                         { userType = Right (typeRep @a)                         , expectedType = Right (typeRep @c)                         , callingFunctionName = Just "ifoldrColumn"                         , errorColumnName = Nothing                         }-                    ) +-- | Left fold over a column, skipping null rows. foldlColumn ::     forall a b.     (Columnable a, Columnable b) =>     (b -> a -> b) -> b -> Column -> Either DataFrameException b foldlColumn f acc = \case-    BoxedColumn _ column -> foldlWorker column-    UnboxedColumn _ column -> foldlWorker column+    BoxedColumn bm column -> foldlWorker bm column+    UnboxedColumn bm (column :: VU.Vector c)+        | Just Refl <- testEquality (typeRep @a) (typeRep @c)+        , Nothing <- nullsIn (VU.length column) bm ->+            pure (fromMaybe VU.foldl' compiledFoldl f acc column)+        | otherwise -> foldlWorker bm column     c@(PackedText _ _) -> foldlColumn f acc (materializePacked c)     c@(MergedColumn _ _) -> foldlColumn f acc (materializeMerged c)   where     foldlWorker ::         forall c v.         (Typeable c, VG.Vector v c) =>+        Maybe Bitmap ->         v c ->         Either DataFrameException b-    foldlWorker vec = case testEquality (typeRep @a) (typeRep @c) of-        Just Refl -> pure $ VG.foldl' f acc vec+    foldlWorker bm vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl -> pure $ case nullsIn (VG.length vec) bm of+            Nothing -> VG.foldl' f acc vec+            Just nulls -> VG.ifoldl' (\r i x -> if bitmapTestBit nulls i then f r x else r) acc vec         Nothing ->             Left $                 TypeMismatchException-                    ( MkTypeErrorContext+                    MkTypeErrorContext                         { userType = Right (typeRep @a)                         , expectedType = Right (typeRep @c)-                        , callingFunctionName = Just "ifoldrColumn"+                        , callingFunctionName = Just "foldlColumn"                         , errorColumnName = Nothing                         }-                    ) +-- | Left fold seeded with the first valid row, skipping null rows. foldl1Column ::     forall a.     (Columnable a) =>     (a -> a -> a) -> Column -> Either DataFrameException a foldl1Column f = \case-    BoxedColumn _ column -> foldl1Worker column-    UnboxedColumn _ column -> foldl1Worker column+    BoxedColumn bm column -> foldl1Worker bm column+    UnboxedColumn bm (column :: VU.Vector c)+        | Just Refl <- testEquality (typeRep @a) (typeRep @c)+        , Nothing <- nullsIn (VU.length column) bm ->+            pure (fromMaybe VU.foldl1' compiledFoldl1 f column)+        | otherwise -> foldl1Worker bm column     c@(PackedText _ _) -> foldl1Column f (materializePacked c)     c@(MergedColumn _ _) -> foldl1Column f (materializeMerged c)   where     foldl1Worker ::         forall c v.         (Typeable c, VG.Vector v c) =>+        Maybe Bitmap ->         v c ->         Either DataFrameException a-    foldl1Worker vec = case testEquality (typeRep @a) (typeRep @c) of-        Just Refl -> pure $ VG.foldl1' f vec+    foldl1Worker bm vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl -> case nullsIn (VG.length vec) bm of+            Nothing -> pure (VG.foldl1' f vec)+            Just nulls ->+                let step r i x+                        | not (bitmapTestBit nulls i) = r+                        | otherwise = Just $! maybe x (`f` x) r+                 in maybe+                        (Left (EmptyDataSetException "foldl1Column"))+                        Right+                        (VG.ifoldl' step Nothing vec)         Nothing ->             Left $                 TypeMismatchException-                    ( MkTypeErrorContext+                    MkTypeErrorContext                         { userType = Right (typeRep @a)                         , expectedType = Right (typeRep @c)                         , callingFunctionName = Just "foldl1Column"                         , errorColumnName = Nothing                         }-                    ) -{- | O(n) Seedless fold over groups using the first element of each group as seed.-Like 'foldDirectGroups' but for the case where no initial accumulator is available.--}+-- | O(n) Seedless fold over groups using each group's first element as seed. foldl1DirectGroups ::     forall a.     (Columnable a) =>@@ -669,31 +780,46 @@ foldl1DirectGroups f col valueIndices offsets     | VU.length offsets <= 1 = pure $ fromVector @a VB.empty     | otherwise = case col of-        UnboxedColumn _ (vec :: VU.Vector d) -> UnboxedColumn Nothing <$> foldl1Worker vec-        BoxedColumn _ (vec :: VB.Vector d) -> BoxedColumn Nothing <$> foldl1Worker vec+        UnboxedColumn bm (vec :: VU.Vector d) -> uncurry UnboxedColumn <$> foldl1Worker bm vec+        BoxedColumn bm (vec :: VB.Vector d) -> uncurry BoxedColumn <$> foldl1Worker bm vec         PackedText _ _ -> foldl1DirectGroups f (materializePacked col) valueIndices offsets         MergedColumn _ _ -> foldl1DirectGroups f (materializeMerged col) valueIndices offsets   where+    nGroups = VU.length offsets - 1     foldl1Worker ::         forall c v.         (Typeable c, VG.Vector v c) =>+        Maybe Bitmap ->         v c ->-        Either DataFrameException (v c)-    foldl1Worker vec = case testEquality (typeRep @a) (typeRep @c) of-        Just Refl ->-            Right $-                VG.generate (VU.length offsets - 1) foldGroup+        Either DataFrameException (Maybe Bitmap, v c)+    foldl1Worker bm vec = case testEquality (typeRep @a) (typeRep @c) of+        Just Refl -> Right $ case nullsIn (VG.length vec) bm of+            Nothing -> (Nothing, VG.generate nGroups foldGroup)+            Just nulls ->+                -- A group of all-null rows has no seed, so it is null.+                let groups = VG.generate nGroups (foldValidGroup nulls) :: VB.Vector (Maybe c)+                    empty = VU.filter (isNothing . VB.unsafeIndex groups) (VU.enumFromN 0 nGroups)+                    placeholder k = VG.unsafeIndex vec (VU.unsafeIndex valueIndices (VU.unsafeIndex offsets k))+                 in ( Just (buildBitmapFromNulls' nGroups empty)+                    , VG.generate nGroups (\k -> fromMaybe (placeholder k) (VB.unsafeIndex groups k))+                    )           where+            row i = VG.unsafeIndex vec (VU.unsafeIndex valueIndices i)             foldGroup k =                 let !s = VU.unsafeIndex offsets k                     !e = VU.unsafeIndex offsets (k + 1)-                    !seed = VG.unsafeIndex vec (VU.unsafeIndex valueIndices s)-                 in go (s + 1) e seed+                 in go (s + 1) e (row s)             go !i !e !acc                 | i >= e = acc-                | otherwise =-                    go (i + 1) e $!-                        f acc (VG.unsafeIndex vec (VU.unsafeIndex valueIndices i))+                | otherwise = go (i + 1) e $! f acc (row i)+            foldValidGroup nulls k = goValid (VU.unsafeIndex offsets k) Nothing+              where+                !e = VU.unsafeIndex offsets (k + 1)+                goValid !i acc+                    | i >= e = acc+                    | not (bitmapTestBit nulls (VU.unsafeIndex valueIndices i)) =+                        goValid (i + 1) acc+                    | otherwise = goValid (i + 1) (Just $! maybe (row i) (`f` row i) acc)         Nothing ->             Left $                 TypeMismatchException@@ -705,10 +831,7 @@                         } {-# INLINEABLE foldl1DirectGroups #-} -{- | O(n) fold over groups by scanning the column linearly (rowToGroup[i] = group-of row i). Random writes hit the small per-group accumulator array; when @acc@ is-unboxable that array is unboxed, avoiding pointer indirection.--}+-- | O(n) fold over groups by scanning the column linearly. foldLinearGroups ::     forall b acc.     (Columnable b, Columnable acc) =>@@ -721,29 +844,32 @@ foldLinearGroups f seed col rowToGroup nGroups     | nGroups == 0 = Right (fromVector @acc VB.empty)     | otherwise = case col of-        UnboxedColumn _ (vec :: VU.Vector d) -> foldLinearWorker vec-        BoxedColumn _ (vec :: VB.Vector d) -> foldLinearWorker vec+        UnboxedColumn bm (vec :: VU.Vector d) -> foldLinearWorker bm vec+        BoxedColumn bm (vec :: VB.Vector d) -> foldLinearWorker bm vec         PackedText _ _ ->             foldLinearGroups f seed (materializePacked col) rowToGroup nGroups         MergedColumn _ _ ->             foldLinearGroups f seed (materializeMerged col) rowToGroup nGroups   where+    -- Null rows are skipped; a group with no valid rows keeps the seed.     foldLinearWorker ::         forall c v.         (Typeable c, VG.Vector v c) =>+        Maybe Bitmap ->         v c ->         Either DataFrameException Column-    foldLinearWorker vec = case testEquality (typeRep @b) (typeRep @c) of+    foldLinearWorker bm vec = case testEquality (typeRep @b) (typeRep @c) of         Just Refl ->-            Right $-                unsafePerformIO $-                    runWith-                        ( \readAt writeAt ->-                            VG.iforM_ vec $ \row x -> do-                                let !k = VG.unsafeIndex rowToGroup row-                                cur <- readAt k-                                writeAt k $! f cur x-                        )+            let nulls = nullsIn (VG.length vec) bm+             in Right $+                    unsafePerformIO $+                        runWith+                            ( \readAt writeAt ->+                                VG.iforM_ vec $ \row x -> when (validAt nulls row) $ do+                                    let !k = VG.unsafeIndex rowToGroup row+                                    cur <- readAt k+                                    writeAt k $! f cur x+                            )         Nothing ->             Left $                 TypeMismatchException@@ -767,33 +893,34 @@     {-# INLINE runWith #-} {-# INLINEABLE foldLinearGroups #-} +-- | The first valid row. headColumn :: forall a. (Columnable a) => Column -> Either DataFrameException a headColumn = \case-    BoxedColumn _ col -> headWorker col-    UnboxedColumn _ col -> headWorker col+    BoxedColumn bm col -> headWorker bm col+    UnboxedColumn bm col -> headWorker bm col     c@(PackedText _ _) -> headColumn (materializePacked c)     c@(MergedColumn _ _) -> headColumn (mergedHead c)   where     headWorker ::         forall c v.         (Typeable c, VG.Vector v c) =>+        Maybe Bitmap ->         v c ->         Either DataFrameException a-    headWorker vec = case testEquality (typeRep @a) (typeRep @c) of+    headWorker bm vec = case testEquality (typeRep @a) (typeRep @c) of         Just Refl ->-            if VG.null vec-                then Left (EmptyDataSetException "headColumn")-                else pure (VG.head vec)+            case VU.find (validAt bm) (VU.enumFromN 0 (VG.length vec)) of+                Nothing -> Left (EmptyDataSetException "headColumn")+                Just i -> pure (VG.unsafeIndex vec i)         Nothing ->             Left $                 TypeMismatchException-                    ( MkTypeErrorContext+                    MkTypeErrorContext                         { userType = Right (typeRep @a)                         , expectedType = Right (typeRep @c)                         , callingFunctionName = Just "headColumn"                         , errorColumnName = Nothing                         }-                    )  -- | An internal, column version of zip. zipColumns :: Column -> Column -> Column@@ -801,10 +928,10 @@ zipColumns l r@(MergedColumn _ _) = zipColumns l (materializeMerged r) zipColumns l@(PackedText _ _) r = zipColumns (materializePacked l) r zipColumns l r@(PackedText _ _) = zipColumns l (materializePacked r)-zipColumns (BoxedColumn _ column) (BoxedColumn _ other) = zipVectors column other-zipColumns (BoxedColumn _ column) (UnboxedColumn _ other) = zipVectors column other-zipColumns (UnboxedColumn _ column) (BoxedColumn _ other) = zipVectors column other-zipColumns (UnboxedColumn _ column) (UnboxedColumn _ other) = zipVectors column other+zipColumns (BoxedColumn bmL column) (BoxedColumn bmR other) = withBitmap (bothValid bmL bmR) (zipVectors column other)+zipColumns (BoxedColumn bmL column) (UnboxedColumn bmR other) = withBitmap (bothValid bmL bmR) (zipVectors column other)+zipColumns (UnboxedColumn bmL column) (BoxedColumn bmR other) = withBitmap (bothValid bmL bmR) (zipVectors column other)+zipColumns (UnboxedColumn bmL column) (UnboxedColumn bmR other) = withBitmap (bothValid bmL bmR) (zipVectors column other) {-# INLINE zipColumns #-}  zipVectors ::@@ -827,24 +954,55 @@     forall a b c.     (Columnable a, Columnable b, Columnable c) =>     (a -> b -> c) -> Column -> Column -> Either DataFrameException Column-zipWithColumns f (UnboxedColumn bmL (column :: VU.Vector d)) (UnboxedColumn bmR (other :: VU.Vector e)) = case testEquality (typeRep @a) (typeRep @d) of-    Just Refl -> case testEquality (typeRep @b) (typeRep @e) of-        Just Refl-            | isNothing bmL-            , isNothing bmR ->-                pure $ case sUnbox @c of-                    STrue ->-                        let !n = min (VU.length column) (VU.length other)-                         in UnboxedColumn Nothing $-                                parGenerateUnboxed n $ \i ->-                                    f (VU.unsafeIndex column i) (VU.unsafeIndex other i)-                    SFalse -> fromVector $ VB.zipWith f (VG.convert column) (VG.convert other)-        _ -> zipWithColumnsGeneral f (UnboxedColumn bmL column) (UnboxedColumn bmR other)-    Nothing -> zipWithColumnsGeneral f (UnboxedColumn bmL column) (UnboxedColumn bmR other)--- TODO: mchavinda - reuse pattern from interpret where we augment the--- error at the end.-zipWithColumns f left right = zipWithColumnsGeneral f left right+zipWithColumns f left right = case (left, right) of+    (PackedText _ _, _) -> zipWithColumns f (materializePacked left) right+    (_, PackedText _ _) -> zipWithColumns f left (materializePacked right)+    (MergedColumn _ _, _) -> zipWithColumns f (materializeMerged left) right+    (_, MergedColumn _ _) -> zipWithColumns f left (materializeMerged right)+    ( UnboxedColumn Nothing (l :: VU.Vector d)+        , UnboxedColumn Nothing (r :: VU.Vector e)+        )+            | Just Refl <- testEquality (typeRep @a) (typeRep @d)+            , Just Refl <- testEquality (typeRep @b) (typeRep @e)+            , STrue <- sUnbox @c ->+                pure . UnboxedColumn Nothing $ case testEquality (typeRep @a) (typeRep @b) of+                    Just Refl -> fromMaybe zipLoop compiledZip f l r+                    Nothing -> zipLoop f l r+    (BoxedColumn bmL l, BoxedColumn bmR r) -> elements bmL l bmR r+    (BoxedColumn bmL l, UnboxedColumn bmR r) -> elements bmL l bmR r+    (UnboxedColumn bmL l, BoxedColumn bmR r) -> elements bmL l bmR r+    (UnboxedColumn bmL l, UnboxedColumn bmR r) -> elements bmL l bmR r+  where+    -- A row is null if either side is; f never sees a null row.+    elements ::+        forall d e vl vr.+        (Typeable d, Typeable e, VG.Vector vl d, VG.Vector vr e) =>+        Maybe Bitmap -> vl d -> Maybe Bitmap -> vr e -> Either DataFrameException Column+    elements bmL l bmR r+        | Just Refl <- testEquality (typeRep @a) (typeRep @d)+        , Just Refl <- testEquality (typeRep @b) (typeRep @e) =+            let !n = min (VG.length l) (VG.length r)+                g i = f (VG.unsafeIndex l i) (VG.unsafeIndex r i)+             in pure $ case (bothValid bmL bmR, sUnbox @c) of+                    (Nothing, STrue) -> UnboxedColumn Nothing (parGenerateUnboxed n g)+                    (Nothing, SFalse) -> fromVector @c (VB.generate n g)+                    (Just bm, STrue) -> UnboxedColumn (Just bm) $ case nullsIn n (Just bm) of+                        Nothing -> parGenerateUnboxed n g+                        Just nulls -> generateValidUnboxed nulls n g+                    (Just bm, SFalse) -> generateValidBoxed bm n g+        | otherwise = zipWithColumnsGeneral f left right +-- | Valid where both are valid; a missing bitmap means every row is valid.+bothValid :: Maybe Bitmap -> Maybe Bitmap -> Maybe Bitmap+bothValid (Just l) (Just r) = Just (andBitmaps l r)+bothValid l Nothing = l+bothValid Nothing r = r++withBitmap :: Maybe Bitmap -> Column -> Column+withBitmap bm (BoxedColumn _ v) = BoxedColumn bm v+withBitmap bm (UnboxedColumn _ v) = UnboxedColumn bm v+withBitmap _ c = c+ zipWithColumnsGeneral ::     forall a b c.     (Columnable a, Columnable b, Columnable c) =>@@ -861,16 +1019,10 @@         Left e -> Left e         Right right' -> pure $ fromVector $ VB.zipWith f left' right' {-# INLINE zipWithColumnsGeneral #-}-{-# INLINE zipWithColumns #-} --- writeColumn and freezeColumn' (CSV-ingest helpers) moved to--- DataFrame.IO.Internal.MutableColumn so the core column module does not--- need to depend on DataFrame.Internal.Parsing.+-- writeColumn and freezeColumn' moved to DataFrame.IO.Internal.MutableColumn. -{- | Freeze a mutable column into an @Either Text a@ column: every recorded-null position becomes @Left rawText@ (preserving the original input), every-other position becomes @Right v@. Used by CSV readers under 'EitherRead' mode.--}+-- | Freeze a mutable column into an @Either Text a@ column. freezeColumnEither :: [(Int, T.Text)] -> MutableColumn -> IO Column freezeColumnEither nulls (MBoxedColumn col) = do     frozen <- VB.unsafeFreeze col@@ -894,9 +1046,7 @@                     Nothing -> Right (c VU.! i) {-# INLINE freezeColumnEither #-} -{- | Promote a non-nullable column to a nullable one (add an all-valid bitmap).-No-op when already nullable.--}+-- | Promote a non-nullable column to a nullable one (all-valid bitmap). ensureOptional :: Column -> Column ensureOptional c@(MergedColumn _ _) = ensureOptional (materializeMerged c) ensureOptional c@(BoxedColumn (Just _) _) = c@@ -909,7 +1059,7 @@ 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.+-- | Fill the end of a column, up to n, with null rows. expandColumn :: Int -> Column -> Column expandColumn n c@(MergedColumn a b)     | n <= min (columnLength a) (columnLength b) = c@@ -943,7 +1093,7 @@                 VU.unsafeFreeze mv          in UnboxedColumn newBm newCol --- | Fills the beginning of a column, up to n, with null rows. Does nothing if column has length >= n.+-- | Fill the beginning of a column, up to n, with null rows. leftExpandColumn :: Int -> Column -> Column leftExpandColumn n c@(MergedColumn a b)     | n <= min (columnLength a) (columnLength b) = c@@ -980,9 +1130,7 @@                 VU.unsafeFreeze mv          in UnboxedColumn newBm newCol -{- | Concatenates two columns.-Returns Nothing if the columns are of different types.--}+-- | Concatenates two columns; fails if the column types differ. mappendColumns :: Column -> Column -> Either DataFrameException Column mappendColumns left right = case (left, right) of     (MergedColumn _ _, _) -> mappendColumns (materializeMerged left) right@@ -1031,15 +1179,9 @@                         }                     ) -{- | Like 'mappendColumns' but also combines columns of different types by wrapping-values in 'Either' (e.g. @[1,2]@ and @["a","b"]@ become-@[Left 1, Left 2, Right "a", Right "b"]@).--}+-- ike 'mappendColumns' but wraps values in 'Either' across types. -{- | O(n) Concatenate a list of same-type columns in a single allocation.-All columns must have the same constructor and element type (as they will-within a single Parquet column). Calls 'error' on mismatch.--}+-- | O(n) Concatenate a list of same-type columns in a single allocation. concatManyColumns :: [Column] -> Column concatManyColumns [] = fromList ([] :: [Maybe Int]) concatManyColumns [c] = c@@ -1155,22 +1297,7 @@ mappendColumnsEither (UnboxedColumn _ left) (BoxedColumn _ right) =     BoxedColumn Nothing $ fmap Left (VG.convert left) <> fmap Right right ----------------------------------------------------------------------------------- Fused multi-column gather----------------------------------------------------------------------------------{- | Gather ONE in-bounds index vector through several columns in a single-parallel pass. Result columns are identical to @map (atIndicesStable ixs)@;-columns whose shape the fused kernel does not cover (bitmapped, boxed, merged,-unusual element types) fall back to per-column 'atIndicesStable'. All fused-outputs are backed by one shared deferred computation: forcing ANY of them-runs the single pass that fills ALL of them. The pass reads the index vector-once per block instead of once per column, and every iteration of a block-issues each column's independent random load back-to-back, so their cache/TLB-misses overlap instead of forming one latency chain per column (the aggregate-key materialization of a 1e8-group result was 6 sequential latency-bound-passes without this).--}+-- | Gather one in-bounds index vector through several columns in one pass. atIndicesStableMulti :: VU.Vector Int -> [Column] -> [Column] atIndicesStableMulti ixs cols =     let specs = map mgClassify cols@@ -1194,9 +1321,7 @@       MGInt !(VU.Vector Int)     | -- | Clean unboxed 'Double' payload.       MGDouble !(VU.Vector Double)-    | {- | Packed text with an 'Int32' selector (base row count necessarily-      fits 'Int32'); carries the payload for rebuilding and the canon flag.-      -}+    | -- | Packed text with an 'Int32' selector and an 'Int32'-sized base.       MGSel32 !PackedTextData !(VU.Vector Int32)     | -- | Packed text with an 'Int' selector and an 'Int32'-sized base.       MGSel64To32 !PackedTextData !(VU.Vector Int)@@ -1221,11 +1346,7 @@ mgInt32Max :: Int mgInt32Max = fromIntegral (maxBound :: Int32) -{- | Run the fused pass over the fusable specs; element @r@ of the result is-spec @r@'s gathered column. Deferred as one shared thunk (see-'atIndicesStableMulti'). Pure w.r.t. its immutable inputs, so the-'unsafePerformIO' is safe.--}+-- | Run the fused pass over the fusable specs, one gathered column per spec. multiGatherRun :: VU.Vector Int -> [MGSpec] -> VB.Vector Column multiGatherRun ixs specs = unsafePerformIO $ do     let !n = VU.length ixs@@ -1242,11 +1363,7 @@     VB.fromList <$> mapM snd opened {-# NOINLINE multiGatherRun #-} -{- | Allocate a spec's destination; return its block-fill action and its-finalizer. Selector gathers reproduce 'packedGather' exactly (composition-clamps against the source selector length and base row count); unboxed gathers-reproduce the unclamped 'parBackpermuteUnboxed'.--}+-- | Allocate a spec's destination; return its block-fill action and finalizer. mgOpen :: Int -> VU.Vector Int -> MGSpec -> IO (Int -> Int -> IO (), IO Column) mgOpen n ixs spec = case spec of     MGInt v -> do@@ -1370,3 +1487,43 @@             VUM.unsafeWrite dst i out             go (i + 1) {-# NOINLINE mgFillSel64To64 #-}++-- | A packed text column of @f i@ per row, built chunk by chunk in parallel.+packedMapText :: Int -> (Int -> T.Text) -> Column+packedMapText n f = unsafePerformIO $ do+    chunks <-+        if shouldParallelize parThreshold n+            then parallelChunks parThreshold n (\lo hi -> pure $! packRows lo hi)+            else pure [packRows 0 n]+    let totalBytes = sum [VU.unsafeIndex o (VU.length o - 1) | (_, o) <- chunks]+        totalRows = sum [VU.length o - 1 | (_, o) <- chunks]+    pure $! runST $ do+        arr <- A.new (max 1 totalBytes)+        offs <- VUM.unsafeNew (totalRows + 1)+        VUM.unsafeWrite offs 0 0+        let splice !byteBase !rowBase ((a, o) : rest) = do+                let rows = VU.length o - 1+                    used = VU.unsafeIndex o rows+                A.copyI used arr byteBase a 0+                forM_ [1 .. rows] $ \i ->+                    VUM.unsafeWrite offs (rowBase + i) (byteBase + VU.unsafeIndex o i)+                splice (byteBase + used) (rowBase + rows) rest+            splice _ _ [] = pure ()+        splice 0 0 chunks+        farr <- A.unsafeFreeze arr+        foffs <- VU.unsafeFreeze offs+        pure (PackedText Nothing (mkPackedContiguous farr foffs))+  where+    -- One chunk's rows as (bytes, offsets); offsets are chunk-relative.+    packRows lo hi =+        let texts = VB.generate (hi - lo) (\i -> f (lo + i))+            lens = VU.generate (hi - lo) (textBytes . VB.unsafeIndex texts)+            offs = VU.scanl' (+) 0 lens+            bytes = VU.unsafeIndex offs (hi - lo)+            arr = runST $ do+                m <- A.new (max 1 bytes)+                VB.iforM_ texts $ \i (Text src off len) ->+                    A.copyI len m (VU.unsafeIndex offs i) src off+                A.unsafeFreeze m+         in (arr, offs)+    textBytes (Text _ _ len) = len

@@ -15,10 +15,7 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} -{- |-Predicates and introspection over a 'Column': representation tests, null/type-queries, and the human-readable type descriptions used in error messages.--}+-- | Predicates and introspection over a 'Column'. module DataFrame.Internal.Column.Properties where  import qualified Data.Text as T@@ -82,9 +79,7 @@     Just Refl -> True isNumeric (PackedText _ _) = False -{- | Whether the column stores element type @a@. For nullable columns, also-'True' when @a = Maybe b@ and the column stores @b@ internally.--}+-- | Whether the column stores element type @a@, looking through 'Maybe'. hasElemType :: forall a. (Columnable a) => Column -> Bool hasElemType = \case     BoxedColumn bm (_column :: VB.Vector b) -> checkBoxed bm (typeRep @b)@@ -116,9 +111,7 @@     PackedText (Just _) _ -> "NullableBoxed"     MergedColumn _ _ -> columnVersionString (mergedHead column) -{- | An internal/debugging function to get the type stored in the outermost vector-of a column.--}+-- | Debugging helper: the type stored in a column's outermost vector. columnTypeString :: Column -> String columnTypeString column = case column of     BoxedColumn Nothing (_ :: VB.Vector a) -> show (typeRep @a)

@@ -22,16 +22,11 @@  type Columnable' a = (Typeable a, Show a, Eq a) -{- | Inline replacement for @Data.These.These@ to keep @dataframe-core@ free-of the @these@ package dependency. Only the three constructors and the-derived classes are used internally.--}+-- | Inline replacement for @Data.These.These@, avoiding the @these@ package. data These a b = This a | That b | These a b     deriving (Eq, Ord, Show, Read, Functor, Foldable, Traversable) -{- | A type with column representations used to select the-"right" representation when specializing the `toColumn` function.--}+-- | Column representations used to select the right 'toColumn' behaviour. data Rep     = RBoxed     | RUnboxed@@ -136,9 +131,7 @@  type FloatingIf a = When (FloatingTypes a) (Real a, Fractional a) -{- | Numeric type promotion: resolves the common type for mixed arithmetic.-Double dominates over Float/Int; Float dominates over Int; same types stay unchanged.--}+-- | Numeric type promotion: resolves the common type for mixed arithmetic. type family Promote (a :: Type) (b :: Type) :: Type where     Promote a a = a     Promote Double _ = Double@@ -151,9 +144,7 @@     Promote _ Int32 = Int32     Promote a _ = a -{- | Like 'Promote', but integral × integral → Double for use with './' .-Double\/Float still dominate; any two integral types (same or mixed) become Double.--}+-- | Like 'Promote' but integral by integral gives Double, for './'. type family PromoteDiv (a :: Type) (b :: Type) :: Type where     PromoteDiv Double _ = Double     PromoteDiv _ Double = Double

@@ -49,11 +49,7 @@ shouldParallelize threshold n = n >= threshold && capabilities > 1 {-# INLINE shouldParallelize #-} -{- | Row count below which a fan-out does not pay for itself, for the fixed-cost-per-row loops (grouping, the aggregation kernels). Pair it with-'shouldParallelize'. Kernels with a materially different per-row cost —-the join probe, the radix sort — set their own thresholds.--}+-- | Row count below which a fan-out does not pay for itself. parThreshold :: Int parThreshold = 200000 @@ -68,8 +64,7 @@         , len > 0         ]   where-    -- A non-positive width is a caller bug; clamp to one chunk rather than-    -- return [] and silently skip the rows.+    -- A non-positive width is a caller bug; clamp to one chunk.     !k' = max 1 k     (!q, !r) = n `quotRem` k' @@ -79,11 +74,7 @@     | otherwise = splitChunkRange capabilities n {-# INLINE chunksFor #-} -{- | Adjacent pairs of a precomputed bounds vector of length @caps + 1@. Empty-ranges are NOT dropped: the bounds are the caller's and the slot count is-often meaningful. Endpoints are forced here, on the spawning thread, so a-worker never starts by evaluating an index thunk that retains @bs@.--}+-- | Adjacent pairs of a precomputed bounds vector of length @caps + 1@. boundsChunks :: Int -> VU.Vector Int -> [(Int, Int)] boundsChunks caps bs =     [ (lo, hi)

@@ -2,10 +2,7 @@ {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} -{- | A flat, unboxed, open-addressing (linear-probe) hash table mapping a row's-key-hash to a dense group id, re-verifying the real key on every hash hit to-reject collisions. Runs in any 'PrimMonad' ('ST' for grouping, 'IO' per worker).--}+-- | A flat open-addressing hash table mapping row hash to a dense group id. module DataFrame.Internal.Data.HashTable (     HashTable (..),     newHashTable,@@ -18,9 +15,7 @@ import qualified Data.Vector.Unboxed.Mutable as VUM import Data.Word (Word64) -{- | An open-addressing linear-probe table. @htMask@ is @capacity - 1@ (capacity-is a power of two) and maps a hash to its home slot.--}+-- | An open-addressing linear-probe table with a power-of-two capacity. data HashTable s = HashTable     { htHash :: !(VUM.MVector s Int)     , htGroup :: !(VUM.MVector s Int)@@ -28,10 +23,7 @@     , 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.--}+-- | Smallest power of two strictly greater than @n@, at least 2. nextPow2Above :: Int -> Int nextPow2Above n = go 2   where@@ -40,10 +32,7 @@         | 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@).--}+-- | Allocate an empty table for up to @n@ groups at load factor under 0.5. newHashTable :: (PrimMonad m) => Int -> m (HashTable (PrimState m)) newHashTable n = do     let !cap = nextPow2Above (2 * max 1 n)@@ -53,10 +42,7 @@     pure (HashTable h g r (cap - 1)) {-# INLINE newHashTable #-} -{- | Look up @row@ (with precomputed @hash@) and return its dense group id: an-empty slot starts a new group via @nextGroup@, a stored-hash match is re-verified-with @eqRow@ before reuse. The 'Bool' is 'True' when a new group was created.--}+-- | Look up @row@ by hash and return its group id, adding one if new. htInsert ::     (PrimMonad m) =>     HashTable (PrimState m) ->@@ -94,11 +80,7 @@                     else go ((slot + 1) .&. mask) {-# INLINE htInsert #-} -{- | Home slot of a hash: the top @log2 cap@ bits of a Fibonacci multiply. The-row hash's final FxHash step is a multiply, which leaves its LOW bits poorly-diffused — raw-text keys cluster into contiguous linear-probe pileups when-slotted by @hash .&. mask@, so the home slot must come from the top bits.--}+-- | Home slot of a hash: the top bits of a Fibonacci multiply. homeSlot :: Int -> Int -> Int homeSlot !mask !hash =     fromIntegral

@@ -1,11 +1,6 @@ {-# LANGUAGE BangPatterns #-} -{- | Packed-text payload + byte-slice primitives. A 'PackedTextData' shares one-UTF-8 byte buffer across all rows of a string column, with @n+1@ row offsets, so-no per-row 'Data.Text.Text' header is materialized until decode is demanded.-Offsets and selection vectors are stored 'Int32' whenever their values fit-(Arrow-style), halving the per-row footprint of large string columns.--}+-- | Packed-text payload and byte-slice primitives over one shared buffer. module DataFrame.Internal.Data.PackedText (     PackedTextData (..),     PackedOffsets (..),@@ -47,9 +42,7 @@  ) import System.IO.Unsafe (unsafePerformIO) -{- | Row byte-offsets, physically 'Int32' when every value fits (total buffer-bytes < 2^31) and 'Int' otherwise. Values are non-negative byte positions.--}+-- | Row byte-offsets, 'Int32' when every value fits and 'Int' otherwise. data PackedOffsets     = Offs32 {-# UNPACK #-} !(VU.Vector Int32)     | Offs64 {-# UNPACK #-} !(VU.Vector Int)@@ -66,9 +59,7 @@ offCount (Offs64 v) = VU.length v {-# INLINE offCount #-} -{- | A selection layer mapping logical rows to base rows; @-1@ marks an-invalid/null row. 'Int32' when the base row count fits.--}+-- | A selection layer mapping logical rows to base rows; @-1@ marks null. data PackedSel     = Sel32 {-# UNPACK #-} !(VU.Vector Int32)     | Sel64 {-# UNPACK #-} !(VU.Vector Int)@@ -84,17 +75,7 @@ selLength (Sel64 v) = VU.length v {-# INLINE selLength #-} -{- | A shared UTF-8 byte buffer plus @n+1@ row offsets (base row @r@ spans bytes-@[offsets!r, offsets!(r+1))@); validity lives in the column's bitmap. @ptSel@ is-an optional selection layer letting a gather/join/sort result share the buffer.--@ptCanonicalSel@ marks a selection that is a canonical dictionary encoding:-equal byte slices always map to the same base row (codes). Set by dictionary-compaction; preserved by gather/take over an already-canonical selection (a-row keeps its code); 'False' for a gather over an unselected base, where two-logical rows can select different but equal-byted base rows. Grouping keys on-codes directly when it holds.--}+-- | A shared UTF-8 byte buffer plus @n+1@ row offsets and optional selection. data PackedTextData = PackedTextData     { ptBytes :: {-# UNPACK #-} !A.Array     , ptOffsets :: !PackedOffsets@@ -113,9 +94,7 @@     | otherwise = Offs64 offs {-# INLINE mkOffsets #-} -{- | Narrow an 'Int' base-row vector (@-1@ sentinels allowed) when the base-row count fits in 'Int32'.--}+-- | Narrow an 'Int' base-row vector to 'Int32' when the row count fits. mkSel :: Int -> VU.Vector Int -> PackedSel mkSel base rows     | base <= int32Max = Sel32 (VU.map fromIntegral rows)@@ -132,24 +111,14 @@ mkPackedContiguous32 arr offs = PackedTextData arr (Offs32 offs) Nothing False {-# INLINE mkPackedContiguous32 #-} -{- | Reindex a packed payload by a selection vector, sharing the byte buffer;-logical row @i@ becomes base row @indices!i@. A negative or out-of-range index-decodes to the empty slice. Composes with an existing selection; canonicality-survives composition (a kept row keeps its code) but not a first selection-over the unselected base.--}+-- | Reindex a packed payload by a selection vector, sharing the byte buffer. packedGather :: VU.Vector Int -> PackedTextData -> PackedTextData packedGather indices (PackedTextData arr offs msel canon) =     let !base = offCount offs - 1         canon' = case msel of             Nothing -> False             Just _ -> canon-        {- Base row for logical output row i composes any existing selection;-        only the SELECTOR is rebuilt (the byte buffer and offsets are shared-        untouched — a canonical dict column keeps its bytes). Each (source-        width x output width) shape runs a monomorphic closure-free kernel,-        generated directly at its final 'mkSel' width (one parallel pass).-        Same values and widths as the historical closure-driven build. -}+        -- Only the selector is rebuilt; bytes and offsets stay shared.         sel' = case msel of             Nothing                 | base <= int32Max -> Sel32 (gatherBaseTo32 indices base)@@ -162,8 +131,7 @@                 | otherwise -> Sel64 (gatherSel64To64 indices s base)      in PackedTextData arr offs (Just sel') canon' --- Monomorphic 'packedGather' selector kernels. All reproduce exactly--- @clamp r = if r >= 0 && r < base then r else -1@ over the composed pick.+-- Monomorphic 'packedGather' selector kernels; out-of-range picks give -1.  gatherBaseTo32 :: VU.Vector Int -> Int -> VU.Vector Int32 gatherBaseTo32 indices !base =@@ -227,18 +195,9 @@                     else -1 {-# NOINLINE gatherSel64To64 #-} -{- | Selector generate over 'parallelChunks_': one contiguous index chunk per-capability written into disjoint slices of one buffer — element @i@ depends-only on @f i@, so the result is bit-identical to 'VU.generate' at any @-N@.-Same policy as 'DataFrame.Internal.Column.Operations.parGenerateUnboxed', which-sits above this module.--}+-- elector generate over 'parallelChunks_', one chunk per capability. -{- | 'parGenSel' with an INLINE body: each monomorphic NOINLINE kernel above-gets its own copy of the fill loop with the pick function inlined — no unknown-closure call (or boxed result allocation) per element. Bit-identical results;-callers must be NOINLINE so the 'unsafePerformIO' runs once per call.--}+-- | 'parGenSel' with an INLINE body so each kernel inlines its pick function. parGenSelInline :: (VU.Unbox c) => Int -> (Int -> c) -> VU.Vector c parGenSelInline n f     | not (shouldParallelize parThreshold n) = VU.generate n f@@ -253,8 +212,6 @@ {-# INLINE parGenSelInline #-}  parGenSel :: (VU.Unbox c) => Int -> (Int -> c) -> VU.Vector c-{-# SPECIALIZE parGenSel :: Int -> (Int -> Int32) -> VU.Vector Int32 #-}-{-# SPECIALIZE parGenSel :: Int -> (Int -> Int) -> VU.Vector Int #-} parGenSel n f     | not (shouldParallelize parThreshold n) = VU.generate n f     | otherwise = unsafePerformIO $ do@@ -267,10 +224,7 @@         VU.unsafeFreeze mv {-# NOINLINE parGenSel #-} -{- | Take the first @k@ logical rows, sharing the byte buffer via a capped-selection layer. O(k), no byte copy or decode — cheap @take@/display on a-large packed column.--}+-- | Take the first @k@ logical rows, sharing the byte buffer, in O(k). packedTake :: Int -> PackedTextData -> PackedTextData packedTake k (PackedTextData arr offs msel canon) =     let !base = offCount offs - 1@@ -294,7 +248,7 @@ packedLength (PackedTextData _ _ (Just sel) _) = selLength sel {-# INLINE packedLength #-} --- | Raw byte slice for logical row @i@: @(buffer, offset, length)@. The hot accessor.+-- | Raw byte slice for logical row @i@: @(buffer, offset, length)@. packedSlice :: PackedTextData -> Int -> (A.Array, Int, Int) packedSlice p@(PackedTextData arr offs _ _) i =     let !r = baseRow p i@@ -304,18 +258,13 @@                 let o = offAt offs r in (arr, o, offAt offs (r + 1) - o) {-# INLINE packedSlice #-} -{- | The shared buffer + contiguous @n+1@ offsets when the payload is the-unselected base; a selected (gathered) payload returns 'Nothing' (its rows are-non-contiguous). Lets contiguous consumers skip the selection indirection.--}+-- | The shared buffer and contiguous offsets when the payload is unselected. packedRowOffsets :: PackedTextData -> Maybe (A.Array, PackedOffsets) packedRowOffsets (PackedTextData arr offs Nothing _) = Just (arr, offs) packedRowOffsets _ = Nothing {-# INLINE packedRowOffsets #-} -{- | On-demand single 'Data.Text.Text' for row @i@, using the same-validate-or-lenient decode as the freeze path so output is bit-identical.--}+-- | On-demand single 'Data.Text.Text' for row @i@, similar to freeze. packedIndexText :: PackedTextData -> Int -> T.Text packedIndexText p i =     let (arr, o, l) = packedSlice p i@@ -330,10 +279,7 @@     | 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.--}+-- | Byte-wise equality of two slices, agreeing with 'Text' equality. sliceEqBytes :: A.Array -> Int -> Int -> A.Array -> Int -> Int -> Bool sliceEqBytes a ao al b bo bl     | al /= bl = False@@ -345,10 +291,7 @@         | 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.--}+-- | Unsigned byte-lexicographic comparison of two slices (memcmp semantics). sliceCmpBytes :: A.Array -> Int -> Int -> A.Array -> Int -> Int -> Ordering sliceCmpBytes a ao al b bo bl = go 0   where

@@ -1,8 +1,6 @@ {-# LANGUAGE BangPatterns #-} -{- | UTF-8 validation and @decodeUtf8Lenient@-parity slice decoding used by-'DataFrame.Internal.Column.Builder' to turn shared byte buffers into 'Text'.--}+-- | UTF-8 validation and lenient slice decoding for shared byte buffers. module DataFrame.Internal.Data.PackedText.Utf8 (     isValidUtf8Slice,     isUtf8Boundary,@@ -30,18 +28,12 @@ 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.--}+-- | Whether a byte may start a code point (is not a continuation byte). isUtf8Boundary :: Word8 -> Bool isUtf8Boundary w = w < 0x80 || w >= 0xC0 {-# INLINE isUtf8Boundary #-} -{- | Decode a byte slice exactly like @decodeUtf8Lenient@: 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.--}+-- | Decode a byte slice exactly like @decodeUtf8Lenient@. lenientDecodeSlice :: A.Array -> Int -> Int -> T.Text lenientDecodeSlice arr off len = T.pack (go off)   where@@ -59,10 +51,7 @@             | j >= end = Nothing             | otherwise = loop (utf8DecodeContinue (A.unsafeIndex arr j) st cp) (j + 1) -{- | Slice forced 'Text' values off a shared array; row @i@ spans bytes-@[offs!i, offs!(i+1))@. Fast path validates the whole span once when every field-starts on a code-point boundary; else per-field validation with lenient decode.--}+-- | Slice forced 'Text' values off a shared array by row offsets. sliceTextVector :: A.Array -> VU.Vector Int -> VB.Vector T.Text sliceTextVector arr offs = VB.create $ do     mv <- VBM.unsafeNew n

@@ -25,7 +25,7 @@     empty,     columnNames,     insertColumn,-    fromNamedColumns,+    fromColumns,     getColumn,     unsafeGetColumn,     null,@@ -35,8 +35,10 @@     getRowAsText,     showElement,     stripJust,+    (=:), ) where +import qualified Data.List as L import qualified Data.Map as M import qualified Data.Text as T import qualified Data.Vector as V@@ -60,11 +62,13 @@  ) import DataFrame.Internal.Column (     Column (..),+    Columnable,     columnLength,     columnToTextVec,     columnTypeString,     expandColumn,     forceColumn,+    fromList,     materializeMerged,     parBackpermute2Int,     sliceColumn,@@ -84,32 +88,15 @@     -- ^ Keeps the column names in the order they were inserted in.     , dataframeDimensions :: (Int, Int)     -- ^ (rows, columns)-    , derivingExpressions :: M.Map T.Text UExpr+    , derivingExpressions :: M.Map T.Text SomeExpr     } -{- | Force evaluation of all columns in a DataFrame. Replacement for the removed-@instance NFData DataFrame@; used by the IO and lazy-executor strict paths.--}+-- | Force evaluation of all columns in a DataFrame. forceDataFrame :: DataFrame -> DataFrame forceDataFrame df@(DataFrame cols idx dims _exprs) =     V.foldl' (\() c -> forceColumn c) () cols `seq` idx `seq` dims `seq` df -{- | A record that contains information about how and what-rows are grouped in the dataframe. This can only be used with-`aggregate`.--Laziness contract: every field except 'valueIndices' and 'groupRepRows' is-computed eagerly by the grouping paths. 'valueIndices' may be a lazy thunk (the-low-cardinality direct grouping defers the O(n) stable placement pass until a-consumer — grouped median/top-k gathers, set ops, the interpreter's-group-slicing — actually demands the permutation); forcing it always yields the-unique stable counting-sort permutation of 'rowToGroup', so WHAT it evaluates to-is independent of when it is forced. 'groupRepRows' is the per-group-representative row (the first original row of each group, in canonical group-order); aggregation uses it to materialize the key columns without demanding-'valueIndices'. It may also be a thunk; its value always equals-@VU.map (valueIndices !) (VU.init offsets)@.--}+-- | Row grouping of a dataframe; can only be used with 'aggregate'. data GroupedDataFrame = GroupedInternal     { fullDataframe :: DataFrame     , groupedColumns :: [T.Text]@@ -117,19 +104,12 @@     -- ^ Rows sorted by group id (stable); possibly an unevaluated thunk.     , offsets :: VU.Vector Int     , rowToGroup :: VU.Vector Int-    {- ^ rowToGroup[i] = group index for row i.  Length n (one per row).-    Built once in 'groupBy'; reused by every aggregation.-    -}+    -- ^ rowToGroup[i] is the group index for row i, built once in 'groupBy'.     , groupRepRows :: VU.Vector Int     -- ^ First original row of each group, length nGroups. See laziness note.     } -{- | The historical five-field view of 'GroupedDataFrame'. Matching ignores-'groupRepRows'; building derives it lazily from @valueIndices@/@offsets@ (the-thunk only forces them if something actually reads the representative rows).-The direct grouping paths construct 'GroupedInternal' directly instead so the-representative rows never demand the placement pass.--}+-- | The historical five-field view of 'GroupedDataFrame'. pattern Grouped ::     DataFrame ->     [T.Text] ->@@ -146,11 +126,7 @@                 vis                 offs                 rtg-                {- Parallel gather (still a deferred thunk; forcing it forces-                vis/offs as before). Indices are grouping-produced and-                in-bounds by construction: offs has nGroups+1 entries and-                offs!g < length vis for every non-empty group. Values are-                identical to the historical @VU.map (vis !) (VU.init offs)@. -}+                -- Parallel gather; indices are in-bounds by construction.                 (parBackpermute2Int vis offs)  {-# COMPLETE Grouped #-}@@ -191,10 +167,7 @@                 | otherwise = T.empty          in T.unpack (body <> footer) -{- | Configures how a 'DataFrame' is rendered as text: 'maxRows' caps rendered-rows, 'maxColumns' collapses middle columns past the limit into an ellipsis, and-'maxCellWidth' truncates long cells. A non-positive field means \"no limit\".--}+-- | Configures how a 'DataFrame' is rendered as truncated text. data TruncateConfig = TruncateConfig     { maxRows :: Int     , maxColumns :: Int@@ -261,10 +234,7 @@             (map clipCell finalTypes)             (map (V.map clipCell) finalCols) -{- | Decide which columns survive horizontal truncation and where to splice the-ellipsis column. Splits with the extra column on the left for odd 'maxColumns';-inserts the ellipsis only when it actually saves space.--}+-- | Decide which columns survive horizontal truncation. pickColumns ::     Maybe TruncateConfig ->     Int ->@@ -309,10 +279,7 @@ columnNames = map fst . sortBy (compare `on` snd) . M.toList . columnIndices {-# INLINE columnNames #-} -{- | Insert a column into a DataFrame. If a column with the same name already-exists it is replaced in-place; otherwise the column is appended at the end.-Other columns are expanded (padded with nulls) to match the new row count.--}+-- | Insert a column into a DataFrame, replacing any column of the same name. insertColumn :: T.Text -> Column -> DataFrame -> DataFrame insertColumn name column d =     let@@ -334,22 +301,18 @@                     (n, c + 1)                     exprs --- | Build a DataFrame from a list of @(name, column)@ pairs using 'insertColumn'.-fromNamedColumns :: [(T.Text, Column)] -> DataFrame-fromNamedColumns = foldl (\df (name, column) -> insertColumn name column df) empty--{- | Safely retrieves a column by name from the dataframe.--Returns 'Nothing' if the column does not exist.+-- | Build a DataFrame from @(name, column)@ pairs using 'insertColumn'.+fromColumns :: [(T.Text, Column)] -> DataFrame+fromColumns = L.foldl' (\df (name, column) -> insertColumn name column df) empty -==== __Examples__+mkColumn :: (Columnable a) => T.Text -> [a] -> (T.Text, Column)+mkColumn name column = (name, fromList column) ->>> getColumn "age" df-Just (UnboxedColumn ...)+infix 1 =:+(=:) :: (Columnable a) => T.Text -> [a] -> (T.Text, Column)+(=:) = mkColumn ->>> getColumn "nonexistent" df-Nothing--}+-- | Safely retrieve a column by name, or 'Nothing' if it does not exist. getColumn :: T.Text -> DataFrame -> Maybe Column getColumn name df     | null df = Nothing@@ -357,19 +320,13 @@         i <- columnIndices df M.!? name         columns df V.!? i -{- | Retrieve a column by name, throwing 'ColumnsNotFoundException' if it does not-exist. Unsafe version of 'getColumn'; use when the column is certain to exist.--}+-- | Retrieve a column by name, throwing 'ColumnsNotFoundException' if absent. unsafeGetColumn :: T.Text -> DataFrame -> Column unsafeGetColumn name df = case getColumn name df of     Nothing -> throw $ ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)     Just col -> col -{- | Checks if the dataframe is empty (has no columns).--Returns 'True' if the dataframe has no columns, 'False' otherwise.-Note that a dataframe with columns but no rows is not considered null.--}+-- | Checks if the dataframe is empty (has no columns). null :: DataFrame -> Bool null df = V.null (columns df) 

@@ -1,7 +1,4 @@-{- | A minimal Wadler/Leijen-style document combinator and width-aware renderer.-A 'Doc' describes a layout abstractly; 'render' chooses where soft breaks become-newlines to fit a target width. 'Group' lays a region flat when it fits.--}+-- | A minimal Wadler/Leijen-style document combinator and width-aware renderer. module DataFrame.Internal.Display.Pretty (     Doc,     text,@@ -38,9 +35,7 @@ text :: String -> Doc text = Text -{- | A soft break: a single space when its enclosing 'group' fits the width,-otherwise a newline + current indentation.--}+-- | A soft break: a space when its group fits the width, else a newline. line :: Doc line = Line @@ -74,10 +69,7 @@ parens :: Doc -> Doc parens d = Text "(" <> d <> Text ")" -{- | Render @d@ bare when it fits flat on the current line, wrapped in parens when-it must break across lines. Keeps operator grouping unambiguous once a-sub-expression wraps, without parenthesis noise on one-line expressions.--}+-- | Render @d@ bare when it fits flat, in parens when it must break. parensWhenBroken :: Doc -> Doc parensWhenBroken d = Group (Alt d (parens d)) 

@@ -25,10 +25,7 @@ import qualified DataFrame.Internal.Display.Pretty as P import Type.Reflection (Typeable, typeOf, typeRep) -{- | Operators are an open typeclass: built-ins get their own 'Typeable' type so the-simplifier can match them by 'cast', and users can add instances. The generic-'UnUDF'/'BinUDF' carriers cover UDFs, dynamic-named, and arithmetic ops.--}+-- | Operators are an open typeclass so the simplifier can match them by 'cast'. class (Typeable op) => UnaryOp op where     unaryFn :: op a b -> a -> b     unaryName :: op a b -> T.Text@@ -108,25 +105,25 @@     Agg :: (Columnable a, Columnable b) => AggStrategy a b -> Expr b -> Expr a     Over :: (Columnable a) => [T.Text] -> Expr a -> Expr a -data UExpr where-    UExpr :: (Columnable a) => Expr a -> UExpr+data SomeExpr where+    SomeExpr :: (Columnable a) => Expr a -> SomeExpr -instance Show UExpr where-    show :: UExpr -> String-    show (UExpr expr) = show expr+instance Show SomeExpr where+    show :: SomeExpr -> String+    show (SomeExpr expr) = show expr -toUExpr :: (Columnable a) => Expr a -> UExpr-toUExpr = UExpr+toSomeExpr :: (Columnable a) => Expr a -> SomeExpr+toSomeExpr = SomeExpr -fromUExpr :: forall a. (Columnable a) => UExpr -> Maybe (Expr a)-fromUExpr (UExpr (expr :: Expr b)) = do+fromSomeExpr :: forall a. (Columnable a) => SomeExpr -> Maybe (Expr a)+fromSomeExpr (SomeExpr (expr :: Expr b)) = do     Refl <- testEquality (typeRep @a) (typeRep @b)     pure expr -type NamedExpr = (T.Text, UExpr)+type NamedExpr = (T.Text, SomeExpr)  toNamedExpr :: (Columnable a) => T.Text -> Expr a -> NamedExpr-toNamedExpr exprName expr = (exprName, UExpr expr)+toNamedExpr exprName expr = (exprName, SomeExpr expr)  instance (Num a, Columnable a) => Num (Expr a) where     (+) :: Expr a -> Expr a -> Expr a@@ -377,18 +374,15 @@         (Agg op inner) -> Agg op (replaceExpr old new inner)         (Over keys inner) -> Over keys (replaceExpr old new inner) -{- | Simultaneously substitute 'Col' references from a name→expression map in a-single parallel pass, so a swap like @{a ↦ col b, b ↦ col a}@ works. Raw-text-references (in 'CastWith', 'Over' keys) are left untouched; type mismatch raises.--}+-- | Substitute 'Col' references from a name-to-expression map in one pass. substituteColumns ::-    forall a. (Columnable a) => M.Map T.Text UExpr -> Expr a -> Expr a+    forall a. (Columnable a) => M.Map T.Text SomeExpr -> Expr a -> Expr a substituteColumns subs = go   where     go :: forall b. (Columnable b) => Expr b -> Expr b     go e@(Col name) = case M.lookup name subs of         Nothing -> e-        Just (UExpr (repl :: Expr c)) -> case testEquality (typeRep @b) (typeRep @c) of+        Just (SomeExpr (repl :: Expr c)) -> case testEquality (typeRep @b) (typeRep @c) of             Just Refl -> repl             Nothing ->                 error $@@ -429,17 +423,11 @@ getColumns (Agg _strategy expr) = getColumns expr getColumns (Over keys inner) = keys <> getColumns inner -{- | Render an expression as readable, width-aware pseudo-code at the default-width ('P.defaultWidth'). See 'prettyPrintWidth' to control wrapping.--}+-- | Render an expression as readable pseudo-code at the default width. prettyPrint :: Expr a -> String prettyPrint = prettyPrintWidth P.defaultWidth -{- | Render an expression as readable, width-aware pseudo-code: long binary chains-wrap onto aligned continuation lines, conditionals are laid out like Python-statements (@if cond@ / @else if cond@ / @else@ with 4-space indented branches, so-nested conditionals nest visually), and sub-exprs are parenthesized by precedence.--}+-- | Render an expression as readable, width-aware pseudo-code. prettyPrintWidth :: Int -> Expr a -> String prettyPrintWidth width = P.render width . toDoc 0   where

@@ -21,7 +21,7 @@     ),     Expr (Binary, Col, If, Lit, Unary),     NamedExpr,-    UExpr (UExpr),+    SomeExpr (SomeExpr),     UnUDF (MkUnaryOp),  ) import DataFrame.Internal.Expression.Operators.Nullable (@@ -50,7 +50,7 @@ (|>) = (&)  as :: (Columnable a) => Expr a -> T.Text -> NamedExpr-as expr colName = (colName, UExpr expr)+as expr colName = (colName, SomeExpr expr)  name :: (Show a) => Expr a -> T.Text name (Col n) = n@@ -64,6 +64,13 @@ ifThenElse :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a ifThenElse = If +caseWhen :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a+caseWhen = If++-- | Closes a 'caseWhen' chain with its fallback value.+orElse :: Expr a -> Expr a+orElse = id+ lit :: (Columnable a) => a -> Expr a lit = Lit @@ -247,9 +254,7 @@  -- Nullable-aware arithmetic operators -{- | Nullable-aware addition. Works for all combinations of nullable\/non-nullable operands.-@col \@Int "x" .+ col \@(Maybe Int) "y"  -- :: Expr (Maybe Int)@--}+-- | Nullable-aware addition for nullable and non-nullable operands. (.+) ::     ( NumericWidenOp (BaseType a) (BaseType b)     , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b)@@ -293,12 +298,9 @@     Expr (WidenResultDiv a b) (./) = lift2Decorated (applyNull2 (divArithOp (/))) "nulldiv" (Just "./") False 7 --- Nullable-aware comparison operators (three-valued logic: Nothing if either operand is Nothing)+-- Three-valued comparisons: Nothing if either operand is Nothing. -{- | Nullable-aware equality. Widens numeric operands to their common type,-so @Expr Double .== Expr Int@ typechecks. Returns @Maybe Bool@ when either-operand is nullable.--}+-- | Nullable-aware equality that widens numeric operands to a common type. (.==) ::     ( NumericWidenOp (BaseType a) (BaseType b)     , NullLift2Op a b Bool (NullCmpResult a b)@@ -342,9 +344,7 @@     Expr (NullCmpResult a b) (.>) = Binary NullGt -{- | Nullable-aware less-than-or-equal. Widens numeric operands to their-common type, so @Expr Double .<= Expr Int@ typechecks.--}+-- | Nullable-aware less-than-or-equal that widens numeric operands. (.<=) ::     ( NumericWidenOp (BaseType a) (BaseType b)     , NullLift2Op a b Bool (NullCmpResult a b)@@ -355,7 +355,7 @@     Expr (NullCmpResult a b) (.<=) = Binary NullLeq --- | Nullable-aware greater-than-or-equal. Widens numeric operands to their common type.+-- | Nullable-aware greater-than-or-equal that widens numeric operands. (.>=) ::     ( NumericWidenOp (BaseType a) (BaseType b)     , NullLift2Op a b Bool (NullCmpResult a b)@@ -372,7 +372,7 @@ (.||.) :: Expr Bool -> Expr Bool -> Expr Bool (.||.) = lift2Decorated (||) "or" (Just ".||.") True 2 --- | Nullable-aware logical AND. Returns @Maybe Bool@ when either operand is nullable.+-- | Nullable-aware logical AND, giving @Maybe Bool@ if either side is nullable. (.&&) ::     (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>     Expr a ->@@ -380,7 +380,7 @@     Expr (NullCmpResult a b) (.&&) = Binary NullAnd --- | Nullable-aware logical OR. Returns @Maybe Bool@ when either operand is nullable.+-- | Nullable-aware logical OR, giving @Maybe Bool@ if either side is nullable. (.||) ::     (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>     Expr a ->

@@ -9,21 +9,7 @@ {-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE UndecidableSuperClasses #-} -{- | Nullable-aware arithmetic and comparison operators ('.+', '.==', …) that work-transparently across nullable (@Maybe a@) and non-nullable (@a@) operands.-Functional dependencies infer the result type without annotations.--@--- Mixing nullable and non-nullable columns:-F.col \@Int \"x\" '.+' F.col \@(Maybe Int) \"y\"  -- :: Expr (Maybe Int)---- Both non-nullable (existing behaviour preserved):-F.col \@Int \"x\" '.+' F.col \@Int \"y\"           -- :: Expr Int---- Comparison with three-valued logic:-F.col \@(Maybe Int) \"x\" '.==' F.col \@Int \"y\"  -- :: Expr (Maybe Bool)-@--}+-- | Nullable-aware operators working across @Maybe a@ and @a@ operands. module DataFrame.Internal.Expression.Operators.Nullable (     -- * Type family     BaseType,@@ -61,21 +47,12 @@ import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.Column.Types (Promote, PromoteDiv) -{- | Strip one layer of 'Maybe'.--@-BaseType (Maybe a) = a-BaseType a         = a   -- for any non-Maybe type-@--}+-- | Strip one layer of 'Maybe'. type family BaseType a where     BaseType (Maybe a) = a     BaseType a = a -{- | Arithmetic binary operations that work over nullable and non-nullable operand-types. The functional dependency @a b -> c@ infers the result; the 'OVERLAPPABLE'-non-nullable instance yields to the specific @(Maybe a, Maybe a)@ one.--}+-- | Arithmetic binary operations over nullable and non-nullable operands. class     ( Columnable a     , Columnable b@@ -84,35 +61,20 @@     NullableArithOp a b c         | a b -> c     where-    {- | Lift an arithmetic function over the inner (non-Maybe) values.-    'Nothing' short-circuits: any 'Nothing' operand produces 'Nothing'.-    -}+    -- | Lift an arithmetic function over inner values; 'Nothing' wins.     nullArithOp ::         (BaseType a -> BaseType a -> BaseType a) ->         a ->         b ->         c -{- | Compute the result type of a nullable comparison.--@-NullCmpResult (Maybe a) b = Maybe Bool-NullCmpResult a (Maybe b) = Maybe Bool   -- when a is apart from Maybe-NullCmpResult a b         = Bool-@--Used by the comparison operators ('.==', '.<', etc.) so GHC infers the-return type without an explicit annotation.--}+-- | Compute the result type of a nullable comparison. type family NullCmpResult a b where     NullCmpResult (Maybe a) b = Maybe Bool     NullCmpResult a (Maybe b) = Maybe Bool     NullCmpResult a b = Bool -{- | Comparison binary operations over nullable and non-nullable operands. No-functional dependency on @e@; overlapping/overlappable instance pragmas pick the-unique most-specific instance from the concrete operand types.--}+-- | Comparison binary operations over nullable and non-nullable operands. class     ( Columnable a     , Columnable b@@ -120,18 +82,14 @@     ) =>     NullableCmpOp a b e     where-    {- | Lift a comparison function over the inner values (three-valued logic).-    Returns 'Nothing' when either operand is 'Nothing'.-    -}+    -- | Lift a comparison over inner values with three-valued logic.     nullCmpOp ::         (BaseType a -> BaseType a -> Bool) ->         a ->         b ->         e -{- | Non-nullable × Non-nullable: apply directly, no wrapping.-Arithmetic result is @a@; comparison result is @Bool@.--}+-- | Non-nullable by non-nullable: apply directly, no wrapping. instance     {-# OVERLAPPABLE #-}     (Columnable a, a ~ BaseType a) =>@@ -202,25 +160,9 @@     nullCmpOp _f _ Nothing = Nothing     nullCmpOp f (Just x) (Just y) = Just (f x y) --- ------------------------------------------------------------------------------ Generalized nullable lift (unary)--- -----------------------------------------------------------------------------{- | Lift a unary function over a column expression, propagating 'Nothing' (applied-directly when non-nullable, under 'Just' when @a = Maybe x@). Use via-'DataFrame.Functions.nullLift'.--}--{- | Compute the result type of a nullable unary lift.--@-NullLift1Result (Maybe a) r = Maybe r-NullLift1Result a         r = r        -- for any non-Maybe a-@+-- ift a unary function over a column expression, propagating 'Nothing'. -Used by 'DataFrame.Functions.nullLift' so GHC can infer the return type-without an explicit annotation.--}+-- | Compute the result type of a nullable unary lift. type family NullLift1Result a r where     NullLift1Result (Maybe a) r = Maybe r     NullLift1Result a r = r@@ -251,25 +193,9 @@     applyNull1 _ Nothing = Nothing     applyNull1 f (Just x) = Just (f x) --- ------------------------------------------------------------------------------ Generalized nullable lift (binary)--- -----------------------------------------------------------------------------{- | Lift a binary function over two column expressions, propagating 'Nothing': the-result is @Maybe r@ if either operand is nullable, else @r@. Use via-'DataFrame.Functions.nullLift2'.--}--{- | Compute the result type of a nullable binary lift.--@-NullLift2Result (Maybe a) b         r = Maybe r-NullLift2Result a         (Maybe b) r = Maybe r   -- when a is apart from Maybe-NullLift2Result a         b         r = r-@+-- ift a binary function over two column expressions, propagating 'Nothing'. -Used by 'DataFrame.Functions.nullLift2' so GHC can infer the return type.--}+-- | Compute the result type of a nullable binary lift. type family NullLift2Result a b r where     NullLift2Result (Maybe a) b r = Maybe r     NullLift2Result a (Maybe b) r = Maybe r@@ -321,15 +247,7 @@     applyNull2 _ _ Nothing = Nothing     applyNull2 f (Just x) (Just y) = Just (f x y) --- ------------------------------------------------------------------------------ Numeric widening--- -----------------------------------------------------------------------------{- | Widen two numeric base types to their promoted common type.--When @a ~ b@ the coercions are identity; otherwise one operand is widened-(e.g. 'Int' → 'Double').--}+-- | Widen two numeric base types to their promoted common type. class (Columnable (Promote a b)) => NumericWidenOp a b where     widen1 :: a -> Promote a b     widen2 :: b -> Promote a b@@ -364,7 +282,7 @@     widen1 = id     widen2 = fromIntegral --- | Apply an arithmetic function after widening both operands to their common type.+-- | Apply an arithmetic function after widening both operands. widenArithOp ::     forall a b.     (NumericWidenOp a b) =>@@ -374,7 +292,7 @@     Promote a b widenArithOp f x y = f (widen1 @a @b x) (widen2 @a @b y) --- | Apply a comparison function after widening both operands to their common type.+-- | Apply a comparison function after widening both operands. widenCmpOp ::     forall a b.     (NumericWidenOp a b) =>@@ -384,17 +302,10 @@     Bool widenCmpOp f x y = f (widen1 @a @b x) (widen2 @a @b y) --- | Result type of a widening binary operator, accounting for nullable wrappers.+-- | Result type of a widening binary operator, accounting for nullables. type WidenResult a b = NullLift2Result a b (Promote (BaseType a) (BaseType b)) --- ------------------------------------------------------------------------------ Division widening (integral × integral → Double)--- -----------------------------------------------------------------------------{- | Like 'NumericWidenOp' but uses 'PromoteDiv': integral×integral → Double.-Floating types still dominate (Double > Float), and any two integral types-(same or mixed) are both widened to Double.--}+-- | Like 'NumericWidenOp' but uses 'PromoteDiv' (integral gives Double). class (Columnable (PromoteDiv a b)) => DivWidenOp a b where     divWiden1 :: a -> PromoteDiv a b     divWiden2 :: b -> PromoteDiv a b@@ -462,6 +373,6 @@     PromoteDiv a b divArithOp f x y = f (divWiden1 @a @b x) (divWiden2 @a @b y) --- | Result type of a division-widening binary operator, accounting for nullable wrappers.+-- | Result type of a division-widening binary operator. type WidenResultDiv a b =     NullLift2Result a b (PromoteDiv (BaseType a) (BaseType b))

@@ -0,0 +1,111 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE UndecidableInstances #-}++-- | Class evidence for the supported column types, recovered from 'Typeable'.+module DataFrame.Internal.Expression.Reflect (+    Dict (..),+    RealUnbox,+    numDict,+    ordDict,+    fracDict,+    floatingDict,+    realDict,+    integralDict,+    realUnboxDict,+) where++import Data.Foldable (asum)+import Data.Int (Int16, Int32, Int64, Int8)+import Data.Kind (Constraint, Type)+import qualified Data.Text as T+import Data.Type.Equality (testEquality, (:~:) (Refl))+import qualified Data.Vector.Unboxed as VU+import Data.Word (Word16, Word32, Word64, Word8)+import Type.Reflection (Typeable, typeRep)++-- | A captured class dictionary.+data Dict (c :: Constraint) where+    Dict :: (c) => Dict c++-- | 'Real' and 'VU.Unbox' together, for variance, median and similar.+class (Real a, VU.Unbox a) => RealUnbox a++instance (Real a, VU.Unbox a) => RealUnbox a++-- | Evidence for @c a@ when @a@ is the concrete type @b@.+via ::+    forall b (c :: Type -> Constraint) a.+    (Typeable a, Typeable b, c b) => Maybe (Dict (c a))+via = case testEquality (typeRep @a) (typeRep @b) of+    Just Refl -> Just Dict+    Nothing -> Nothing++-- | The fixed-width integral types.+ints ::+    forall (c :: Type -> Constraint) a.+    ( Typeable a+    , c Int+    , c Int8+    , c Int16+    , c Int32+    , c Int64+    , c Word+    , c Word8+    , c Word16+    , c Word32+    , c Word64+    ) =>+    [Maybe (Dict (c a))]+ints =+    [ via @Int @c @a+    , via @Int8 @c @a+    , via @Int16 @c @a+    , via @Int32 @c @a+    , via @Int64 @c @a+    , via @Word @c @a+    , via @Word8 @c @a+    , via @Word16 @c @a+    , via @Word32 @c @a+    , via @Word64 @c @a+    ]++floats ::+    forall (c :: Type -> Constraint) a.+    (Typeable a, c Double, c Float) => [Maybe (Dict (c a))]+floats = [via @Double @c @a, via @Float @c @a]++numDict :: forall a. (Typeable a) => Maybe (Dict (Num a))+numDict = asum (ints @Num @a ++ [via @Integer @Num @a] ++ floats @Num @a)++ordDict :: forall a. (Typeable a) => Maybe (Dict (Ord a))+ordDict =+    asum+        ( ints @Ord @a+            ++ [via @Integer @Ord @a]+            ++ floats @Ord @a+            ++ [via @Bool @Ord @a, via @Char @Ord @a, via @T.Text @Ord @a, via @String @Ord @a]+        )++fracDict :: forall a. (Typeable a) => Maybe (Dict (Fractional a))+fracDict = asum (floats @Fractional @a)++floatingDict :: forall a. (Typeable a) => Maybe (Dict (Floating a))+floatingDict = asum (floats @Floating @a)++realDict :: forall a. (Typeable a) => Maybe (Dict (Real a))+realDict = asum (ints @Real @a ++ [via @Integer @Real @a] ++ floats @Real @a)++integralDict :: forall a. (Typeable a) => Maybe (Dict (Integral a))+integralDict = asum (ints @Integral @a ++ [via @Integer @Integral @a])++-- | 'Integer' is deliberately absent: it is 'Real' but not 'VU.Unbox'.+realUnboxDict :: forall a. (Typeable a) => Maybe (Dict (RealUnbox a))+realUnboxDict = asum (ints @RealUnbox @a ++ floats @RealUnbox @a)

@@ -122,9 +122,7 @@             Nothing -> Nothing asBoolLit _ = Nothing -{- | Polymorphic boolean literal: @Lit b@ for @Expr Bool@, @Lit (Just b)@ for-@Expr (Maybe Bool)@.--}+-- | Polymorphic boolean literal for @Expr Bool@ and @Expr (Maybe Bool)@. litBoolish :: forall a. (Columnable a) => Bool -> Maybe (Expr a) litBoolish v =     case testEquality (typeRep @a) (typeRep @Bool) of@@ -214,9 +212,7 @@                     Nothing -> Nothing litDouble _ = Nothing -{- | True for a column lifted from an integral type (never NaN): @toDouble (col …)@-or a column whose type is itself integral.--}+-- | True for a column lifted from an integral type (never NaN). integralColE :: forall c. (Columnable c) => Expr c -> Bool integralColE (Unary op _) = unaryName op == "toDouble" integralColE _ =@@ -258,9 +254,7 @@ litFalseGated UnknownOnNull = Nothing litFalseGated _ = litBoolish False -{- | Tautology to literal True is sound only for total (never-null) atoms; the-exhaustive-cover form additionally needs a non-NaN (integral) column.--}+-- | Tautology to literal True is sound only for total (never-null) atoms. litTrueTotal :: (Columnable a) => NullK -> Maybe (Expr a) litTrueTotal Total = litBoolish True litTrueTotal _ = Nothing@@ -317,9 +311,7 @@             | cA == CEq, cB == CEq -> if tA == tB then Just a else Nothing             | otherwise -> Nothing -{- | Build @col == t@ for the point-collapse rule; only strict @Expr Bool@ over a-@Double@ column (otherwise bail).--}+-- | Build @col == t@ for the point-collapse rule over a @Double@ column. pointEq :: forall a. (Columnable a) => Expr a -> Double -> Maybe (Expr a) pointEq atom lo = case testEquality (typeRep @a) (typeRep @Bool) of     Just Refl -> (\colE -> colE .==. Lit lo) <$> recoverColD atom@@ -351,10 +343,6 @@ outside t CLeq tb = t > tb outside _ _ _ = False --- ------------------------------------------------------------------------------ Path-condition entailment for fitted-tree pruning.--- ---------------------------------------------------------------------------- -- | A known same-column threshold fact accumulated along a tree path. data PredFact = PredFact !String !Cmp !Double @@ -362,10 +350,7 @@ factTrue :: Expr Bool -> Maybe PredFact factTrue e = (\a -> PredFact (aKey a) (aCmp a) (aThr a)) <$> atomOf e -{- | The fact a branch's false edge establishes (the negated condition). Only-sound for non-NaN (integral) columns — a NaN row takes the false edge too,-so @¬(x>t)@ is not a clean @x<=t@ bound for floats.--}+-- | The fact a branch's false edge establishes (the negated condition). factFalse :: Expr Bool -> Maybe PredFact factFalse e = do     a <- atomOf e@@ -380,9 +365,7 @@ negCmp CGeq = Just CLt negCmp _ = Nothing -{- | @entails facts cond@: 'Just' 'True' when the path facts force @cond@ true,-'Just' 'False' when they force it false, 'Nothing' when undecided.--}+-- | Whether the path facts force @cond@ true, false, or leave it undecided. entails :: [PredFact] -> Expr Bool -> Maybe Bool entails facts cond = do     a <- atomOf cond@@ -396,10 +379,7 @@         (d : _) -> Just d         [] -> Nothing -{- | Does the fact's solution set sit inside @cond@ ('Just' 'True'), disjoint-from it ('Just' 'False'), or neither ('Nothing')? Boundary strictness is-honoured: e.g. @x<=t@ does NOT entail @x<t@, and @x>=t ∧ x<=t@ is not empty.--}+-- | Does a fact's solution set sit inside, apart from, or across @cond@? factImplies :: (Cmp, Double) -> (Cmp, Double) -> Maybe Bool factImplies (fc, ft) (cc, tc)     | isLower fc, isLower cc, subset = Just True

@@ -33,6 +33,7 @@ import DataFrame.Internal.Column (     Column (..),     materializeMerged,+    widenToInt,  ) import DataFrame.Internal.Column.Bitmap (     Bitmap,@@ -41,16 +42,14 @@ import DataFrame.Internal.Column.Encode (dictEncodeColumnUpTo) import DataFrame.Internal.Column.Types import DataFrame.Internal.Control.Concurrent (-    capabilities,     chunksFor,     forkJoin,     parThreshold,     shouldParallelize,-    splitChunkRange,  ) import DataFrame.Internal.Data.HashTable (htInsert, newHashTable) import DataFrame.Internal.Data.PackedText (-    PackedSel,+    PackedOffsets (..),     PackedTextData (..),     offCount,     packedLength,@@ -60,11 +59,9 @@  ) import DataFrame.Internal.DataFrame (DataFrame (..), GroupedDataFrame (..)) import DataFrame.Internal.Grouping.Direct (-    DirectGrouping (..),     ascendingCodeGroups,     directGroupThreshold,     rangeOf,-    tryDirectGroupColumn,  ) import qualified DataFrame.Internal.Grouping.Direct as GD import DataFrame.Internal.Grouping.Partitioned (@@ -75,10 +72,7 @@ import System.IO.Unsafe (unsafePerformIO) import Type.Reflection (typeRep) -{- | O(k * n) group the dataframe by the given key columns, bucketing rows with an-open-addressing hash table that re-verifies keys on each hash hit. Groups are-numbered in first-appearance order; 'valueIndices'/'offsets' follow by counting sort.--}+-- | O(k * n) group the dataframe by the given key columns via a hash table. groupBy ::     [T.Text] ->     DataFrame ->@@ -103,23 +97,7 @@   where     !n = nRows df -{- | Low-cardinality direct-indexed grouping fast path-('DataFrame.Internal.Grouping.Direct'): fires for key lists-where every key is a single clean small-range @Int@ column or a canonical-dict-encoded text column, and the product of the key domains stays within-'directGroupThreshold' (the keys fuse into one mixed-radix code; a single key-degenerates to its own code). Returns 'Nothing' on any other key shape, falling-back to the hash path.--Narrow domains build @offsets@/@groupRepRows@ eagerly (histogram-sized work-only) and leave BOTH per-row outputs lazy: @valueIndices@ ('visFromCodes')-only materializes for consumers that gather (median/top-k, set ops,-interpreter slices), and @rowToGroup@ ('rtgFromCodes') only for the streaming-scatter aggregations — each aggregate pays for exactly one O(n) output pass,-not both. Wide domains run the two-level radix engine instead-('DataFrame.Internal.Grouping.Direct.directLayoutLazy'): @rowToGroup@ eager,-@valueIndices@ deferred (see 'fusedDirectGroup').--}+-- | Low-cardinality direct-indexed grouping fast path for multi-key lists. tryDirectGroup :: [T.Text] -> DataFrame -> Maybe GroupedDataFrame tryDirectGroup [] _ = Nothing tryDirectGroup names df = do@@ -131,22 +109,16 @@             ([name], [col]) -> tryDictGroup (nRows df) df [name] col             _ -> Nothing -{- | One key column of a fused multi-key direct grouping: a per-row component-code in @[0, fkDomain)@ (negative marks an invalid/corrupt code, which aborts-the direct path).--}+-- | One key column of a fused multi-key direct grouping. data FusedKey = FusedKey     { fkCode :: Int -> Int     , fkDomain :: !Int     } -{- | Classify a key column for the fused multi-key direct path: a clean non-null-unboxed @Int@ of small range, or a non-null canonical dict-encoded text column-of small dictionary. Anything else falls back to the hash group-by.--}+-- | Classify a key column for the fused multi-key direct path. fusedKey :: Column -> Maybe FusedKey-fusedKey (UnboxedColumn Nothing (v :: VU.Vector a))-    | Just Refl <- testEquality (typeRep @a) (typeRep @Int)+fusedKey col@(UnboxedColumn Nothing (_ :: VU.Vector a))+    | Just v <- asIntKeys col     , not (VU.null v) =         let (!mn, !mx) = rangeOf v             !range = mx - mn + 1@@ -168,24 +140,14 @@                 else Nothing fusedKey _ = Nothing -{- | Fuse the per-key codes into one mixed-radix code per row-(@((k1*d2)+k2)*d3+...@) and feed the direct counting-sort machinery. Group-order: ascending fused code (lexicographic in key order) — ascending value-order for @Int@ keys (mirroring the order the single-@Int@-key direct path-always had) and ascending dictionary code for dict-encoded text keys (the-dictionary's first-appearance order, a fixed property of the column). The-ascending order keeps @codeToGroup@ an identity map whenever the domain is-fully occupied, so the deferred @rowToGroup@ pass skips its per-row random-remap lookup; ranking dict groups by string hash instead (the historical-order) profiled ~0.6s slower per 1e8 rows at 1e6 groups.+-- | An integral key column as 'Int's; floats are refused.+asIntKeys :: Column -> Maybe (VU.Vector Int)+asIntKeys (UnboxedColumn Nothing (v :: VU.Vector a)) = case sFloating @a of+    STrue -> Nothing+    SFalse -> widenToInt v+asIntKeys _ = Nothing -On the narrow-domain engine, @valueIndices@ and @rowToGroup@ are passed to the-constructor as unevaluated applications of 'visFromCodes' / 'rtgFromCodes'-(constructor arguments are not forced even under @-XStrict@, and the fields-are lazy at their definition site), so each per-row output pass is deferred-until a consumer demands it. The wide-domain engine defers only-@valueIndices@ (see the branch comment below).--}+-- | Fuse per-key codes into one mixed-radix code per row and group directly. fusedDirectGroup ::     [T.Text] -> DataFrame -> [FusedKey] -> Maybe GroupedDataFrame fusedDirectGroup names df keys = do@@ -194,13 +156,7 @@         codeAt' = fusedCodeAt keys     if GD.useTwoLevel n domain         then do-            {- Wide domains (> ~1024 codes at parallel scale): the two-level-            radix engine — no pass random-writes a multi-megabyte table per-            worker, unlike the per-chunk direct histograms below (measured ~2x-            on the eager layout at 1e6 codes / 1e8 rows). It builds-            @rowToGroup@ eagerly (the streaming aggregations force it first-            thing anyway); only @valueIndices@ stays deferred, reconstructed-            from @rowToGroup@ by the same engine on demand. -}+            -- Wide domains use the radix engine; valueIndices stays deferred.             (rtg, offs, reps, nGroups) <-                 GD.directLayoutLazy codeAt' n domain ascendingCodeGroups             Just@@ -225,10 +181,7 @@                     reps                 ) -{- | Product of the per-key domains, 'Nothing' once it (or any factor) passes-'directGroupThreshold'. Factors are capped before multiplying, so the running-product never exceeds @threshold^2@ and cannot overflow.--}+-- | Product of the per-key domains, or 'Nothing' past the direct threshold. fusedDomain :: [Int] -> Maybe Int fusedDomain = go 1   where@@ -238,10 +191,7 @@         | acc * d > directGroupThreshold = Nothing         | otherwise = go (acc * d) ds -{- | Per-row fused mixed-radix code; @-1@ when any component code is invalid-(only possible for corrupt dict codes), making 'groupCodesMaybe' bail to the-hash path. Valid components compose to a code in @[0, product of domains)@.--}+-- | Per-row fused mixed-radix code; @-1@ when any component code is invalid. fusedCodeAt :: [FusedKey] -> (Int -> Int) fusedCodeAt [] = const (-1) fusedCodeAt (k0 : ks0) = go (fkCode k0) ks0@@ -259,10 +209,7 @@                 )                 ks -{- | Dictionary-encode a single text key to dense int codes, then derive-@valueIndices@/@offsets@ by counting sort. Profiled slower than the fused hash-group-by on every db-benchmark question, so it always falls back ('dictGroupEnabled').--}+-- | Dictionary-encode a single text key, then group by counting sort. tryDictGroup ::     Int -> DataFrame -> [T.Text] -> Column -> Maybe GroupedDataFrame tryDictGroup n df names col@@ -272,23 +219,15 @@         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.--}+-- | Master switch for the single-key dict-encode grouping path. dictGroupEnabled :: Bool dictGroupEnabled = False -{- | Cardinality ceiling for the single-key dict-encode probe: it bails to 'Nothing'-once the distinct count passes this. Only consulted when 'dictGroupEnabled' is 'True'.--}+-- | Cardinality ceiling for the single-key dict-encode probe. 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.--}+-- | The sequential grouping path: one open-addressing table over all rows. groupBySeq :: [T.Text] -> DataFrame -> GroupedDataFrame groupBySeq names df =     let !n = nRows df@@ -300,10 +239,7 @@         (vis, os) = indicesFromGroups rtg nGroups      in Grouped df names vis os rtg -{- | The parallel partitioned grouping path (see 'DataFrame.Internal.Grouping.Partitioned'):-forks one task per capability, producing output bit-for-bit identical to-'groupBySeq'. Pure via 'unsafePerformIO' (deterministic thread fan-out only).--}+-- | The parallel partitioned grouping path, identical in output to sequential. groupByPar :: [T.Text] -> DataFrame -> GroupedDataFrame groupByPar names df =     let !n = nRows df@@ -312,15 +248,9 @@         selectedCols = map (materializeMerged . (columns df V.!)) indicesToGroup         !eqRow = eqKeyRow df indicesToGroup         (vis, os) = unsafePerformIO $ do-            -- Parallel row-hash kernel, bit-identical to 'computeHashes' at the-            -- same dict-code setting (grouping always hashes canonical dict-            -- columns by code; see 'hashPacked').             hashes <- computeRowHashesWithIO True n selectedCols             parallelAssignGroups n hashes eqRow-     in -- rowToGroup is passed as an UNFORCED constructor argument (this module-        -- is -XStrict, so it must not be let-bound): gather-style aggregation-        -- over huge group counts never reads it, and the deferred pass writes-        -- values identical to the eager build.+     in -- rowToGroup stays an unforced argument (the module is -XStrict).         Grouped df names vis os (rtgFromVisOffs n vis os) {-# NOINLINE groupByPar #-} @@ -329,10 +259,7 @@ keyColIndices names df =     M.elems $ M.filterWithKey (\k _ -> k `elem` names) (columnIndices df) -{- | Assign every row to a dense group id in first-appearance order. Returns-@(rowToGroup, repHash, repRow)@ — the hash and representative row of each group.-'eqKeyRow' re-verifies the real key on each hash hit so colliding keys stay apart.--}+-- | Assign every row a dense group id in first-appearance order. assignGroups ::     DataFrame -> [Int] -> Int -> (VU.Vector Int, VU.Vector Int, VU.Vector Int) assignGroups df indicesToGroup n = runST $ do@@ -358,18 +285,12 @@     repRow <- VU.unsafeFreeze (VUM.slice 0 nGroups repRowM)     pure (frozen, repHash, repRow) -{- | Map each first-appearance group id to its canonical id: groups ordered by-ascending representative hash (tie-broken by representative row), making group-order a deterministic function of the key set so set ops commute. O(g), no sort.--}+-- | Map each first-appearance group id to its canonical id (by hash). canonicalRemap :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int canonicalRemap repHash _repRow =     runST (rankByHash (pure . VU.unsafeIndex repHash) (VU.length repHash)) -{- | Compute the FNV row-hash of the key columns into a fresh unboxed vector,-mixing 'nullSalt' for null slots so a missing value never collides with a-present one of the same bits.--}+-- | Compute the FNV row-hash of the key columns into an unboxed vector. computeHashes :: DataFrame -> [Int] -> Int -> ST s (VU.Vector Int) computeHashes df indicesToGroup n = do     mh <- VUM.replicate n fnvOffset@@ -419,20 +340,16 @@             error "computeHashes: MergedColumn is normalized before hashing"     VU.unsafeFreeze mh -{- | Build the row-key equality predicate over the selected key columns.-@eqKeyRow df idxs a b@ is 'True' iff rows @a@ and @b@ agree on all key columns-(validity first, a null equals only a null). Used to reject hash collisions.--}+-- | Build the row-key equality predicate over the selected key columns. 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+eqKeyRow df indicesToGroup = case map (colEqRow . (columns df V.!)) indicesToGroup of+    [p] -> p+    preds ->+        let go [] _ _ = True+            go (p : ps) a b = p a b && go ps a b+         in go preds -{- | Per-column row equality respecting nulls. Two rows are equal at a column-when both are null, or both are valid and their values compare equal.--}+-- | Per-column row equality respecting nulls. colEqRow :: Column -> (Int -> Int -> Bool) colEqRow c@(MergedColumn _ _) = colEqRow (materializeMerged c) colEqRow (UnboxedColumn bm v) =@@ -442,22 +359,38 @@     let eqV a b = V.unsafeIndex v a == V.unsafeIndex v b      in withNulls bm eqV colEqRow (PackedText bm p) =-    -- A canonical dictionary selection assigns equal strings the same code,-    -- so two rows are byte-equal iff their codes agree.+    -- A canonical dictionary selection gives equal strings equal codes.     let eqV = case ptSel p of             Just sel                 | ptCanonicalSel p ->                     \a b -> selAt sel a == selAt sel b-            _ -> \a b ->-                let (arrA, oA, lA) = packedSlice p a-                    (arrB, oB, lB) = packedSlice p b-                 in sliceEqBytes arrA oA lA arrB oB lB+            _ -> packedEqRow p      in withNulls bm eqV {-# INLINE colEqRow #-} -{- | Wrap a value-equality with null handling: equal iff both valid and the-values agree, or both null.--}+-- | Byte equality of two packed-text rows, with the layout matched once.+packedEqRow :: PackedTextData -> (Int -> Int -> Bool)+packedEqRow (PackedTextData arr offs sel _) = case (sel, offs) of+    (Nothing, Offs32 o) -> \a b -> a == b || eqBase (off32 o) a b+    (Nothing, Offs64 o) -> \a b -> a == b || eqBase (VU.unsafeIndex o) a b+    (Just s, Offs32 o) -> eqSel s (off32 o)+    (Just s, Offs64 o) -> eqSel s (VU.unsafeIndex o)+  where+    off32 o i = fromIntegral (VU.unsafeIndex o i)+    -- Base rows; equal ones share bytes, and -1 is the empty string.+    eqSel s off a b =+        let !ra = selAt s a+            !rb = selAt s b+         in ra == rb || eqBase off ra rb+    eqBase off ra rb =+        let span r = if r < 0 then (0, 0) else let o = off r in (o, off (r + 1) - o)+            (!oa, !la) = span ra+            (!ob, !lb) = span rb+         in sliceEqBytes arr oa la arr ob lb+    {-# INLINE eqSel #-}+    {-# INLINE eqBase #-}++-- | Wrap a value-equality with null handling (both null counts as equal). withNulls :: Maybe Bitmap -> (Int -> Int -> Bool) -> (Int -> Int -> Bool) withNulls Nothing eqV = eqV withNulls (Just bm) eqV = \a b ->@@ -467,10 +400,7 @@         _ -> 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.--}+-- | Derive @(valueIndices, offsets)@ from @rowToGroup@ by stable counting sort. indicesFromGroups :: VU.Vector Int -> Int -> (VU.Vector Int, VU.Vector Int) indicesFromGroups rtg nGroups = runST $ do     let !n = VU.length rtg@@ -512,37 +442,15 @@     frozenVis <- VU.unsafeFreeze vis     pure (frozenVis, offs) ----------------------------------------------------------------------------------- Deferred-placement direct grouping----------------------------------------------------------------------------------{- | Contiguous per-worker row ranges: one chunk per capability above the-parallel threshold, a single chunk otherwise.--}+-- | Contiguous per-worker row ranges, one chunk per capability when parallel. directRowChunks :: Int -> [(Int, Int)] directRowChunks = chunksFor parThreshold -{- | Like 'directRowChunks' but over a code/group domain (merge/seed passes),-which pays for a fan-out at a much lower width than the row passes do.--}+-- | Like 'directRowChunks' but over a code/group domain. directCodeSlices :: Int -> [(Int, Int)] directCodeSlices = chunksFor 4096 -{- | The eager phases of the direct counting-sort grouping — WITHOUT either-per-row output pass: per-chunk validated code histograms (parallel), per-code-totals and first occurrences (parallel over code slices), the caller-chosen-code->group mapping, the offsets prefix scan and per-group representative-rows. Returns-@(offsets, groupRepRows, counts, codeToGroup, chunkHists, nGroups)@ — the last-three feed the deferred @valueIndices@ placement ('visFromCodes') and-@rowToGroup@ ('rtgFromCodes') thunks, so a consumer pays only for the per-row-output it actually demands. 'Nothing' when any row's code falls outside-@[0, card)@ (fall back to hashing).--Pure w.r.t. its immutable inputs: the fork fan-out is a fixed function of the-row count and capability count, and every merge runs in fixed chunk order, so-the result is deterministic and the 'unsafePerformIO' is safe.--}+-- | The eager phases of the direct counting-sort grouping. directLayoutLazy ::     (Int -> Int) ->     Int ->@@ -559,9 +467,7 @@ directLayoutLazy codeAt n card mkGroups     | n <= 0 || card <= 0 = Nothing     | n < packedRowLimit = unsafePerformIO $ do-        -- Packed variant: count and first-occurrence row share one word per-        -- code, keeping phase 1 at a single accumulator array per chunk-        -- (measured ~0.15s/1e8 rows cheaper than a second firstOcc array).+        -- Packed variant: count and first row share one word per code.         let chunks = directRowChunks n         parts <- forkJoin [histFirstChunkPacked codeAt card lo hi | (lo, hi) <- chunks]         if not (all snd parts)@@ -569,8 +475,7 @@             else finishLayout card mkGroups (map fst parts) $ \histsM totalsM firstAllM lo hi ->                 sumFirstSlicePacked histsM totalsM firstAllM lo hi     | otherwise = unsafePerformIO $ do-        -- Fallback for gigantic frames where a row index does not fit the-        -- packed word: separate count and firstOcc arrays, same results.+        -- Fallback for huge frames: separate count and firstOcc arrays.         let chunks = directRowChunks n         parts <- forkJoin [histFirstChunk codeAt card lo hi | (lo, hi) <- chunks]         if not (all (\(_, _, ok) -> ok) parts)@@ -585,11 +490,7 @@                     ) {-# NOINLINE directLayoutLazy #-} -{- | Shared tail of 'directLayoutLazy': run the totals/first-occurrence merge-(which also normalizes each chunk histogram to plain counts, see-'sumFirstSlicePacked'), derive the group mapping, offsets and representative-rows, and freeze the retained chunk histograms.--}+-- | Shared tail of 'directLayoutLazy': merge, normalise and build the layout. finishLayout ::     Int ->     (VU.Vector Int -> (VU.Vector Int, Int)) ->@@ -633,10 +534,7 @@     hists <- mapM VU.unsafeFreeze histsM     pure (Just (offs, reps, counts, codeToGroup, hists, nGroups)) -{- | Rows must satisfy @row + 1 < 2^31@ for the packed count/first-row encoding-(count in the high bits, first row + 1 in the low 31). Above it (a >2e9-row-frame, >17GB per Int column) the unpacked variant runs instead.--}+-- | Row limit for the packed count/first-row encoding (@row + 1 < 2^31@). packedRowLimit :: Int packedRowLimit = 0x7FFFFFFF @@ -644,10 +542,7 @@ packedCountOne :: Int packedCountOne = 0x80000000 -{- | Whether @codeToGroup@ maps every code to itself (fully occupied ascending-domain — e.g. a dense Int key covering its whole range). The rowToGroup pass-then skips the random remap lookup entirely.--}+-- | Whether @codeToGroup@ maps every code to itself. isIdentityMap :: VU.Vector Int -> Bool isIdentityMap m = go 0   where@@ -657,10 +552,7 @@         | VU.unsafeIndex m i /= i = False         | otherwise = go (i + 1) -{- | Histogram one row chunk with the packed encoding: slot @c@ holds-@count(c) * 2^31 + (firstRow(c) + 1)@ (zero = never seen). One accumulator-array per chunk. Reports 'False' as soon as any code escapes @[0, card)@.--}+-- | Histogram one row chunk with the packed count and first-row encoding. histFirstChunkPacked ::     (Int -> Int) -> Int -> Int -> Int -> IO (VUM.IOVector Int, Bool) histFirstChunkPacked codeAt card lo hi = do@@ -681,11 +573,7 @@     ok <- go lo     pure (acc, ok) -{- | Per-code totals and overall first occurrences from the PACKED chunk-histograms, rewriting each histogram slot to its plain count in place (the-placement thunk then sees ordinary counts). Chunks are ordered by row range, so-the first chunk with a nonzero slot holds the code's globally first row.--}+-- | Per-code totals and first occurrences from the packed chunk histograms. sumFirstSlicePacked ::     [VUM.IOVector Int] ->     VUM.IOVector Int ->@@ -711,11 +599,7 @@             VUM.unsafeWrite firstAll c fo             go (c + 1) -{- | Histogram one row chunk into a private @card@-slot count plus the chunk's-first occurrence of each code, reporting 'False' as soon as any code escapes-@[0, card)@ (the counts are then abandoned). Fallback for frames beyond-'packedRowLimit'.--}+-- | Histogram one row chunk into a private count plus first occurrences. histFirstChunk ::     (Int -> Int) ->     Int ->@@ -739,10 +623,7 @@     ok <- go lo     pure (acc, firstOcc, ok) -{- | Per-code totals over one code slice, plus the overall first occurrence of-each code: the chunks are ordered by row range, so the first chunk with a-nonzero count for a code holds its globally first row.--}+-- | Per-code totals over one code slice, plus each code's first occurrence. sumFirstSlice ::     [VUM.IOVector Int] ->     [VUM.IOVector Int] ->@@ -770,9 +651,7 @@             VUM.unsafeWrite firstAll c fo             go (c + 1) -{- | Exclusive prefix scan of per-group counts (gathered through @codeToGroup@)-into the offsets array of length @nGroups + 1@.--}+-- | Exclusive prefix scan of per-group counts into the offsets array. scanGroupOffsets :: VU.Vector Int -> VU.Vector Int -> Int -> IO (VU.Vector Int) scanGroupOffsets counts codeToGroup nGroups = do     let !card = VU.length counts@@ -797,10 +676,7 @@     scan 0 0     VU.unsafeFreeze offsM -{- | @reps[codeToGroup c] = firstAll c@ for every occupied code: each group is-exactly one occupied code, so this is a disjoint parallel write and equals-@vis[offs[g]]@ (the group's first row in original order).--}+-- | Scatter representative rows for every occupied code, in parallel. scatterRepsSlice ::     VU.Vector Int ->     VU.Vector Int ->@@ -821,16 +697,7 @@                 (VU.unsafeIndex firstAll c)             go (c + 1) -{- | Deferred @rowToGroup@: one parallel per-row pass mapping each row's code-through @codeToGroup@ (skipping the lookup entirely when the map is the-identity, i.e. a fully occupied ascending domain). Only the streaming-aggregation paths force this; a purely gather-driven consumer (median, top-k)-never pays for it.--Pure w.r.t. its immutable inputs and deterministic (fixed chunking), so the-'unsafePerformIO' behind a lazy field is safe: whenever and however many times-the thunk is forced it yields the same vector.--}+-- | Deferred @rowToGroup@: one parallel per-row pass through @codeToGroup@. rtgFromCodes :: (Int -> Int) -> VU.Vector Int -> Int -> VU.Vector Int rtgFromCodes codeAt codeToGroup n = unsafePerformIO $ do     rtgM <- VUM.new n@@ -868,21 +735,7 @@             VUM.unsafeWrite rtgM i (codeAt i)             go (i + 1) -{- | Deferred stable placement: build the @valueIndices@ permutation from the-per-row codes and the RETAINED phase-1 chunk histograms — the same-seed-cursors-then-place structure (and cost) the eager path used, minus the-@rowToGroup@ writes. Each chunk's code-indexed cursor starts at the group-offset plus everything earlier chunks (in row order) place there, so rows keep-original order within each group: the result is the unique group-major,-original-row-order permutation, bit-identical to the eager placement at any-chunk count.--Pure w.r.t. its immutable inputs and deterministic (fixed chunking, fixed merge-order), so the 'unsafePerformIO' behind a lazy field is safe: whenever and-however many times the thunk is forced it yields the same vector. The thunk-retains the chunk histograms (capabilities x card words) until forced or the-grouping is dropped.--}+-- | Deferred stable placement building @valueIndices@ from per-row codes. visFromCodes ::     (Int -> Int) ->     VU.Vector Int ->@@ -894,8 +747,7 @@     VU.Vector Int visFromCodes codeAt counts codeToGroup offs hists n card = unsafePerformIO $ do     let chunks = directRowChunks n-    -- Private mutable copies of the retained histograms, rewritten in place-    -- into the per-chunk write cursors.+    -- Private mutable copies of the histograms, rewritten into write cursors.     cursors <- mapM VU.thaw hists     _ <-         forkJoin@@ -911,10 +763,7 @@     VU.unsafeFreeze vis {-# NOINLINE visFromCodes #-} -{- | Rewrite each chunk's histogram copy in place into its disjoint write-cursor: chunk w's run for code c starts at the offset of c's group plus what-earlier chunks (in row order) place there.--}+-- | Rewrite each chunk's histogram copy in place into its write cursor. seedCursorSlice ::     VU.Vector Int ->     VU.Vector Int ->@@ -938,9 +787,7 @@             loop cursors (VU.unsafeIndex offs g)             go (c + 1) -{- | Stable placement over one row chunk via the chunk's advancing-code-indexed cursors.--}+-- | Stable placement over one row chunk via its advancing cursors. placeVisChunk ::     (Int -> Int) -> VUM.IOVector Int -> VUM.IOVector Int -> Int -> Int -> IO () placeVisChunk codeAt cursor vis lo hi = go lo@@ -954,9 +801,7 @@             VUM.unsafeWrite cursor c (pos + 1)             go (i + 1) -{- | Fold a value-mix over an unboxed column into the running hash vector,-respecting the null bitmap: a null slot mixes a fixed 'nullSalt' sentinel.--}+-- | Fold a value-mix over an unboxed column into the running hash vector. hashUnboxed ::     (VU.Unbox a) =>     VUM.MVector s Int ->@@ -984,14 +829,7 @@             v {-# INLINE hashUnboxed #-} -{- | Hash a packed-text column over its raw UTF-8 byte slices (no per-row-'Data.Text.Text'), mixing 'nullSalt' for null rows. Shares 'mixBytes' with-'mixText' so packed and boxed Text columns hash identically. A canonical-dict-encoded column (equal strings share a code) instead mixes its 'Int' code-with one 'mixInt' per row; 'DataFrame.Internal.RowHash.packedRange' applies the-same rule under the grouping setting so 'groupBySeq' and 'groupByPar' bucket-identically.--}+-- | Hash a packed-text column over its raw UTF-8 byte slices. hashPacked ::     VUM.MVector s Int -> Maybe Bitmap -> PackedTextData -> ST s () hashPacked mh bm p = case ptSel p of@@ -1027,9 +865,7 @@ nRows :: DataFrame -> Int nRows = fst . dataframeDimensions -{- | Build the rowToGroup lookup vector from valueIndices and offsets.-rowToGroup[i] = k means row i belongs to group k.--}+-- | Build the rowToGroup lookup vector from valueIndices and offsets. buildRowToGroup :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int buildRowToGroup n vis os = runST $ do     rtg <- VUM.new n

@@ -4,28 +4,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | Low-cardinality direct-indexed grouping fast path: when every row's key-reduces to a dense @Int@ code in a small domain, the code itself indexes a dense-accumulator (no hashing/probing). All O(n) passes run chunked across-capabilities — per-chunk histograms feed prefix-summed disjoint write cursors —-so the stable within-group row order of the sequential counting sort is-reproduced exactly.--Two engines cover the code-domain spectrum with bit-identical results:--* narrow domains (@card <= 'twoLevelCardThreshold'@) index per-worker-  histogram\/cursor tables directly — they stay cache-resident;--* wide domains use a two-level radix split (top code bits pick one of ~@2^10@-  buckets, cursors cache-resident) so no pass ever random-writes a-  multi-megabyte table per worker.--'directLayoutLazy' is the aggregation entry point: it builds @rowToGroup@,-@offsets@ and the per-group representative rows eagerly but skips the O(n)-stable placement entirely; 'visFromRowToGroup' reconstructs @valueIndices@-on demand (its value is the unique stable counting-sort permutation, so WHEN it-runs is unobservable).--}+-- | Low-cardinality direct-indexed grouping fast path for dense Int key codes. module DataFrame.Internal.Grouping.Direct (     directGroupThreshold,     tryDirectGroupColumn,@@ -59,17 +38,11 @@     splitChunkRange,  ) -{- | Largest key code DOMAIN (single-key value range, or the product of per-key-domains for a fused multi-key code) the direct grouping path accepts. A @2^20@-slot-histogram is 8MB; the low-cardinality questions sit far below it (id4 range 100,-id6 range 1e5). Wider domains fall back to the hash group-by.--}+-- | Largest key code domain the direct grouping path accepts. 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.--}+-- | The grouping layout the hash path also produces. data DirectGrouping = DirectGrouping     { dgRowToGroup :: !(VU.Vector Int)     , dgValueIndices :: !(VU.Vector Int)@@ -77,10 +50,7 @@     , dgNGroups :: !Int     } -{- | Take the direct path if the (single) key column is a clean non-null unboxed-@Int@ column with a small value range. Returns 'Nothing' to fall back to the-hash group-by on anything else (boxed/text keys, nullable, wide ranges, empty).--}+-- | Take the direct path for a clean small-range Int key, else 'Nothing'. tryDirectGroupColumn :: Column -> Maybe DirectGrouping tryDirectGroupColumn (UnboxedColumn Nothing (v :: VU.Vector a))     | Just Refl <- testEquality (typeRep @a) (typeRep @Int)@@ -125,38 +95,17 @@ shouldPar :: Int -> Bool shouldPar = shouldParallelize parThreshold -{- | Contiguous per-worker row ranges: one chunk per capability above the-parallel threshold, a single chunk otherwise (the sequential fallback runs the-same code on the calling thread).--}+-- | Contiguous per-worker row ranges, one chunk per capability when parallel. rowChunks :: Int -> [(Int, Int)] rowChunks = chunksFor parThreshold -{- | Like 'rowChunks' but over the code domain (parallel merge/seed passes),-which pays for a fan-out at a much lower width than the row passes do.--}+-- | Like 'rowChunks' but over the code domain. codeSlices :: Int -> [(Int, Int)] codeSlices = chunksFor 4096 -{- | Run each action on its own thread and collect the results in order;-rethrow the first failure. A single action runs on the calling thread.--}------------------------------------------------------------------------------------ Full grouping (eager valueIndices): compatibility entry point----------------------------------------------------------------------------------{- | Build the grouping by counting sort on a per-row code in @[0, card)@.-Returns 'Nothing' when any row's code falls outside @[0, card)@ (fall back to-hashing).--@mkGroups counts@ must return a dense group id for every code with a nonzero-count (other slots are never read; occupied codes must get distinct ids) and-the group count; it decides group order.+-- un each action on its own thread and collect results in order. -Equivalent to 'directLayoutLazy' plus a forced 'visFromRowToGroup'; kept for-callers that want the whole layout eagerly.--}+-- | Build the grouping by counting sort on a per-row code in @[0, card)@. groupCodesMaybe ::     (Int -> Int) ->     Int ->@@ -168,19 +117,7 @@     let !vis = visFromRowToGroup n nGroups offs rtg     Just (DirectGrouping rtg vis offs nGroups) ----------------------------------------------------------------------------------- Eager layout without placement: rowToGroup + offsets + group rep rows----------------------------------------------------------------------------------{- | The layout every aggregation needs, WITHOUT the O(n) stable placement:-@(rowToGroup, offsets, groupRepRows, nGroups)@, all four computed eagerly.-@groupRepRows[g]@ is the first original row of group @g@ (what-@valueIndices[offsets[g]]@ evaluates to). Pair with 'visFromRowToGroup' for a-deferred @valueIndices@. Returns 'Nothing' when any row's code falls outside-@[0, card)@.--@mkGroups@ contract as in 'groupCodesMaybe'.--}+-- | The layout every aggregation needs, without the O(n) stable placement. directLayoutLazy ::     (Int -> Int) ->     Int ->@@ -193,9 +130,7 @@     | otherwise = unsafePerformIO (layoutNarrow codeAt n card mkGroups) {-# NOINLINE directLayoutLazy #-} -{- | Narrow domains (and the sequential small-@n@ fallback): per-chunk direct-histograms with first-occurrence tracking, merged over code slices.--}+-- | Narrow domains: per-chunk histograms with first-occurrence tracking. layoutNarrow ::     (Int -> Int) ->     Int ->@@ -218,10 +153,7 @@             firstRow <- VU.unsafeFreeze firstRowM             finishLayout codeAt n card mkGroups counts firstRow -{- | Histogram one row chunk into a private @card@-slot count, recording the-chunk's first row of each code, and reporting invalid codes (third component-'False'; the arrays are then abandoned).--}+-- | Histogram one row chunk into a private count, recording first rows. histFirstChunk ::     (Int -> Int) ->     Int ->@@ -245,11 +177,7 @@     ok <- go lo     pure (acc, firstV, ok) -{- | @totals[c] = Σ_w counts_w[c]@ and @firstRow[c]@ = the first chunk's first-occurrence (chunks are in row order, so that IS the global first row of @c@),-over one code slice. Every @totals@ slot is written; @firstRow[c]@ only where-the count is nonzero (never read otherwise).--}+-- | Merge per-chunk counts and first rows over one slice of the code domain. mergeSlice ::     [(VUM.IOVector Int, VUM.IOVector Int)] ->     VUM.IOVector Int ->@@ -272,11 +200,7 @@             VUM.unsafeWrite totals c s             go (c + 1) -{- | Wide domains: two-level radix. Rows are bucketed by the top code bits-(cache-resident cursors) as packed @(code, row)@ words; each bucket's slice —-in original row order — yields its exact per-code counts and first rows from an-L1-resident table. No placement pass runs here.--}+-- | Wide domains: a two-level radix over packed @(code, row)@ words. layoutWide ::     (Int -> Int) ->     Int ->@@ -299,10 +223,7 @@             firstRow <- VU.unsafeFreeze firstRowM             finishLayout codeAt n card mkGroups counts firstRow -{- | Shared tail of 'directLayoutLazy': group mapping, offsets, representative-rows gathered through the code-to-group table, and the parallel @rowToGroup@-pass (one sequential read of the codes, one table lookup each).--}+-- | Shared tail of 'directLayoutLazy': mapping, offsets and representatives. finishLayout ::     (Int -> Int) ->     Int ->@@ -327,8 +248,7 @@     repLoop 0     reps <- VU.unsafeFreeze repsM     rtgM <- VUM.unsafeNew n-    -- A fully occupied ascending domain maps every code to itself; skipping-    -- the per-row random table lookup then leaves one sequential read+write.+    -- An identity code map lets us skip the per-row table lookup.     let identity = isIdentityMap codeToGroup     _ <-         forkJoin@@ -351,9 +271,7 @@         | VU.unsafeIndex m i /= i = False         | otherwise = go (i + 1) -{- | Exclusive prefix scan of per-group counts (gathered through @codeToGroup@)-into the offsets array of length @nGroups + 1@.--}+-- | Exclusive prefix scan of per-group counts into the offsets array. scanOffsets :: VU.Vector Int -> VU.Vector Int -> Int -> IO (VU.Vector Int) scanOffsets counts codeToGroup nGroups = do     let !card = VU.length counts@@ -400,19 +318,7 @@             VUM.unsafeWrite rtgM i (codeAt i)             go (i + 1) ----------------------------------------------------------------------------------- Deferred stable placement: valueIndices from rowToGroup----------------------------------------------------------------------------------{- | The unique stable counting-sort permutation of @rowToGroup@: rows sorted by-group id, original order within each group — exactly what the eager engines'-placement pass produces. Pure (integer bookkeeping only, deterministic), so it-can sit under a lazy 'DataFrame.Internal.DataFrame.valueIndices' field.--Preconditions (all guaranteed by the grouping paths): @rtg@ has length @n@ with-every value in @[0, nGroups)@, and @offs@ is the group-count prefix array of-length @nGroups + 1@ with @offs[nGroups] == n@.--}+-- | The unique stable counting-sort permutation of @rowToGroup@. visFromRowToGroup ::     Int -> Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int visFromRowToGroup n nGroups offs rtg@@ -421,9 +327,7 @@     | otherwise = unsafePerformIO (visNarrow n nGroups offs rtg) {-# NOINLINE visFromRowToGroup #-} -{- | Narrow group domains: per-chunk group histograms prefix-summed (chunk-order) into disjoint cursors seeded from @offs@, then parallel stable placement.--}+-- | Narrow group domains: per-chunk histograms, then parallel stable placement. visNarrow :: Int -> Int -> VU.Vector Int -> VU.Vector Int -> IO (VU.Vector Int) visNarrow n nGroups offs rtg = do     let chunks = rowChunks n@@ -441,8 +345,7 @@                 go lo             | (lo, hi) <- chunks             ]-    -- Rewrite each chunk histogram into its write cursor:-    -- cursor_w[g] = offs[g] + Σ_{w'<w} hist_w'[g].+    -- Each chunk histogram becomes its write cursor: offs[g] + earlier chunks.     _ <-         forkJoin             [ let seed !g@@ -474,10 +377,7 @@             ]     VU.unsafeFreeze visM -{- | Wide group domains: two-level radix. Group ids ascend with buckets, so the-bucket-sorted layout written at @offs@-seeded cursors IS @valueIndices@ — each-bucket writes one contiguous region.--}+-- | Wide group domains: two-level radix placement into @valueIndices@. visWide :: Int -> Int -> VU.Vector Int -> VU.Vector Int -> IO (VU.Vector Int) visWide n nGroups offs rtg = do     mPacked <- packBucketed False (VU.unsafeIndex rtg) n nGroups@@ -492,9 +392,7 @@                 (placeBucketOffsPar shift nGroups bucketStart offs packed visM)             VU.unsafeFreeze visM -{- | Stable placement of one bucket's packed slice at cursors seeded straight-from the group offsets (the bucket's groups own a contiguous @vis@ region).--}+-- | Stable placement of one bucket's packed slice at offset-seeded cursors. placeBucketOffs ::     Int ->     Int ->@@ -527,10 +425,7 @@                 place (pos + 1)     place s -{- | 'placeBucketOffs' for an oversized bucket: per-sub-chunk histograms-prefix-summed (sub-chunks in row order) onto the offset-seeded cursors keep the-placement identical to the serial walk.--}+-- | 'placeBucketOffs' for an oversized bucket, split into sub-chunks. placeBucketOffsPar ::     Int ->     Int ->@@ -577,22 +472,11 @@             ]     pure () ----------------------------------------------------------------------------------- Two-level radix plumbing----------------------------------------------------------------------------------{- | Above this code-domain size the parallel passes switch to the two-level-radix engine; below it, per-worker direct tables stay cache-resident and are-faster (no bucket store).--}+-- | Code-domain size above which parallel passes use two-level radix. twoLevelCardThreshold :: Int twoLevelCardThreshold = 1024 -{- | The two-level engine packs @(code, row)@ into one machine word: row in the-low 'packShift' bits, code above them. Codes are capped at 'directGroupThreshold'-(@2^20@) by every caller, so the packed value stays well within 63 bits; the-guards in 'useTwoLevel' keep the narrow engine for anything larger.--}+-- | The two-level engine packs @(code, row)@ into one machine word. packShift :: Int packShift = 40 @@ -617,21 +501,13 @@     | x <= 1 = 0     | otherwise = 64 - countLeadingZeros (x - 1) -{- | The bucket store: rows partitioned by the top code bits, each bucket-holding packed @(code, row)@ words in original row order.--}+-- | The bucket store: rows partitioned by the top code bits. newtype PackedBuckets = PackedBuckets (VUM.IOVector Int)  packedVec :: PackedBuckets -> VUM.IOVector Int packedVec (PackedBuckets v) = v -{- | Partition rows into ~@2^'bucketBits'@ buckets by the top bits of their-code, as packed @(code, row)@ words: per-chunk bucket histograms (validating-every code when asked), prefix-summed in (bucket, chunk) order into disjoint-cursors, then a parallel scatter. Chunks are processed in row order, so every-bucket keeps its rows in ascending original row order. Returns 'Nothing' iff-validation was requested and some code fell outside @[0, card)@.--}+-- | Partition rows into buckets by top code bits as packed words. packBucketed ::     Bool ->     (Int -> Int) ->@@ -671,10 +547,7 @@                     ]             pure (Just (shift, bucketStart, PackedBuckets packed)) -{- | Histogram one row chunk by bucket (top code bits) into a private-@nBuckets@-slot count; with @validate@, report 'False' as soon as any code-escapes @[0, card)@ (the counts are then abandoned).--}+-- | Histogram one row chunk by bucket (top code bits) into a private count. bucketHist ::     Bool ->     (Int -> Int) ->@@ -701,7 +574,7 @@     ok <- go lo     pure (acc, ok) --- | Scatter one row chunk's packed @(code, row)@ words through its bucket cursor.+-- | Scatter one row chunk's packed words through its bucket cursor. scatterPacked ::     (Int -> Int) ->     Int ->@@ -722,12 +595,7 @@             VUM.unsafeWrite cursor b (pos + 1)             go (i + 1) -{- | Drive one action per bucket: buckets far above the fair per-worker share-run first through @big@ (internally parallel, one at a time), the rest are-pulled off a shared counter by one worker per capability. Every bucket —-including empty ones — is visited exactly once, so per-bucket passes may rely-on covering their whole output slice.--}+-- | Drive one action per bucket, with oversized buckets run first. overBuckets :: VU.Vector Int -> Int -> (Int -> IO ()) -> (Int -> IO ()) -> IO () overBuckets bucketStart n small big = do     let !nBuckets = VU.length bucketStart - 1@@ -737,11 +605,7 @@     pooledIndices capabilities nBuckets $ \b ->         when (size b < bigCut) (small b) -{- | One bucket's exact per-code counts and first rows from its (row-ordered)-packed slice, via an L1-resident table spanning only the bucket's code range.-Writes the bucket's whole slice of @counts@ (zeros included); @firstRow@ only-where the count is nonzero.--}+-- | One bucket's exact per-code counts and first rows from its packed slice. countFirstBucket ::     Int ->     Int ->@@ -840,10 +704,7 @@                 go (pos + 1)     go lo -{- | The ascending-code group order: walk the counts in code order, assigning a-dense group id to each non-empty code (empty codes get no id and no output-group). The single-Int-key path keeps its groups in ascending value order.--}+-- | The ascending-code group order: dense ids for non-empty codes. ascendingCodeGroups :: VU.Vector Int -> (VU.Vector Int, Int) ascendingCodeGroups counts = runST $ do     let !card = VU.length counts

@@ -2,19 +2,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE Strict #-} -{- | Partitioned group-by: rows are counting-sorted into partitions by the top-hash bits, then one task per capability groups its partitions independently,-after which the group numbering is canonicalized to first-appearance order.-Output is bit-for-bit identical to the sequential-'DataFrame.Internal.Grouping.groupBy'.--The name is the mechanism, not the threading: this is a genuinely different-algorithm from the sequential single-hash-table path, not that path with a-fork\/join wrapped around it. (Its sibling-"DataFrame.Internal.Grouping.Direct" is also internally parallel.) Whether to-take this path is 'DataFrame.Internal.Grouping.groupBy''s decision, not this-module's.--}+-- | Partitioned group-by: counting-sort rows by hash bits, group per partition. module DataFrame.Internal.Grouping.Partitioned (     parallelAssignGroups,     rtgFromVisOffs,@@ -22,8 +10,9 @@ ) where  import Control.Concurrent (getNumCapabilities)-import Control.Monad (forM_, when)-import Data.Bits (countLeadingZeros, unsafeShiftR)+import Control.Monad (when)+import Data.Bits (countLeadingZeros, unsafeShiftL, unsafeShiftR)+import qualified Data.List as L import qualified Data.Vector as V import qualified Data.Vector.Mutable as VM import qualified Data.Vector.Unboxed as VU@@ -43,10 +32,7 @@  ) import System.IO.Unsafe (unsafePerformIO) -{- | Sign-preserving unsigned remap: ascending 'Word64' order of @key h@ equals-ascending signed-'Int' order of @h@, so partitioning and sorting on it reproduce-the sequential @compare \`on\` repHash@ ordering exactly.--}+-- | Sign-preserving unsigned remap of a hash, ordering like signed 'Int'. key :: Int -> Word64 key h = fromIntegral h + 0x8000000000000000 {-# INLINE key #-}@@ -56,15 +42,7 @@ partIx shift h = fromIntegral (key h `unsafeShiftR` shift) {-# INLINE partIx #-} -{- | Number of partitions: a power of two, at least @4 * caps@ (P >> cores for-skew tolerance), floored at 256 — and grown with the row count (up to-'maxPartitions') so a partition's worst-case hash table (every row a distinct-group: @nextPow2Above (2 * rows/p)@ slots x 3 arrays) stays cache-resident-instead of thrashing DRAM. Partitioning is by the top hash bits and canonical-ranking is ascending unsigned key both across and within partitions, so the-grouping output is bit-for-bit identical at ANY partition count; only the-constant matters for speed.--}+-- | Number of partitions: a power of two scaled by capabilities and rows. numPartitionsFor :: Int -> Int -> Int numPartitionsFor caps n = go 1   where@@ -87,12 +65,7 @@ intLog2 x = 63 - countLeadingZeros x {-# INLINE intLog2 #-} -{- | Parallel group assignment. @parallelAssignGroups n hashes eqRow@ returns-@(valueIndices, offsets)@ in canonical group order. @eqRow a b@ must report-whether rows @a@ and @b@ share all key columns (null-aware). @rowToGroup@ is-NOT built here any more: gather-style aggregation over huge group counts never-reads it, so callers derive it on demand with 'rtgFromVisOffs'.--}+-- | Parallel group assignment, returning @(valueIndices, offsets)@ canonically. parallelAssignGroups ::     Int ->     VU.Vector Int ->@@ -102,8 +75,11 @@     caps <- getNumCapabilities     let !p = numPartitionsFor caps n         !shift = 64 - intLog2 p-    (partStart, sortedRows, sortedHash) <- partitionRows n hashes p shift+    link <- VUM.new (max 1 n)+    (partStart, sortedRows, sortedHash, sortedLink) <-+        partitionRows n hashes eqRow link p shift     localGid <- VUM.new (max 1 n)+    gidByRow <- VUM.new (max 1 n)     canonBoxes <- VM.replicate p (VU.empty :: VU.Vector Int)     nLocalGroups <- VUM.replicate p (0 :: Int)     runPartitions@@ -113,39 +89,29 @@         sortedRows         sortedHash         eqRow+        sortedLink+        gidByRow         localGid         canonBoxes         nLocalGroups     (globalBase, canonOf, nGroups) <- canonicalize p canonBoxes nLocalGroups     assemble n p partStart sortedRows localGid globalBase canonOf nGroups ----------------------------------------------------------------------------------- Phase 1: counting sort by partition----------------------------------------------------------------------------------{- | Bucket every row index into its partition by a counting sort. Returns the-exclusive prefix-sum @partStart@ (length @p+1@, @partStart[p] == n@), the row-indices laid out partition-by-partition in @sortedRows@, and each sorted-position's hash in @sortedHash@ (same layout) so the grouping loop reads its-hashes sequentially instead of a random @hashes[row]@ per row.--Runs chunked across capabilities: per-chunk partition histograms are prefix-summed (in chunk order) into disjoint per-chunk write cursors, so the scatter-threads never contend and each partition keeps its rows in ascending original-row order — bit-for-bit the sequential counting sort's layout.--}+-- | Bucket every row index into its partition by a counting sort. partitionRows ::     Int ->     VU.Vector Int ->+    (Int -> Int -> Bool) ->+    VUM.IOVector Int ->     Int ->     Int ->-    IO (VU.Vector Int, VU.Vector Int, VU.Vector Int)-partitionRows n hashes p shift = do+    IO (VU.Vector Int, VU.Vector Int, VU.Vector Int, VU.Vector Int)+partitionRows n hashes eqRow link p shift = do     caps <- getNumCapabilities     let chunks = rowChunks caps n-    cursors <- forkJoin [histChunk hashes p shift lo hi | (lo, hi) <- chunks]-    -- Exclusive prefix over partitions (outer) and chunks (inner): partStart-    -- from the totals, and each chunk's histogram rewritten into its cursor.+    cursors <-+        forkJoin [histChunk hashes eqRow link p shift lo hi | (lo, hi) <- chunks]+    -- Exclusive prefix over partitions and chunks; histograms become cursors.     partStartM <- VUM.new (p + 1)     let seed !pp !acc             | pp >= p = VUM.unsafeWrite partStartM p acc@@ -161,14 +127,16 @@     seed 0 0     sortedM <- VUM.new (max 1 n)     sortedHashM <- VUM.new (max 1 n)+    sortedLinkM <- VUM.new (max 1 n)     forkJoin_-        [ scatterChunk hashes shift cur sortedM sortedHashM lo hi+        [ scatterChunk hashes link shift cur sortedM sortedHashM sortedLinkM lo hi         | ((lo, hi), cur) <- zip chunks cursors         ]     partStart <- VU.unsafeFreeze partStartM     sortedRows <- VU.unsafeFreeze sortedM     sortedHash <- VU.unsafeFreeze sortedHashM-    pure (partStart, sortedRows, sortedHash)+    sortedLink <- VU.unsafeFreeze sortedLinkM+    pure (partStart, sortedRows, sortedHash, sortedLink)  -- | Contiguous near-equal row chunks, one per capability; empties dropped. rowChunks :: Int -> Int -> [(Int, Int)]@@ -182,32 +150,74 @@   where     !per = (n + max 1 caps - 1) `div` max 1 caps --- | Per-partition counts of one row chunk.-histChunk :: VU.Vector Int -> Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)-histChunk hashes p shift lo hi = do+-- | Per-partition counts of one row chunk, linking rows to recent equal keys.+histChunk ::+    VU.Vector Int ->+    (Int -> Int -> Bool) ->+    VUM.IOVector Int ->+    Int ->+    Int ->+    Int ->+    Int ->+    IO (VUM.IOVector Int)+histChunk hashes eqRow link p shift lo hi = do     acc <- VUM.replicate p (0 :: Int)-    let go !i+    -- Two ways per set, so two hot keys sharing a set do not evict each other.+    cacheHash <- VUM.new (2 * linkCacheSets)+    cacheRow <- VUM.replicate (2 * linkCacheSets) (-1 :: Int)+    let hit !w !h !i = do+            r <- VUM.unsafeRead cacheRow w+            h' <- VUM.unsafeRead cacheHash w+            pure (if r >= 0 && h' == h && eqRow r i then r else -1)+        go !i             | i >= hi = pure acc             | otherwise = do-                let !pp = partIx shift (VU.unsafeIndex hashes i)+                let !h = VU.unsafeIndex hashes i+                    !pp = partIx shift h+                    !w0 = 2 * cacheSet h+                r0 <- hit w0 h i+                r <- if r0 >= 0 then pure r0 else hit (w0 + 1) h i+                VUM.unsafeWrite link i r+                when (r < 0) $ do+                    -- Way 0 holds the newest key; the older one moves to way 1.+                    VUM.unsafeRead cacheHash w0 >>= VUM.unsafeWrite cacheHash (w0 + 1)+                    VUM.unsafeRead cacheRow w0 >>= VUM.unsafeWrite cacheRow (w0 + 1)+                    VUM.unsafeWrite cacheHash w0 h+                    VUM.unsafeWrite cacheRow w0 i                 c <- VUM.unsafeRead acc pp                 VUM.unsafeWrite acc pp (c + 1)                 go (i + 1)     go lo -{- | Scatter one row chunk into @sortedM@/@sortedHashM@ through the chunk's-private cursor.--}+-- | Cache set of a hash, mixing the weak low bits first.+cacheSet :: Int -> Int+cacheSet h =+    fromIntegral+        ( (fromIntegral h * 0x9E3779B97F4A7C15 :: Word64)+            `unsafeShiftR` (64 - linkCacheBits)+        )+{-# INLINE cacheSet #-}++linkCacheBits :: Int+linkCacheBits = 12++-- | Sets in each chunk's cache of recent keys (a power of two).+linkCacheSets :: Int+linkCacheSets = 1 `unsafeShiftL` linkCacheBits++-- | Scatter one row chunk through the chunk's private cursor. scatterChunk ::     VU.Vector Int ->+    VUM.IOVector Int ->     Int ->     VUM.IOVector Int ->     VUM.IOVector Int ->     VUM.IOVector Int ->+    VUM.IOVector Int ->     Int ->     Int ->     IO ()-scatterChunk hashes shift cursor sortedM sortedHashM lo hi = go lo+scatterChunk hashes link shift cursor sortedM sortedHashM sortedLinkM lo hi = go lo   where     go !i         | i >= hi = pure ()@@ -217,17 +227,19 @@             pos <- VUM.unsafeRead cursor pp             VUM.unsafeWrite sortedM pos i             VUM.unsafeWrite sortedHashM pos h+            VUM.unsafeRead link i >>= VUM.unsafeWrite sortedLinkM pos             VUM.unsafeWrite cursor pp (pos + 1)             go (i + 1) ----------------------------------------------------------------------------------- Phase 2: per-partition grouping (parallel)--------------------------------------------------------------------------------+-- | Partition indices, largest first, so skewed partitions start early.+largestFirst :: VU.Vector Int -> VU.Vector Int+largestFirst partStart =+    let !p = VU.length partStart - 1+        size i = VU.unsafeIndex partStart (i + 1) - VU.unsafeIndex partStart i+     in VU.fromList+            (map snd (L.sortBy (flip compare) [(size i, i) | i <- [0 .. p - 1]])) -{- | Group each partition with its own hash table, then rank its local groups into-canonical order — all inside the parallel worker. Forks @caps@ workers pulling-partition indices off a shared counter; disjoint keys mean no cross-partition merge.--}+-- | Group each partition with its own hash table and rank its local groups. runPartitions ::     Int ->     Int ->@@ -235,45 +247,56 @@     VU.Vector Int ->     VU.Vector Int ->     (Int -> Int -> Bool) ->+    VU.Vector Int ->     VUM.IOVector Int ->+    VUM.IOVector Int ->     VM.IOVector (VU.Vector Int) ->     VUM.IOVector Int ->     IO ()-runPartitions caps p partStart sortedRows sortedHash eqRow localGid canonBoxes nLocalGroups =-    pooledIndices caps p groupPartition+runPartitions caps p partStart sortedRows sortedHash eqRow sortedLink gidByRow localGid canonBoxes nLocalGroups =+    pooledIndices caps p (groupPartition . VU.unsafeIndex order)   where+    !order = largestFirst partStart     groupPartition !pp = do         let !s = VU.unsafeIndex partStart pp             !e = VU.unsafeIndex partStart (pp + 1)             !sz = e - s         when (sz > 0) $ do-            ht <- newHashTable sz-            repHashM <- VUM.new sz+            -- Linked rows skip the table, so size it by the unlinked ones.+            let countUnlinked !pos !k+                    | pos >= e = pure k+                    | otherwise =+                        countUnlinked (pos + 1) (if VU.unsafeIndex sortedLink pos < 0 then k + 1 else k)+            nUnlinked <- countUnlinked s 0+            ht <- newHashTable nUnlinked+            repHashM <- VUM.new nUnlinked             let loop !pos !nextGid                     | pos >= e = pure nextGid                     | otherwise = do                         let !row = VU.unsafeIndex sortedRows pos                             !h = VU.unsafeIndex sortedHash pos-                        (gid, isNew) <- htInsert ht eqRow nextGid row h-                        VUM.unsafeWrite localGid pos gid-                        if isNew+                            !r = VU.unsafeIndex sortedLink pos+                        if r >= 0                             then do-                                VUM.unsafeWrite repHashM nextGid h-                                loop (pos + 1) (nextGid + 1)-                            else loop (pos + 1) nextGid+                                -- r: an earlier unlinked row of this partition.+                                gid <- VUM.unsafeRead gidByRow r+                                VUM.unsafeWrite localGid pos gid+                                loop (pos + 1) nextGid+                            else do+                                (gid, isNew) <- htInsert ht eqRow nextGid row h+                                VUM.unsafeWrite localGid pos gid+                                VUM.unsafeWrite gidByRow row gid+                                if isNew+                                    then do+                                        VUM.unsafeWrite repHashM nextGid h+                                        loop (pos + 1) (nextGid + 1)+                                    else loop (pos + 1) nextGid             ng <- loop s 0             VUM.unsafeWrite nLocalGroups pp ng             canon <- rankByHash (VUM.unsafeRead repHashM) ng             VM.unsafeWrite canonBoxes pp canon ----------------------------------------------------------------------------------- Phase 3: global base ids + assembly----------------------------------------------------------------------------------{- | Exclusive prefix sum of the per-partition group counts into @globalBase@-(@globalBase[pp]@ = first global id of partition @pp@). Ranks were computed in-'runPartitions'; prepending the base to each yields the sequential order.--}+-- | Exclusive prefix sum of per-partition group counts into @globalBase@. canonicalize ::     Int ->     VM.IOVector (VU.Vector Int) ->@@ -292,19 +315,7 @@     canonOf <- V.unsafeFreeze canonBoxes     pure (globalBase, canonOf, total) -{- | Build the final @(valueIndices, offsets)@: the global group id of a-sorted position is @globalBase[pp] + canonOf[pp][localGid]@. @valueIndices@ orders-rows by group, @offsets@ the boundaries. (@rowToGroup@, the per-original-row-inverse, is no longer built here — 'rtgFromVisOffs' derives it on demand, so-aggregations that never read it skip its full random-write pass.)--Each partition owns a disjoint @sortedRows@ range and a disjoint global group-id-range, and its rows are exactly its groups' rows — so its first group's offset is-its own @partStart@ and every pass (group ids, offsets, placement) runs per-partition on parallel workers with no shared writes. @sortedRows@ keeps ascending-original row order inside a partition, so per-group row order matches the-sequential pass exactly.--}+-- | Build the final @(valueIndices, offsets)@ from global group ids. assemble ::     Int ->     Int ->@@ -342,8 +353,7 @@                         VUM.unsafeWrite counts g (c + 1)                         pass1 (pos + 1)             pass1 s-            -- Offsets for our group range (they start at our partStart);-            -- counts becomes the per-group write cursor.+            -- Offsets for our group range; counts becomes the write cursor.             let offsLoop !g !acc                     | g >= gEnd = pure ()                     | otherwise = do@@ -363,19 +373,14 @@                         VUM.unsafeWrite counts g (c + 1)                         pass2 (pos + 1)             pass2 s-    pooledIndices caps p doPartition+    let !order = largestFirst partStart+    pooledIndices caps p (doPartition . VU.unsafeIndex order)     VUM.unsafeWrite offsM nGroups n     offs <- VU.unsafeFreeze offsM     vis <- VU.unsafeFreeze visM     pure (vis, offs) -{- | Deferred @rowToGroup@ from @(valueIndices, offsets)@:-@rtg[vis[i]] = g@ for every @i@ in group @g@'s range. @vis@ is a permutation,-so any split of the position space writes disjoint slots; each worker binary-searches its first group and then walks group ranges. Values are identical to-the @rowToGroup@ the assembly pass used to build inline. Pure w.r.t. its-immutable inputs, so the 'unsafePerformIO' is safe.--}+-- | Deferred @rowToGroup@ from @(valueIndices, offsets)@. rtgFromVisOffs :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int rtgFromVisOffs n vis offs = unsafePerformIO $ do     let !nGroups = VU.length offs - 1

@@ -1,14 +1,16 @@ {-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TupleSections #-} {-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_GHC -Wno-orphans #-} @@ -27,6 +29,7 @@  import Data.Bifunctor (first) import qualified Data.Map as M+import qualified Data.Set as S import qualified Data.Text as T import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) import qualified Data.Vector as V@@ -34,20 +37,26 @@ import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM import DataFrame.Errors+import DataFrame.Internal.Aggregation.Kernel.Distinct (+    countDistinctColumn,+    countDistinctGroups,+ ) import DataFrame.Internal.Column-import DataFrame.Internal.Column.Bitmap import DataFrame.Internal.DataFrame import DataFrame.Internal.Expression import qualified DataFrame.Internal.Grouping as G+import qualified DataFrame.Internal.Interpreter.Kernels as K import Type.Reflection (+    SomeTypeRep (..),     Typeable,     typeRep,  )  import Data.Int (Int16, Int32, Int64, Int8)+import Data.Word (Word64)+import GHC.Float (castDoubleToWord64)  -- Specializations for common aggregation types to avoid dictionary overhead.--- foldLinearGroups: mean accumulator {-# SPECIALIZE foldLinearGroups ::     ((Double, Int) -> Double -> (Double, Int)) ->     (Double, Int) ->@@ -228,94 +237,24 @@ {-# SPECIALIZE mapColumn ::     (Int -> Double) -> Column -> Either DataFrameException Column     #-}---- zipWithColumns: binary ops-{-# SPECIALIZE zipWithColumns ::-    (Double -> Double -> Double) ->-    Column ->-    Column ->-    Either DataFrameException Column-    #-}-{-# SPECIALIZE zipWithColumns ::-    (Float -> Float -> Float) ->-    Column ->-    Column ->-    Either DataFrameException Column-    #-}-{-# SPECIALIZE zipWithColumns ::-    (Int -> Int -> Int) -> Column -> Column -> Either DataFrameException Column-    #-}-{-# SPECIALIZE zipWithColumns ::-    (Int8 -> Int8 -> Int8) -> Column -> Column -> Either DataFrameException Column-    #-}-{-# SPECIALIZE zipWithColumns ::-    (Int16 -> Int16 -> Int16) ->-    Column ->-    Column ->-    Either DataFrameException Column-    #-}-{-# SPECIALIZE zipWithColumns ::-    (Int32 -> Int32 -> Int32) ->-    Column ->-    Column ->-    Either DataFrameException Column-    #-}-{-# SPECIALIZE zipWithColumns ::-    (Int64 -> Int64 -> Int64) ->-    Column ->-    Column ->-    Either DataFrameException Column-    #-}--- Bool-returning binary comparators (hot path for Expr Bool used in--- DecisionTree splits)-{-# SPECIALIZE zipWithColumns ::-    (Double -> Double -> Bool) ->-    Column ->-    Column ->-    Either DataFrameException Column-    #-}-{-# SPECIALIZE zipWithColumns ::-    (Float -> Float -> Bool) ->-    Column ->-    Column ->-    Either DataFrameException Column-    #-}-{-# SPECIALIZE zipWithColumns ::-    (Int -> Int -> Bool) ->-    Column ->-    Column ->-    Either DataFrameException Column-    #-}-{-# SPECIALIZE zipWithColumns ::-    (Bool -> Bool -> Bool) ->-    Column ->-    Column ->-    Either DataFrameException Column-    #-}- -- Bool-mapping unary ops (e.g. 'not') {-# SPECIALIZE mapColumn ::     (Bool -> Bool) -> Column -> Either DataFrameException Column     #-} ----------------------------------------------------------------------------------- Value: the unified result type----------------------------------------------------------------------------------{- | The result of interpreting an expression.  Keeps literals as scalars-until the point where a concrete column is needed, avoiding premature-broadcast allocations.--}+-- | The result of interpreting an expression; literals stay scalar. data Value a where     -- | A single value, not yet broadcast to any length.-    Scalar :: (Columnable a) => !a -> Value a-    {- | A flat column (one element per row in the flat case, or one-    element per group after aggregation).-    -}-    Flat :: (Columnable a) => !Column -> Value a-    {- | A grouped column: one 'Column' slice per group.  Only produced-    when interpreting inside a 'GroupCtx'.-    -}+    Scalar ::+        (Columnable a) =>+        !a ->+        -- | A flat column, one element per row (or per group once aggregated).+        Value a+    Flat ::+        (Columnable a) =>+        !Column ->+        -- | A grouped column: one 'Column' slice per group.+        Value a     Group :: (Columnable a) => !(V.Vector Column) -> Value a  instance (Show a) => Show (Value a) where@@ -328,31 +267,19 @@     = FlatCtx DataFrame     | GroupCtx GroupedDataFrame ----------------------------------------------------------------------------------- Materialisation----------------------------------------------------------------------------------{- | Force a 'Value' into a flat 'Column' of the given length.  Scalars-are broadcast; flat columns are returned as-is.--}+-- | Force a 'Value' into a flat 'Column' of the given length. materialize :: forall a. (Columnable a) => Int -> Value a -> Column materialize n (Scalar v) = broadcastScalar @a n v materialize _ (Flat c) = c materialize _ (Group _) =     error "materialize: cannot flatten a grouped value to a single column" -{- | Replicate a scalar to a column of length @n@, choosing the most-efficient representation.--}+-- | Replicate a scalar to a column of length @n@ in the cheapest form. broadcastScalar :: forall a. (Columnable a) => Int -> a -> Column broadcastScalar n v = case sUnbox @a of     STrue -> fromUnboxedVector (VU.replicate n v)     SFalse -> fromVector (V.replicate n v) ----------------------------------------------------------------------------------- Lifting: the core combinators--------------------------------------------------------------------------------- -- | Apply a pure function to a 'Value'. liftValue ::     (Columnable b, Columnable a) =>@@ -362,10 +289,7 @@ liftValue f (Group gs) = Group <$> V.mapM (mapColumn f) gs {-# INLINEABLE liftValue #-} -{- | Apply a binary function to two 'Value's. When one side is a 'Scalar' the-operation degenerates to 'liftValue', recovering the old @Binary op (Lit l) right@-special cases without explicit pattern matches.--}+-- | Apply a binary function to two 'Value's, degenerating when one is scalar. liftValue2 ::     (Columnable c, Columnable b, Columnable a) =>     (c -> b -> a) ->@@ -387,7 +311,7 @@     Left $ InternalException "Group count mismatch in binary operation" {-# INLINEABLE liftValue2 #-} --- | Branch on a boolean 'Value', selecting from two same-typed 'Value's.+-- | If-then-else over values already computed for every row. branchValue ::     forall a.     (Columnable a) =>@@ -397,55 +321,428 @@     Either DataFrameException (Value a) branchValue (Scalar True) l _ = Right l branchValue (Scalar False) _ r = Right r-branchValue cond (Scalar l) (Scalar r) =-    liftValue (\c -> if c then l else r) cond-branchValue cond (Scalar l) r =-    liftValue2 (\c rv -> if c then l else rv) cond r-branchValue cond l (Scalar r) =-    liftValue2 (\c lv -> if c then lv else r) cond l-branchValue (Flat cc) (Flat lc) (Flat rc) =-    Flat <$> branchColumn @a cc lc rc-branchValue (Group cgs) (Group lgs) (Group rgs)-    | V.length cgs == V.length lgs-        && V.length lgs == V.length rgs =-        Group-            <$> V.generateM-                (V.length cgs)-                ( \i ->-                    branchColumn @a (cgs V.! i) (lgs V.! i) (rgs V.! i)-                )-branchValue _ _ _ =-    Left $-        AggregatedAndNonAggregatedException-            "if-then-else branches"-            "mismatched shapes"-{-# INLINEABLE branchValue #-}+branchValue (Flat cc) l r = do+    lc <- alongside cc l+    rc <- alongside cc r+    Flat <$> chooseRows cc lc rc+branchValue (Group cgs) l r =+    Group+        <$> V.imapM+            ( \i cc -> do+                lc <- alongside cc =<< groupAt i l+                rc <- alongside cc =<< groupAt i r+                chooseRows cc lc rc+            )+            cgs -{- | Low-level column branch: given a boolean column and two same-typed-columns, produce the element-wise selection.--}-branchColumn ::+shapeMismatch :: DataFrameException+shapeMismatch =+    AggregatedAndNonAggregatedException "if-then-else branches" "mismatched shapes"++-- | A branch as a column with one value per row of the condition column.+alongside ::     forall a.-    (Columnable a) =>-    Column ->-    Column ->-    Column ->-    Either DataFrameException Column-branchColumn cc lc rc = do-    cs <- toVector @Bool @V.Vector cc-    ls <- toVector @a @V.Vector lc-    rs <- toVector @a @V.Vector rc-    pure $-        fromVector @a $-            V.zipWith3 (\c l r -> if c then l else r) cs ls rs+    (Columnable a) => Column -> Value a -> Either DataFrameException Column+alongside cc (Scalar x) = Right (broadcastScalar (columnLength cc) x)+alongside _ (Flat c) = Right c+alongside _ (Group _) = Left shapeMismatch ----------------------------------------------------------------------------------- Error enrichment--------------------------------------------------------------------------------+-- | The part of a branch that belongs to group @i@.+groupAt :: Int -> Value a -> Either DataFrameException (Value a)+groupAt _ v@(Scalar _) = Right v+groupAt i (Group gs) | i < V.length gs = Right (Flat (gs V.! i))+groupAt _ _ = Left shapeMismatch -{- | Wrap an interpretation step so that any 'TypeMismatchException' gets-annotated with the expression that was being evaluated.--}+chooseRows :: Column -> Column -> Column -> Either DataFrameException Column+chooseRows cc lc rc = do+    cs <- conditionVector cc+    let n = VU.length cs+    both <- mappendColumns lc rc+    pure (atIndicesStable (VU.imap (\i c -> if c then i else n + i) cs) both)++-- | A condition with nulls is an error.+conditionVector :: Column -> Either DataFrameException (VU.Vector Bool)+conditionVector (UnboxedColumn bm (v :: VU.Vector x))+    | Just Refl <- testEquality (typeRep @x) (typeRep @Bool)+    , maybe True (not . bitmapHasNulls (VU.length v)) bm =+        Right v+conditionVector c = toVector @Bool @VU.Vector c++-- | A value as a kernel operand: a literal or a null-free unboxed column.+operand :: forall a. (Columnable a) => Value a -> Maybe (K.Operand a)+operand (Scalar s) = Just (K.Literal s)+operand (Flat (UnboxedColumn Nothing (v :: VU.Vector x)))+    | Just Refl <- testEquality (typeRep @x) (typeRep @a) = Just (K.Column v)+operand _ = Nothing++fromOperand :: (Columnable a, VU.Unbox a) => K.Operand a -> Value a+fromOperand (K.Literal x) = Scalar x+fromOperand (K.Column v) = Flat (UnboxedColumn Nothing v)++-- | A column loop for a built-in binary operator, if it can run.+builtinBinary ::+    forall c b a.+    (Columnable c, Columnable b, Columnable a) =>+    T.Text -> Maybe (Value c -> Value b -> Maybe (Value a))+builtinBinary name = do+    Refl <- testEquality (typeRep @c) (typeRep @b)+    case (comparison, arithmetic, integral, textEquality, textContains) of+        (Just run, _, _, _, _) -> Just run+        (_, Just run, _, _, _) -> Just run+        (_, _, Just run, _, _) -> Just run+        (_, _, _, Just run, _) -> Just run+        (_, _, _, _, found) -> found+  where+    -- Int-only integer division and remainder.+    integral :: (c ~ b) => Maybe (Value c -> Value c -> Maybe (Value a))+    integral = do+        Refl <- testEquality (typeRep @a) (typeRep @Int)+        Refl <- testEquality (typeRep @c) (typeRep @Int)+        op <- case name of+            "div" -> Just div+            "mod" -> Just mod+            _ -> Nothing+        -- Only non-zero literal divisors: lazy evaluation might skip a zero.+        Just $ \l r -> case r of+            Scalar k | k /= 0 -> fromOperand <$> (K.zipOperands op <$> operand l <*> operand r)+            _ -> Nothing+    textContains :: (c ~ b) => Maybe (Value c -> Value c -> Maybe (Value a))+    textContains = do+        Refl <- testEquality (typeRep @c) (typeRep @T.Text)+        Refl <- testEquality (typeRep @a) (typeRep @Bool)+        case name of+            "contains" -> Just packedContains+            _ -> Nothing+    comparison :: (c ~ b) => Maybe (Value c -> Value c -> Maybe (Value a))+    comparison = do+        Refl <- testEquality (typeRep @a) (typeRep @Bool)+        op <- K.cmpOpByName name+        run <- K.compareOperands @c+        Just (\l r -> fromOperand <$> (run op <$> operand l <*> operand r))+    arithmetic :: (c ~ b) => Maybe (Value c -> Value c -> Maybe (Value a))+    arithmetic = do+        Refl <- testEquality (typeRep @a) (typeRep @c)+        op <- K.arithOpByName name+        run <- K.arithOperands @c+        case sUnbox @c of+            STrue -> Just (\l r -> fromOperand <$> (run op <$> operand l <*> operand r))+            SFalse -> Nothing+    textEquality :: (c ~ b) => Maybe (Value c -> Value c -> Maybe (Value a))+    textEquality = do+        Refl <- testEquality (typeRep @c) (typeRep @T.Text)+        Refl <- testEquality (typeRep @a) (typeRep @Bool)+        wantEq <- case name of+            "eq" -> Just True+            "neq" -> Just False+            _ -> Nothing+        Just (packedEquals wantEq)++-- | Column on the left, needle on the right; not used if needle has U+FFFD.+packedContains :: Value T.Text -> Value T.Text -> Maybe (Value Bool)+packedContains (Flat (PackedText Nothing p)) (Scalar t)+    | not (T.any (== '\xFFFD') t) =+        Just (Flat (UnboxedColumn Nothing (K.packedContainsLit p t)))+packedContains _ _ = Nothing++-- | Not used when the literal contains U+FFFD (see 'K.packedEqLit').+packedEquals :: Bool -> Value T.Text -> Value T.Text -> Maybe (Value Bool)+packedEquals wantEq l r = do+    (p, t) <- case (l, r) of+        (Flat (PackedText Nothing p), Scalar t) -> Just (p, t)+        (Scalar t, Flat (PackedText Nothing p)) -> Just (p, t)+        _ -> Nothing+    if T.any (== '\xFFFD') t+        then Nothing+        else Just (Flat (UnboxedColumn Nothing (K.packedEqLit wantEq p t)))++-- | As 'builtinBinary', for unary operators.+builtinUnary ::+    forall b a.+    (Columnable b, Columnable a) => T.Text -> Maybe (Value b -> Maybe (Value a))+builtinUnary "toDouble"+    | Just Refl <- testEquality (typeRep @b) (typeRep @Int)+    , Just Refl <- testEquality (typeRep @a) (typeRep @Double) =+        Just $ \v -> case operand v of+            Just (K.Column x) -> Just (Flat (UnboxedColumn Nothing (K.intToDouble x)))+            _ -> Nothing+builtinUnary "textLength"+    | Just Refl <- testEquality (typeRep @b) (typeRep @T.Text)+    , Just Refl <- testEquality (typeRep @a) (typeRep @Int) =+        Just $ \case+            Flat (PackedText Nothing p) -> Just (Flat (UnboxedColumn Nothing (K.packedCharCount p)))+            _ -> Nothing+builtinUnary _ = Nothing++applyUnary ::+    forall op b a.+    (UnaryOp op, Columnable b, Columnable a) =>+    op b a -> Value b -> Either DataFrameException (Value a)+applyUnary op v = case builtinUnary @b @a (unaryName op) >>= ($ v) of+    Just out -> Right out+    Nothing -> liftValue (fastUnaryFn @b @a (unaryName op) (unaryFn op)) v++applyBinary ::+    forall op c b a.+    (BinaryOp op, Columnable c, Columnable b, Columnable a) =>+    op c b a -> Value c -> Value b -> Either DataFrameException (Value a)+applyBinary op l r = case builtinBinary @c @b @a (binaryName op) >>= (\run -> run l r) of+    Just out -> Right out+    Nothing -> liftValue2 (binaryFn op) l r++-- | The rows of the frame an evaluation covers.+data Rows = AllRows | Rows !(VU.Vector Int)++data Env = Env {envFrame :: DataFrame, envRows :: Rows}++frameRows :: Env -> Int+frameRows env = fst (dataframeDimensions (envFrame env))++rowCount :: Env -> Int+rowCount env = case envRows env of+    AllRows -> frameRows env+    Rows sel -> VU.length sel++-- | Keep only the covered rows of a value computed for the whole frame.+restrict :: Rows -> Value a -> Value a+restrict (Rows sel) (Flat col) = Flat (atIndicesStable sel col)+restrict _ v = v++-- | The rows at these positions of the rows already covered.+narrow :: Env -> VU.Vector Int -> Env+narrow env positions = env{envRows = Rows rows}+  where+    rows = case envRows env of+        AllRows -> positions+        Rows sel -> VU.unsafeBackpermute sel positions++-- | Identifies an expression built only from columns, literals and built-ins.+data Key+    = KeyColumn T.Text SomeTypeRep+    | KeyLiteral Literal+    | KeyUnary T.Text SomeTypeRep Key+    | KeyBinary T.Text SomeTypeRep Key Key+    deriving (Eq, Ord)++data Literal+    = LitDouble Word64+    | LitInt Int+    | LitBool Bool+    | LitText T.Text+    deriving (Eq, Ord)++-- | Keys are only built for expressions up to this many nodes.+maxKeyNodes :: Int+maxKeyNodes = 32++keyOf :: (Columnable a) => Expr a -> Maybe Key+keyOf expr = fst <$> go maxKeyNodes expr+  where+    -- Returns the key and the node budget left.+    go :: forall x. (Columnable x) => Int -> Expr x -> Maybe (Key, Int)+    go budget _ | budget <= 0 = Nothing+    go budget e = case e of+        Col name -> Just (KeyColumn name (SomeTypeRep (typeRep @x)), budget - 1)+        Lit v -> (\l -> (KeyLiteral l, budget - 1)) <$> literalKey v+        Unary op (inner :: Expr b) -> do+            _ <- builtinUnary @b @x (unaryName op)+            (k, left) <- go (budget - 1) inner+            Just (KeyUnary (unaryName op) (SomeTypeRep (typeRep @b)) k, left)+        Binary op (l :: Expr c) (r :: Expr b) -> do+            _ <- builtinBinary @c @b @x (binaryName op)+            -- A division is shareable only with a non-zero literal divisor.+            if binaryName op `elem` ["div", "mod"]+                then case r of+                    Lit (k :: b)+                        | Just Refl <- testEquality (typeRep @b) (typeRep @Int)+                        , k /= 0 ->+                            Just ()+                    _ -> Nothing+                else Just ()+            (kl, afterLeft) <- go (budget - 1) l+            (kr, left) <- go afterLeft r+            Just (KeyBinary (binaryName op) (SomeTypeRep (typeRep @c)) kl kr, left)+        _ -> Nothing++literalKey :: forall a. (Columnable a) => a -> Maybe Literal+literalKey v+    | Just Refl <- testEquality (typeRep @a) (typeRep @Double) =+        Just (LitDouble (castDoubleToWord64 v))+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = Just (LitInt v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @Bool) = Just (LitBool v)+    | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) = Just (LitText v)+    | otherwise = Nothing++-- | The keys of operator nodes that occur more than once in an expression.+repeatedKeys :: (Columnable a) => Expr a -> S.Set Key+repeatedKeys expr = M.keysSet (M.filter (> 1) (count expr M.empty))+  where+    count :: forall x. (Columnable x) => Expr x -> M.Map Key Int -> M.Map Key Int+    count e seen = case e of+        Unary _ inner -> count inner (note e seen)+        Binary _ l r -> count r (count l (note e seen))+        If c l r -> count r (count l (count c seen))+        _ -> seen+    note :: forall x. (Columnable x) => Expr x -> M.Map Key Int -> M.Map Key Int+    note e seen = maybe seen (\k -> M.insertWith (+) k (1 :: Int) seen) (keyOf e)++data Sharing = Sharing+    { repeated :: !(S.Set Key)+    , rowsAsked :: !(M.Map Key Int)+    -- ^ How many rows each repeated key has been computed for so far.+    , computed :: !(M.Map Key Column)+    -- ^ Results kept for the whole frame.+    , bytesKept :: !Int+    }++-- | Results kept for reuse during one evaluation may total this many bytes.+maxBytesKept :: Int+maxBytesKept = 256 * 1024 * 1024++newtype Eval a = Eval+    {runEval :: Sharing -> Either DataFrameException (a, Sharing)}++instance Functor Eval where+    fmap f (Eval m) = Eval (fmap (first f) . m)++instance Applicative Eval where+    pure x = Eval (\s -> Right (x, s))+    Eval mf <*> Eval mx = Eval $ \s -> do+        (f, s') <- mf s+        (x, s'') <- mx s'+        Right (f x, s'')++instance Monad Eval where+    Eval m >>= k = Eval $ \s -> do+        (x, s') <- m s+        runEval (k x) s'++failing :: Either DataFrameException a -> Eval a+failing r = Eval (\s -> fmap (,s) r)++inContext :: (Show e) => Expr e -> Eval a -> Eval a+inContext expr (Eval m) = Eval (addContext expr . m)++evalFlat ::+    (Columnable a) => DataFrame -> Expr a -> Either DataFrameException (Value a)+evalFlat df expr =+    fst+        <$> runEval+            (evalRows (Env df AllRows) expr)+            (Sharing (repeatedKeys expr) M.empty M.empty 0)++evalRows :: forall a. (Columnable a) => Env -> Expr a -> Eval (Value a)+evalRows env expr = case expr of+    Lit v -> pure (Scalar v)+    Unary{} -> evalShared env expr+    Binary{} -> evalShared env expr+    If cond l r -> inContext expr $ do+        c <- evalRows env cond+        case c of+            Scalar holds -> do+                -- The other branch runs on no rows; it can only report columns.+                let (chosen, other) = if holds then (l, r) else (r, l)+                _ <- evalRows (narrow env VU.empty) other+                evalRows env chosen+            Flat cc -> do+                cs <- failing (conditionVector cc)+                let (yes, no) = K.partitionRows cs+                lv <- evalArm env yes l+                rv <- evalArm env no r+                failing (mergeArms cs yes no lv rv)+            Group _ -> failing (Left (InternalException "grouped condition in a flat frame"))+    -- Aggregations and windows need every row; then keep the covered ones.+    _ -> failing (restrict (envRows env) <$> eval (FlatCtx (envFrame env)) expr)++-- | Evaluate an operator node, reusing a result kept for the whole frame.+evalShared :: forall a. (Columnable a) => Env -> Expr a -> Eval (Value a)+evalShared env expr = Eval $ \s -> case keyOf expr of+    Just key+        | S.member key (repeated s) -> case M.lookup key (computed s) of+            Just col -> Right (restrict (envRows env) (Flat col), s)+            Nothing -> runEval (askFor key) s+    _ -> runEval (evalNode env expr) s+  where+    total = frameRows env+    askFor key = Eval $ \s ->+        let asked = M.findWithDefault 0 key (rowsAsked s) + rowCount env+         in if asked >= total && bytesKept s + 8 * total <= maxBytesKept+                then do+                    (full, s') <- runEval (evalNode env{envRows = AllRows} expr) s+                    Right (restrict (envRows env) full, keep key full s')+                else runEval (evalNode env expr) s{rowsAsked = M.insert key asked (rowsAsked s)}+    keep :: Key -> Value a -> Sharing -> Sharing+    keep key (Flat col) s =+        s+            { computed = M.insert key col (computed s)+            , bytesKept = bytesKept s + columnBytes col+            }+    keep _ _ s = s++evalNode :: forall a. (Columnable a) => Env -> Expr a -> Eval (Value a)+evalNode env expr = inContext expr $ case expr of+    Unary op inner -> evalRows env inner >>= failing . applyUnary op+    Binary op l r -> do+        lv <- evalRows env l+        rv <- evalRows env r+        failing (applyBinary op lv rv)+    _ -> evalRows env expr++columnBytes :: Column -> Int+columnBytes (UnboxedColumn _ (v :: VU.Vector x))+    | Just Refl <- testEquality (typeRep @x) (typeRep @Bool) = VU.length v+    | otherwise = 8 * VU.length v+columnBytes c = 16 * columnLength c++-- | A branch computed for every covered row, or only for the rows choosing it.+data Arm a = ForEveryRow (Value a) | ForChosenRows (Value a)++-- | A column or literal is already there for every row; else compute lazily.+evalArm :: (Columnable a) => Env -> VU.Vector Int -> Expr a -> Eval (Arm a)+evalArm env chosen expr = case expr of+    Col _ -> ForEveryRow <$> evalRows env expr+    Lit _ -> ForEveryRow <$> evalRows env expr+    _ -> ForChosenRows <$> evalRows (narrow env chosen) expr++-- | Each row's value from the branch its condition chooses.+mergeArms ::+    forall a.+    (Columnable a) =>+    VU.Vector Bool ->+    VU.Vector Int ->+    VU.Vector Int ->+    Arm a ->+    Arm a ->+    Either DataFrameException (Value a)+mergeArms cs yes no l r = case (K.mergeBranches @a, branch l, branch r, sUnbox @a) of+    (Just run, Just lb, Just rb, STrue) -> Right (Flat (UnboxedColumn Nothing (run cs lb rb)))+    _+        | VU.null no -> Right (Flat lc)+        | VU.null yes -> Right (Flat rc)+        | otherwise -> do+            both <- mappendColumns lc rc+            -- Row i reads the next unused value of its branch.+            let place (y, f) c = if c then (y + 1, f) else (y, f + 1)+                taken = VU.prescanl' place (0, VU.length yes) cs+                source = VU.zipWith (\c (y, f) -> if c then y else f) cs taken+            Right (Flat (atIndicesStable source both))+  where+    branch :: Arm a -> Maybe (K.Branch a)+    branch arm = case arm of+        ForEveryRow v ->+            (\case K.Literal x -> K.Constant x; K.Column xs -> K.EveryRow xs)+                <$> operand v+        ForChosenRows v ->+            (\case K.Literal x -> K.Constant x; K.Column xs -> K.ChosenRows xs)+                <$> operand v+    -- The generic path wants each branch as one value per row choosing it.+    chosenColumn :: VU.Vector Int -> Arm a -> Column+    chosenColumn rows arm = case arm of+        ForEveryRow v -> materialize @a (VU.length rows) (restrict (Rows rows) v)+        ForChosenRows v -> materialize @a (VU.length rows) v+    lc = chosenColumn yes l+    rc = chosenColumn no r++-- | Annotate any 'TypeMismatchException' with the expression being evaluated. addContext ::     (Show a) => Expr a -> Either DataFrameException b -> Either DataFrameException b addContext expr = first (enrichError (show expr))@@ -464,14 +761,7 @@     a <|+> _ = a enrichError _ e = e ----------------------------------------------------------------------------------- Group slicing----------------------------------------------------------------------------------{- | Given a flat column and grouping metadata, produce one 'Column' per-group.  Each result column is an O(1) slice into a sorted copy of the-input — the sort happens once, not per-group.--}+-- | Given a flat column and grouping metadata, produce one 'Column' per group. sliceGroups :: Column -> VU.Vector Int -> VU.Vector Int -> V.Vector Column sliceGroups col os indices = case col of     PackedText _ _ -> sliceGroups (materializePacked col) os indices@@ -516,14 +806,7 @@     return inv {-# INLINE invertPermutation #-} ----------------------------------------------------------------------------------- promoteColumnWith: unified numeric / text coercion for CastWith----------------------------------------------------------------------------------{- | Coerce a column to type @a@, then apply @onResult@ to each element; the handler-selects the mode (like @cast@, @castWithDefault@, or @castEither@). Handles Double/-Float/Int coercion and 'reads'-parses Text; other mismatches return 'Left'.--}+-- | Coerce a column to type @a@, then apply @onResult@ to each element. promoteColumnWith ::     forall a b.     (Columnable a, Columnable b, Read a) =>@@ -721,10 +1004,7 @@                         )                         (V.convert v) -{- | When output type @b@ is @Maybe c@ (or @Maybe (Maybe c)@) and the column stores-plain @c@, wrap each element in 'Just' (the double-Maybe case collapses to a single-@Maybe c@). Returns 'Nothing' when neither condition holds.--}+-- | Wrap each element in 'Just' when the output type is @Maybe c@. tryMaybeWrap ::     forall a b.     (Columnable a, Columnable b) =>@@ -762,14 +1042,7 @@                 , errorColumnName = Nothing                 } ----------------------------------------------------------------------------------- eval: the unified interpreter----------------------------------------------------------------------------------{- | Evaluate an expression in a given context, producing a 'Value'.-This single function replaces both the old @interpret@ (flat) and-@interpretAggregation@ (grouped) code paths.--}+-- | Evaluate an expression in a given context, producing a 'Value'. eval ::     forall a.     (Columnable a) =>@@ -840,14 +1113,17 @@             Flat <$> promoteColumnWith onResult col         Group gs ->             Group <$> V.mapM (promoteColumnWith onResult) gs+eval (FlatCtx df) expr@(Unary{}) = evalFlat df expr+eval (FlatCtx df) expr@(Binary{}) = evalFlat df expr+eval (FlatCtx df) expr@(If{}) = evalFlat df expr eval ctx expr@(Unary op (inner :: Expr b)) = addContext expr $ do     v <- eval @b ctx inner-    liftValue (fastUnaryFn @b @a (unaryName op) (unaryFn op)) v+    applyUnary op v eval ctx expr@(Binary op (left :: Expr c) (right :: Expr b)) =     addContext expr $ do         l <- eval @c ctx left         r <- eval @b ctx right-        liftValue2 (binaryFn op) l r+        applyBinary op l r eval ctx expr@(If cond l r) = addContext expr $ do     c <- eval @Bool ctx cond     lv <- eval @a ctx l@@ -871,9 +1147,7 @@             ( InternalException                 "Over (window function) is not supported inside a grouped context"             )--- Fast path: FoldAgg (seeded) on a bare Col in GroupCtx.--- Avoids the O(n) backpermute in sliceGroups by folding directly over--- permuted indices.  Only matches when inner is exactly (Col name).+-- Fast path: seeded FoldAgg on a bare Col folds over permuted indices.  eval (GroupCtx gdf) expr@(Agg (FoldAgg _ (Just seed) (f :: a -> b -> a)) (Col name :: Expr b)) =     addContext expr $@@ -906,6 +1180,20 @@                                 (M.keys $ columnIndices $ fullDataframe gdf)                     Just col ->                         Flat <$> foldl1DirectGroups @b f col (valueIndices gdf) (offsets gdf)+-- Fast path: countDistinct on a bare Col goes to the partitioned hash kernel.+eval ctx (Agg (MergeAgg "countDistinct" _ _ _ _) (Col name))+    | Just Refl <- testEquality (typeRep @a) (typeRep @Int)+    , Just col <- getColumn name (ctxFrame ctx)+    , Just v <- distinctValue col =+        Right v+  where+    ctxFrame (FlatCtx df) = df+    ctxFrame (GroupCtx gdf) = fullDataframe gdf+    distinctValue col = case ctx of+        FlatCtx _ -> Scalar <$> countDistinctColumn col+        GroupCtx gdf ->+            Flat . fromUnboxedVector+                <$> countDistinctGroups (rowToGroup gdf) (numGroups gdf) col -- Fast path: MergeAgg on a bare Col in GroupCtx.  eval@@ -950,6 +1238,10 @@             Group gs ->                 Flat . fromVector                     <$> V.mapM (foldlColumn @b @a f seed) gs+eval (FlatCtx df) (Agg (MergeAgg "mean" _ _ _ (_ :: acc -> a)) (Col name))+    | Just Refl <- testEquality (typeRep @a) (typeRep @Double)+    , Just v <- fusedMean df name =+        Right (Scalar v) eval     ctx     expr@( Agg@@ -972,6 +1264,11 @@                 Group gs ->                     Flat . fromVector                         <$> V.mapM (fmap finalize . foldlColumn @b step seed) gs+-- Fast path: whole-frame sum of a clean Int/Double column, one parallel loop.+eval (FlatCtx df) (Agg (FoldAgg "sum" Nothing (_ :: a -> b -> a)) inner)+    | Just Refl <- testEquality (typeRep @a) (typeRep @b)+    , Just v <- fusedSum df inner =+        Right (Scalar v) eval ctx expr@(Agg (FoldAgg _ Nothing (f :: a -> b -> a)) inner) =     addContext expr $         case testEquality (typeRep @a) (typeRep @b) of@@ -993,15 +1290,7 @@                         Flat . fromVector                             <$> V.mapM (foldl1Column @a f) gs -{- | The op's element function, with a fast path for @toDouble@ (matched by-'unaryName', like the @toDouble@ peeling in "DataFrame.Internal.Simplify").-The closure captured at 'Expr'-construction time is @realToFrac@, which at an-integral source type without a fired rewrite rule lowers to-@fromRational . toRational@ — a 'Rational' allocation plus 'fromRat' per-element. 'fromIntegral' at the concrete type is the same correctly-rounded-conversion, so the swap is bit-identical; only the constant factor changes.-Non-integral sources and other ops keep the stored function.--}+-- | The op's element function, with a fast path for @toDouble@. fastUnaryFn ::     forall b a.     (Columnable b, Columnable a) =>@@ -1046,26 +1335,19 @@     go 0 !acc _ = acc     go k !acc x = go (k - 1) (step acc x) x -{- | Apply a 'CollectAgg' function to a single column, extracting the-appropriate vector type and applying the aggregation function.--}+-- | Apply a 'CollectAgg' function to a single column. applyCollect ::     forall v b a.     (VG.Vector v b, Typeable v, Columnable b, Columnable a) =>     (v b -> a) -> Column -> Either DataFrameException a applyCollect f col = f <$> toVector @b @v col -{- | Result of interpreting an expression in a grouped context.-Retained for backward compatibility with 'aggregate' and friends.--}+-- | Result of interpreting an expression in a grouped context. data AggregationResult a     = UnAggregated Column     | Aggregated (TypedColumn a) -{- | Interpret an expression against a flat 'DataFrame', producing a typed column.-Calls 'eval' then 'materialize'; 'Lit' values are broadcast here at the boundary-rather than eagerly.--}+-- | Interpret an expression against a flat 'DataFrame', producing a column. interpret ::     forall a.     (Columnable a) =>@@ -1074,10 +1356,7 @@     v <- eval (FlatCtx df) expr     pure $ TColumn $ materialize @a (fst (dataframeDimensions df)) v -{- | Interpret an expression against a 'GroupedDataFrame',-distinguishing aggregated results from bare column references.-Internally calls 'eval'.--}+-- | Interpret an expression against a 'GroupedDataFrame'. interpretAggregation ::     forall a.     (Columnable a) =>@@ -1096,3 +1375,57 @@             Right $ Aggregated $ TColumn col         Group _ ->             Right $ UnAggregated $ BoxedColumn @T.Text Nothing V.empty++-- | The mean of a clean Int or Double column. Empty gives NaN, like 'mean'.+fusedMean :: DataFrame -> T.Text -> Maybe Double+fusedMean df name = do+    UnboxedColumn Nothing (v :: VU.Vector t) <- getColumn name df+    total <- case testEquality (typeRep @t) (typeRep @Int) of+        Just Refl -> Just (K.parSumBy fromIntegral v)+        Nothing -> case testEquality (typeRep @t) (typeRep @Double) of+            Just Refl -> Just (K.parSumBy id v)+            Nothing -> Nothing+    pure $! if VU.null v then 0 / 0 else total / fromIntegral (VU.length v)++-- | The sum of a column, or of an arithmetic combination with a literal.+fusedSum :: forall a. (Columnable a) => DataFrame -> Expr a -> Maybe a+fusedSum df expr = case expr of+    Col name -> withClean name Nothing+    Binary op (Col name :: Expr c) (Lit (k :: b))+        | Just Refl <- testEquality (typeRep @c) (typeRep @a)+        , Just Refl <- testEquality (typeRep @b) (typeRep @a) ->+            withClean name (Just (binaryName op, k, False))+    Binary op (Lit (k :: c)) (Col name :: Expr b)+        | Just Refl <- testEquality (typeRep @c) (typeRep @a)+        , Just Refl <- testEquality (typeRep @b) (typeRep @a) ->+            withClean name (Just (binaryName op, k, True))+    _ -> Nothing+  where+    -- The operator is picked by name at a concrete element type.+    withClean :: T.Text -> Maybe (T.Text, a, Bool) -> Maybe a+    withClean name combo = do+        UnboxedColumn Nothing (v :: VU.Vector t) <- getColumn name df+        Refl <- testEquality (typeRep @t) (typeRep @a)+        case testEquality (typeRep @a) (typeRep @Int) of+            Just Refl -> sumInt combo v+            Nothing -> case testEquality (typeRep @a) (typeRep @Double) of+                Just Refl -> sumDouble combo v+                Nothing -> Nothing++sumInt :: Maybe (T.Text, Int, Bool) -> VU.Vector Int -> Maybe Int+sumInt combo v = case combo of+    Nothing -> Just (K.parSumBy id v)+    Just ("add", k, _) -> Just (K.parSumBy (+ k) v)+    Just ("mult", k, _) -> Just (K.parSumBy (* k) v)+    Just ("sub", k, False) -> Just (K.parSumBy (subtract k) v)+    Just ("sub", k, True) -> Just (K.parSumBy (k -) v)+    _ -> Nothing++sumDouble :: Maybe (T.Text, Double, Bool) -> VU.Vector Double -> Maybe Double+sumDouble combo v = case combo of+    Nothing -> Just (K.parSumBy id v)+    Just ("add", k, _) -> Just (K.parSumBy (+ k) v)+    Just ("mult", k, _) -> Just (K.parSumBy (* k) v)+    Just ("sub", k, False) -> Just (K.parSumBy (subtract k) v)+    Just ("sub", k, True) -> Just (K.parSumBy (k -) v)+    _ -> Nothing