diff --git a/dataframe-core.cabal b/dataframe-core.cabal --- a/dataframe-core.cabal +++ b/dataframe-core.cabal @@ -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, diff --git a/src-internal/DataFrame/Display/Terminal/PrettyPrint.hs b/src-internal/DataFrame/Display/Terminal/PrettyPrint.hs --- a/src-internal/DataFrame/Display/Terminal/PrettyPrint.hs +++ b/src-internal/DataFrame/Display/Terminal/PrettyPrint.hs @@ -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 @\@). --} +-- | Escape a value for a GitHub-flavoured Markdown table cell. escapeMarkdownCell :: T.Text -> T.Text escapeMarkdownCell = T.replace "\n" "
" . T.replace "|" "\\|" diff --git a/src-internal/DataFrame/Internal/Aggregation/Kernel/Dense.hs b/src-internal/DataFrame/Internal/Aggregation/Kernel/Dense.hs --- a/src-internal/DataFrame/Internal/Aggregation/Kernel/Dense.hs +++ b/src-internal/DataFrame/Internal/Aggregation/Kernel/Dense.hs @@ -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)] -> diff --git a/src-internal/DataFrame/Internal/Aggregation/Kernel/Distinct.hs b/src-internal/DataFrame/Internal/Aggregation/Kernel/Distinct.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/Internal/Aggregation/Kernel/Distinct.hs @@ -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 #-} diff --git a/src-internal/DataFrame/Internal/Aggregation/Kernel/Fused.hs b/src-internal/DataFrame/Internal/Aggregation/Kernel/Fused.hs --- a/src-internal/DataFrame/Internal/Aggregation/Kernel/Fused.hs +++ b/src-internal/DataFrame/Internal/Aggregation/Kernel/Fused.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Aggregation/Kernel/Moments.hs b/src-internal/DataFrame/Internal/Aggregation/Kernel/Moments.hs --- a/src-internal/DataFrame/Internal/Aggregation/Kernel/Moments.hs +++ b/src-internal/DataFrame/Internal/Aggregation/Kernel/Moments.hs @@ -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 #-} diff --git a/src-internal/DataFrame/Internal/Aggregation/Kernel/Scatter.hs b/src-internal/DataFrame/Internal/Aggregation/Kernel/Scatter.hs --- a/src-internal/DataFrame/Internal/Aggregation/Kernel/Scatter.hs +++ b/src-internal/DataFrame/Internal/Aggregation/Kernel/Scatter.hs @@ -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 #-} diff --git a/src-internal/DataFrame/Internal/Aggregation/Plan.hs b/src-internal/DataFrame/Internal/Aggregation/Plan.hs --- a/src-internal/DataFrame/Internal/Aggregation/Plan.hs +++ b/src-internal/DataFrame/Internal/Aggregation/Plan.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Aggregation/Reduction.hs b/src-internal/DataFrame/Internal/Aggregation/Reduction.hs --- a/src-internal/DataFrame/Internal/Aggregation/Reduction.hs +++ b/src-internal/DataFrame/Internal/Aggregation/Reduction.hs @@ -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) diff --git a/src-internal/DataFrame/Internal/Algorithms/Hash.hs b/src-internal/DataFrame/Internal/Algorithms/Hash.hs --- a/src-internal/DataFrame/Internal/Algorithms/Hash.hs +++ b/src-internal/DataFrame/Internal/Algorithms/Hash.hs @@ -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 #-} diff --git a/src-internal/DataFrame/Internal/Algorithms/Rank/Radix.hs b/src-internal/DataFrame/Internal/Algorithms/Rank/Radix.hs --- a/src-internal/DataFrame/Internal/Algorithms/Rank/Radix.hs +++ b/src-internal/DataFrame/Internal/Algorithms/Rank/Radix.hs @@ -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 #-} diff --git a/src-internal/DataFrame/Internal/Algorithms/Sort/Radix/Parallel.hs b/src-internal/DataFrame/Internal/Algorithms/Sort/Radix/Parallel.hs --- a/src-internal/DataFrame/Internal/Algorithms/Sort/Radix/Parallel.hs +++ b/src-internal/DataFrame/Internal/Algorithms/Sort/Radix/Parallel.hs @@ -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 -> diff --git a/src-internal/DataFrame/Internal/Column.hs b/src-internal/DataFrame/Internal/Column.hs --- a/src-internal/DataFrame/Internal/Column.hs +++ b/src-internal/DataFrame/Internal/Column.hs @@ -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, diff --git a/src-internal/DataFrame/Internal/Column/Base.hs b/src-internal/DataFrame/Internal/Column/Base.hs --- a/src-internal/DataFrame/Internal/Column/Base.hs +++ b/src-internal/DataFrame/Internal/Column/Base.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Column/Bitmap.hs b/src-internal/DataFrame/Internal/Column/Bitmap.hs --- a/src-internal/DataFrame/Internal/Column/Bitmap.hs +++ b/src-internal/DataFrame/Internal/Column/Bitmap.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Column/Builder.hs b/src-internal/DataFrame/Internal/Column/Builder.hs --- a/src-internal/DataFrame/Internal/Column/Builder.hs +++ b/src-internal/DataFrame/Internal/Column/Builder.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Column/Conversion.hs b/src-internal/DataFrame/Internal/Column/Conversion.hs --- a/src-internal/DataFrame/Internal/Column/Conversion.hs +++ b/src-internal/DataFrame/Internal/Column/Conversion.hs @@ -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) => diff --git a/src-internal/DataFrame/Internal/Column/Encode.hs b/src-internal/DataFrame/Internal/Column/Encode.hs --- a/src-internal/DataFrame/Internal/Column/Encode.hs +++ b/src-internal/DataFrame/Internal/Column/Encode.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Column/Merge.hs b/src-internal/DataFrame/Internal/Column/Merge.hs --- a/src-internal/DataFrame/Internal/Column/Merge.hs +++ b/src-internal/DataFrame/Internal/Column/Merge.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Column/Operations.hs b/src-internal/DataFrame/Internal/Column/Operations.hs --- a/src-internal/DataFrame/Internal/Column/Operations.hs +++ b/src-internal/DataFrame/Internal/Column/Operations.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Column/Properties.hs b/src-internal/DataFrame/Internal/Column/Properties.hs --- a/src-internal/DataFrame/Internal/Column/Properties.hs +++ b/src-internal/DataFrame/Internal/Column/Properties.hs @@ -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) diff --git a/src-internal/DataFrame/Internal/Column/Types.hs b/src-internal/DataFrame/Internal/Column/Types.hs --- a/src-internal/DataFrame/Internal/Column/Types.hs +++ b/src-internal/DataFrame/Internal/Column/Types.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Control/Concurrent.hs b/src-internal/DataFrame/Internal/Control/Concurrent.hs --- a/src-internal/DataFrame/Internal/Control/Concurrent.hs +++ b/src-internal/DataFrame/Internal/Control/Concurrent.hs @@ -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) diff --git a/src-internal/DataFrame/Internal/Data/HashTable.hs b/src-internal/DataFrame/Internal/Data/HashTable.hs --- a/src-internal/DataFrame/Internal/Data/HashTable.hs +++ b/src-internal/DataFrame/Internal/Data/HashTable.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Data/PackedText.hs b/src-internal/DataFrame/Internal/Data/PackedText.hs --- a/src-internal/DataFrame/Internal/Data/PackedText.hs +++ b/src-internal/DataFrame/Internal/Data/PackedText.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Data/PackedText/Utf8.hs b/src-internal/DataFrame/Internal/Data/PackedText/Utf8.hs --- a/src-internal/DataFrame/Internal/Data/PackedText/Utf8.hs +++ b/src-internal/DataFrame/Internal/Data/PackedText/Utf8.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/DataFrame.hs b/src-internal/DataFrame/Internal/DataFrame.hs --- a/src-internal/DataFrame/Internal/DataFrame.hs +++ b/src-internal/DataFrame/Internal/DataFrame.hs @@ -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) diff --git a/src-internal/DataFrame/Internal/Display/Pretty.hs b/src-internal/DataFrame/Internal/Display/Pretty.hs --- a/src-internal/DataFrame/Internal/Display/Pretty.hs +++ b/src-internal/DataFrame/Internal/Display/Pretty.hs @@ -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)) diff --git a/src-internal/DataFrame/Internal/Expression.hs b/src-internal/DataFrame/Internal/Expression.hs --- a/src-internal/DataFrame/Internal/Expression.hs +++ b/src-internal/DataFrame/Internal/Expression.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Expression/Operators.hs b/src-internal/DataFrame/Internal/Expression/Operators.hs --- a/src-internal/DataFrame/Internal/Expression/Operators.hs +++ b/src-internal/DataFrame/Internal/Expression/Operators.hs @@ -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 -> diff --git a/src-internal/DataFrame/Internal/Expression/Operators/Nullable.hs b/src-internal/DataFrame/Internal/Expression/Operators/Nullable.hs --- a/src-internal/DataFrame/Internal/Expression/Operators/Nullable.hs +++ b/src-internal/DataFrame/Internal/Expression/Operators/Nullable.hs @@ -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)) diff --git a/src-internal/DataFrame/Internal/Expression/Reflect.hs b/src-internal/DataFrame/Internal/Expression/Reflect.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/Internal/Expression/Reflect.hs @@ -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) diff --git a/src-internal/DataFrame/Internal/Expression/Simplify.hs b/src-internal/DataFrame/Internal/Expression/Simplify.hs --- a/src-internal/DataFrame/Internal/Expression/Simplify.hs +++ b/src-internal/DataFrame/Internal/Expression/Simplify.hs @@ -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 ∧ 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 diff --git a/src-internal/DataFrame/Internal/Grouping.hs b/src-internal/DataFrame/Internal/Grouping.hs --- a/src-internal/DataFrame/Internal/Grouping.hs +++ b/src-internal/DataFrame/Internal/Grouping.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Grouping/Direct.hs b/src-internal/DataFrame/Internal/Grouping/Direct.hs --- a/src-internal/DataFrame/Internal/Grouping/Direct.hs +++ b/src-internal/DataFrame/Internal/Grouping/Direct.hs @@ -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' 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 diff --git a/src-internal/DataFrame/Internal/Grouping/Partitioned.hs b/src-internal/DataFrame/Internal/Grouping/Partitioned.hs --- a/src-internal/DataFrame/Internal/Grouping/Partitioned.hs +++ b/src-internal/DataFrame/Internal/Grouping/Partitioned.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Interpreter.hs b/src-internal/DataFrame/Internal/Interpreter.hs --- a/src-internal/DataFrame/Internal/Interpreter.hs +++ b/src-internal/DataFrame/Internal/Interpreter.hs @@ -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 diff --git a/src-internal/DataFrame/Internal/Interpreter/Kernels.hs b/src-internal/DataFrame/Internal/Interpreter/Kernels.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/Internal/Interpreter/Kernels.hs @@ -0,0 +1,270 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} + +module DataFrame.Internal.Interpreter.Kernels ( + Operand (..), + Branch (..), + CmpOp (..), + cmpOpByName, + ArithOp (..), + arithOpByName, + compareOperands, + arithOperands, + zipOperands, + mergeBranches, + partitionRows, + intToDouble, + packedEqLit, + packedContainsLit, + packedCharCount, + parSumBy, +) where + +import Control.Monad.ST (runST) +import Data.Bits ((.&.)) +import qualified Data.Text as T +import qualified Data.Text.Array as A +import Data.Text.Internal (Text (Text)) +import Data.Type.Equality (testEquality, (:~:) (Refl)) +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM +import Type.Reflection (Typeable, typeRep) + +import DataFrame.Internal.Column.Operations (parGenerateUnboxedInline) +import DataFrame.Internal.Control.Concurrent (parThreshold, parallelChunks) +import DataFrame.Internal.Data.PackedText ( + PackedOffsets (..), + PackedTextData (..), + packedLength, + selAt, + sliceEqBytes, + ) +import System.IO.Unsafe (unsafePerformIO) + +-- | One side of an operation: the same value in every row, or one per row. +data Operand a = Literal !a | Column !(VU.Vector a) + +data CmpOp = CmpLt | CmpLeq | CmpGt | CmpGeq | CmpEq | CmpNeq + deriving (Eq, Show) + +cmpOpByName :: T.Text -> Maybe CmpOp +cmpOpByName name = case name of + "lt" -> Just CmpLt + "leq" -> Just CmpLeq + "gt" -> Just CmpGt + "geq" -> Just CmpGeq + "eq" -> Just CmpEq + "neq" -> Just CmpNeq + _ -> Nothing + +data ArithOp = ArithAdd | ArithSub | ArithMult + deriving (Eq, Show) + +arithOpByName :: T.Text -> Maybe ArithOp +arithOpByName name = case name of + "add" -> Just ArithAdd + "sub" -> Just ArithSub + "mult" -> Just ArithMult + _ -> Nothing + +-- | Runs @k@ with the element type known, only for 'Double' or 'Int'. +withNumeric :: + forall a r. + (Typeable a) => + (forall x. (VU.Unbox x, Ord x, Num x) => a :~: x -> r) -> Maybe r +withNumeric k + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = + Just (k @Double Refl) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = Just (k @Int Refl) + | otherwise = Nothing +{-# INLINE withNumeric #-} + +compareOperands :: + forall a. + (Typeable a) => Maybe (CmpOp -> Operand a -> Operand a -> Operand Bool) +compareOperands = withNumeric @a (\Refl -> compareLoop) + +arithOperands :: + forall a. (Typeable a) => Maybe (ArithOp -> Operand a -> Operand a -> Operand a) +arithOperands = withNumeric @a (\Refl -> arithLoop) + +-- | One branch of an if-then-else: constant, every-row, or chosen-rows values. +data Branch a = Constant !a | EveryRow !(VU.Vector a) | ChosenRows !(VU.Vector a) + +-- | Each row's value from the branch its condition chooses. +mergeBranches :: + forall a. + (Typeable a) => Maybe (VU.Vector Bool -> Branch a -> Branch a -> VU.Vector a) +mergeBranches = withNumeric @a (\Refl -> mergeLoop) + +-- The operator must be fully applied or it is not inlined into the loop. +compareLoop :: + (VU.Unbox a, Ord a) => CmpOp -> Operand a -> Operand a -> Operand Bool +compareLoop op l r = case op of + CmpLt -> zipOperands (<) l r + CmpLeq -> zipOperands (<=) l r + CmpGt -> zipOperands (>) l r + CmpGeq -> zipOperands (>=) l r + CmpEq -> zipOperands (==) l r + CmpNeq -> zipOperands (/=) l r +{-# INLINE compareLoop #-} + +arithLoop :: + (VU.Unbox a, Num a) => ArithOp -> Operand a -> Operand a -> Operand a +arithLoop op l r = case op of + ArithAdd -> zipOperands (+) l r + ArithSub -> zipOperands (-) l r + ArithMult -> zipOperands (*) l r +{-# INLINE arithLoop #-} + +zipOperands :: + (VU.Unbox a, VU.Unbox b, VU.Unbox c) => + (a -> b -> c) -> Operand a -> Operand b -> Operand c +zipOperands f (Literal x) (Literal y) = Literal (f x y) +zipOperands f (Column xs) (Literal y) = + Column + (parGenerateUnboxedInline (VU.length xs) (\i -> f (VU.unsafeIndex xs i) y)) +zipOperands f (Literal x) (Column ys) = + Column (parGenerateUnboxedInline (VU.length ys) (f x . VU.unsafeIndex ys)) +zipOperands f (Column xs) (Column ys) = + Column + ( parGenerateUnboxedInline + (min (VU.length xs) (VU.length ys)) + (\i -> f (VU.unsafeIndex xs i) (VU.unsafeIndex ys i)) + ) +{-# INLINE zipOperands #-} + +mergeLoop :: + (VU.Unbox a) => VU.Vector Bool -> Branch a -> Branch a -> VU.Vector a +mergeLoop cs l r = case (l, r) of + -- With a value per chosen row, the output is filled in row order. + (ChosenRows _, _) -> inOrder + (_, ChosenRows _) -> inOrder + _ -> + parGenerateUnboxedInline + n + (\i -> if VU.unsafeIndex cs i then at l i 0 else at r i 0) + where + n = VU.length cs + -- Row i's value, given how many values this branch has given out. + at (Constant x) _ _ = x + at (EveryRow xs) i _ = VU.unsafeIndex xs i + at (ChosenRows xs) _ used = VU.unsafeIndex xs used + {-# INLINE at #-} + inOrder = runST $ do + out <- VUM.unsafeNew n + let go !i !lUsed !rUsed + | i >= n = pure () + | VU.unsafeIndex cs i = + VUM.unsafeWrite out i (at l i lUsed) >> go (i + 1) (lUsed + 1) rUsed + | otherwise = + VUM.unsafeWrite out i (at r i rUsed) >> go (i + 1) lUsed (rUsed + 1) + go 0 0 0 + VU.unsafeFreeze out +{-# INLINE mergeLoop #-} + +-- | The positions where the condition holds and where it does not. +partitionRows :: VU.Vector Bool -> (VU.Vector Int, VU.Vector Int) +partitionRows cs = runST $ do + let n = VU.length cs + yes <- VUM.unsafeNew n + no <- VUM.unsafeNew n + -- Write to both and advance one side, so there is no data-dependent branch. + let go !i !ny !nn + | i >= n = pure (ny, nn) + | otherwise = do + let c = fromEnum (VU.unsafeIndex cs i) + VUM.unsafeWrite yes ny i + VUM.unsafeWrite no nn i + go (i + 1) (ny + c) (nn + 1 - c) + (ny, nn) <- go 0 0 0 + (,) <$> VU.unsafeFreeze (VUM.take ny yes) <*> VU.unsafeFreeze (VUM.take nn no) + +intToDouble :: VU.Vector Int -> VU.Vector Double +intToDouble v = parGenerateUnboxedInline (VU.length v) (fromIntegral . VU.unsafeIndex v) +{-# NOINLINE intToDouble #-} + +-- | Whether each row of a packed Text column equals a literal. +packedEqLit :: Bool -> PackedTextData -> T.Text -> VU.Vector Bool +packedEqLit wantEq p@(PackedTextData arr offs sel _) (Text la lo ll) = + -- Match the layout once so each row costs two offset reads and a compare. + case (sel, offs) of + (Nothing, Offs32 o) -> gen id (off32 o) + (Nothing, Offs64 o) -> gen id (VU.unsafeIndex o) + (Just s, Offs32 o) -> gen (selAt s) (off32 o) + (Just s, Offs64 o) -> gen (selAt s) (VU.unsafeIndex o) + where + off32 o i = fromIntegral (VU.unsafeIndex o i) + -- A negative base row is the empty string. + gen base off = + parGenerateUnboxedInline (packedLength p) $ \i -> + let !r = base i + !o = if r < 0 then 0 else off r + !l = if r < 0 then 0 else off (r + 1) - o + in (l == ll && sliceEqBytes arr o l la lo ll) == wantEq + {-# INLINE gen #-} +{-# NOINLINE packedEqLit #-} + +-- | Whether each row of a packed Text column contains a literal substring. +packedContainsLit :: PackedTextData -> T.Text -> VU.Vector Bool +packedContainsLit p@(PackedTextData arr offs sel _) (Text la lo ll) + | ll == 0 = VU.replicate (packedLength p) True + | otherwise = case (sel, offs) of + (Nothing, Offs32 o) -> gen id (off32 o) + (Nothing, Offs64 o) -> gen id (VU.unsafeIndex o) + (Just s, Offs32 o) -> gen (selAt s) (off32 o) + (Just s, Offs64 o) -> gen (selAt s) (VU.unsafeIndex o) + where + off32 o i = fromIntegral (VU.unsafeIndex o i) + !first = A.unsafeIndex la lo + gen base off = + parGenerateUnboxedInline (packedLength p) $ \i -> + let !r = base i + in r >= 0 && search (off r) (off (r + 1)) + {-# INLINE gen #-} + -- Scan for the needle's first byte, then compare the rest. + search !o !end = go o + where + !lastStart = end - ll + go !k + | k > lastStart = False + | A.unsafeIndex arr k == first && sliceEqBytes arr k ll la lo ll = True + | otherwise = go (k + 1) +{-# NOINLINE packedContainsLit #-} + +-- | Character count of each row of a packed Text column (UTF-8 lead bytes). +packedCharCount :: PackedTextData -> VU.Vector Int +packedCharCount p@(PackedTextData arr offs sel _) = case (sel, offs) of + (Nothing, Offs32 o) -> gen id (off32 o) + (Nothing, Offs64 o) -> gen id (VU.unsafeIndex o) + (Just s, Offs32 o) -> gen (selAt s) (off32 o) + (Just s, Offs64 o) -> gen (selAt s) (VU.unsafeIndex o) + where + off32 o i = fromIntegral (VU.unsafeIndex o i) + gen base off = + parGenerateUnboxedInline (packedLength p) $ \i -> + let !r = base i + in if r < 0 then 0 else count (off r) (off (r + 1)) 0 + {-# INLINE gen #-} + count !k !end !n + | k >= end = n + | A.unsafeIndex arr k .&. 0xC0 == 0x80 = count (k + 1) end n + | otherwise = count (k + 1) end (n + 1) +{-# NOINLINE packedCharCount #-} + +-- | @sum (map f v)@ as one parallel loop with no intermediate column. +parSumBy :: (VU.Unbox a, Num b) => (a -> b) -> VU.Vector a -> b +parSumBy f v = unsafePerformIO $ do + parts <- + parallelChunks parThreshold (VU.length v) (\lo hi -> pure $! go lo hi 0) + pure $! sum parts + where + go !i !hi !acc + | i >= hi = acc + | otherwise = go (i + 1) hi (acc + f (VU.unsafeIndex v i)) +{-# INLINE parSumBy #-} diff --git a/src-internal/DataFrame/Internal/Row.hs b/src-internal/DataFrame/Internal/Row.hs --- a/src-internal/DataFrame/Internal/Row.hs +++ b/src-internal/DataFrame/Internal/Row.hs @@ -41,9 +41,7 @@ data Any where Value :: (Columnable a) => a -> Any - -- Saves us the extra indirection we get from making Value (Maybe a) - -- and having to unpack it again to check for nulls. - -- Instead, we just have Null as a separate constructor. + -- Null is its own constructor to avoid unpacking a Value (Maybe a). Null :: Any instance Eq Any where @@ -66,7 +64,7 @@ Just Refl -> T.pack v Nothing -> (T.pack . show) v --- | Wraps a value into an \Any\ type. This helps up represent rows as heterogenous lists. +-- | Wraps a value into an @Any@ type to represent rows as heterogeneous lists. toAny :: forall a. (Columnable a) => a -> Any toAny = Value @@ -77,10 +75,7 @@ Refl <- testEquality (typeRep @a) (typeRep @b) pure v -{- | Wrap a column cell into an 'Any', honouring the column's null bitmap: a slot -marked invalid becomes 'Null', any other slot becomes a 'Value'. Only needs -@Columnable a@ (the stored element type), not @Columnable (Maybe a)@. --} +-- | Wrap a column cell into an 'Any', honouring the column's null bitmap. cellAny :: (Columnable a) => Maybe Bitmap -> Int -> a -> Any cellAny Nothing _ x = Value x cellAny (Just bm) i x = if bitmapTestBit bm i then Value x else Null @@ -130,13 +125,7 @@ <> name ) -{- | Convert the whole dataframe to a list of rows, one per row index in natural -order; each row lists all columns ordered by column index. Materializes every -row, so prefer 'toRowVector' for large frames. - ->>> toRowList df -[[("name", "Alice"), ("age", 25), ...], [("name", "Bob"), ("age", 30), ...], ...] --} +-- | Convert the whole dataframe to a list of rows, in natural order. toRowList :: DataFrame -> [[(T.Text, Any)]] toRowList df = let @@ -146,22 +135,11 @@ (zip names . V.toList . mkRowRep df names) [0 .. (fst (dataframeDimensions df) - 1)] -{- | Convert the dataframe to a vector of rows containing only the named columns, -in the given order. An empty name list yields one empty row per dataframe row. - ->>> toRowVector ["name", "age"] df -Vector of rows with only name and age fields --} +-- | Convert the dataframe to a vector of rows with only the named columns. toRowVector :: [T.Text] -> DataFrame -> V.Vector Row toRowVector names df = V.generate (fst (dataframeDimensions df)) (mkRowRep df names) -{- | Given a row gets the value associated with a field. - -==== __Examples__ - ->>> map (rowValue (F.col @Int "age")) (toRowList df) -[25,30, ...] --} +-- | Given a row, get the value associated with a field. rowValue :: forall a. Expr a -> [(T.Text, Any)] -> Maybe a rowValue (Col name) row = lookup name row >>= fromAny @a rowValue _ _ = error "Can only get rowValue of column reference" @@ -184,8 +162,7 @@ BoxedColumn bm column -> cellAny bm 0 (column V.! 0) _ -> error "mkRowFromArgs: materializeMerged is boxed" --- Returns row values in the caller's requested column order, not the --- dataframe's storage order. +-- Rows follow the caller's requested column order, not storage order. mkRowRep :: DataFrame -> [T.Text] -> Int -> Row mkRowRep df names i = V.generate (L.length names) (\index -> get (names' V.! index)) where diff --git a/src-internal/DataFrame/Internal/Row/RowHash.hs b/src-internal/DataFrame/Internal/Row/RowHash.hs --- a/src-internal/DataFrame/Internal/Row/RowHash.hs +++ b/src-internal/DataFrame/Internal/Row/RowHash.hs @@ -5,10 +5,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | Row-hash kernels with a parallel driver, feeding grouping and the join -build/probe. Each row's hash depends only on its own bytes, so hashing disjoint -ranges in parallel is race-free and bit-identical to the sequential pass. --} +-- | Row-hash kernels with a parallel driver, feeding grouping and joins. module DataFrame.Internal.Row.RowHash ( computeRowHashesIO, computeRowHashesWithIO, @@ -54,31 +51,15 @@ selAt, ) -{- | At least this many rows make the fork/coordination overhead of the parallel -hash worth it. Below it the sequential single range is used. Matches the -grouping/join parallel thresholds so the whole pipeline switches together. --} +-- | Row count from which the parallel hash beats the sequential single range. parRowHashThreshold :: Int parRowHashThreshold = 200000 -{- | Compute the per-row key hash over the selected key columns of an @n@-row -frame. Forks one worker per capability over disjoint row ranges when the row -count justifies it, else hashes the single full range; output is capability-independent. - -Dictionary-code hashing is disabled: this is the join entry point, and joins -hash each side by its own representation, so a canonical dict column on one -side of a join against a plain 'T.Text' (or non-canonical packed) column on -the other must byte-hash to keep both sides bucketing identically. --} +-- | Compute the per-row key hash over the selected key columns. computeRowHashesIO :: Int -> [Column] -> IO (VU.Vector Int) computeRowHashesIO = computeRowHashesWithIO False -{- | 'computeRowHashesIO' with an explicit dictionary-code switch. When -@useDictCodes@ is 'True', a canonical dict-encoded 'PackedText' column mixes -its 'Int' code per row instead of its byte slice (equal strings share a code, -so bucketing within one frame is preserved). Only sound when every consumer of -the hashes uses the same rule — the grouping path passes 'True', joins 'False'. --} +-- | 'computeRowHashesIO' with an explicit dictionary-code switch. computeRowHashesWithIO :: Bool -> Int -> [Column] -> IO (VU.Vector Int) computeRowHashesWithIO useDictCodes n selected = do mv <- VUM.unsafeNew (max 1 n) @@ -86,10 +67,7 @@ parallelChunks_ parRowHashThreshold n runRange VU.unsafeFreeze (VUM.slice 0 n mv) -{- | Mix every selected column over the row range @[lo, hi)@ into @mv@, seeding -each slot with 'fnvOffset'. Must match the sequential grouping hash byte-for-byte -(at the same @useDictCodes@ setting) so grouping and joins bucket identically. --} +-- | Mix every selected column over a row range into @mv@. hashRowRange :: Bool -> VUM.IOVector Int -> Int -> Int -> [Column] -> IO () hashRowRange useDictCodes mv lo hi cols = do seedRange mv lo hi @@ -102,10 +80,7 @@ | i >= hi = pure () | otherwise = VUM.unsafeWrite mv i fnvOffset >> go (i + 1) -{- | Fold one column's values over @[lo, hi)@ into the running hashes. The branch -structure mirrors the sequential grouping hash: typed unboxed fast paths, then a -'mixShow' fallback, with the null bitmap mixing 'nullSalt'. --} +-- | Fold one column's values over @[lo, hi)@ into the running hashes. mixColumnRange :: Bool -> VUM.IOVector Int -> Int -> Int -> Column -> IO () mixColumnRange useDictCodes mv lo hi = \case c@(MergedColumn _ _) -> mixColumnRange useDictCodes mv lo hi (materializeMerged c) @@ -131,9 +106,7 @@ Nothing -> boxedRange mv lo hi bm mixShow v PackedText bm p -> packedRange useDictCodes mv lo hi bm p -{- | Mix an unboxed column's range, mixing 'nullSalt' at null slots. @INLINE@d to -specialise on the element type and mixing function per call site. --} +-- | Mix an unboxed column's range, mixing 'nullSalt' at null slots. unboxedRange :: (VU.Unbox a) => VUM.IOVector Int -> @@ -177,13 +150,7 @@ go (i + 1) {-# INLINE boxedRange #-} -{- | Mix a packed-text column's range over its raw UTF-8 byte slices. The -unselected payload is the hot path (indexes the offset vector directly); a -selected payload (a gather/join result) falls back to 'packedSlice'. When -@useDictCodes@ holds and the selection is a canonical dictionary encoding -(equal strings share a code), each row mixes its 'Int' code with one 'mixInt' -instead of walking the string bytes. --} +-- | Mix a packed-text column's range over its raw UTF-8 byte slices. packedRange :: Bool -> VUM.IOVector Int -> diff --git a/src/DataFrame/Core.hs b/src/DataFrame/Core.hs --- a/src/DataFrame/Core.hs +++ b/src/DataFrame/Core.hs @@ -1,13 +1,11 @@ -{- | The curated public surface of @dataframe-core@: the interchange types -('DataFrame', 'Column', 'Row', 'Expr'), element constraints, and the -rendering/serialization verbs. Internal plumbing stays in @dataframe-core:internal@. --} +-- | The curated public surface of dataframe-core: types, constraints, verbs. module DataFrame.Core ( -- * The DataFrame DataFrame, GroupedDataFrame, empty, - fromNamedColumns, + fromColumns, + (=:), insertColumn, columnNames, null, @@ -41,8 +39,8 @@ Expr, NamedExpr, toNamedExpr, - toUExpr, - fromUExpr, + toSomeExpr, + fromSomeExpr, eSize, prettyPrint, prettyPrintWidth, @@ -80,7 +78,7 @@ columnNames, defaultTruncateConfig, empty, - fromNamedColumns, + fromColumns, insertColumn, null, toCsv, @@ -88,16 +86,17 @@ toMarkdown, toMarkdown', toSeparated, + (=:), ) import DataFrame.Internal.Expression ( Expr, NamedExpr, eSize, - fromUExpr, + fromSomeExpr, prettyPrint, prettyPrintWidth, toNamedExpr, - toUExpr, + toSomeExpr, ) import DataFrame.Internal.Row ( Any, diff --git a/src/DataFrame/Expression/Operators.hs b/src/DataFrame/Expression/Operators.hs --- a/src/DataFrame/Expression/Operators.hs +++ b/src/DataFrame/Expression/Operators.hs @@ -1,10 +1,4 @@ -{- | -The public name for dataframe's expression operators. - -This is a re-export shim: the implementation lives in -"DataFrame.Internal.Expression.Operators", which is internal and may be -reorganised without notice. Depend on this module instead. --} +-- | Public name for dataframe's expression operators (re-export shim). module DataFrame.Expression.Operators ( module DataFrame.Internal.Expression.Operators, ) where diff --git a/src/DataFrame/Typed/Freeze.hs b/src/DataFrame/Typed/Freeze.hs --- a/src/DataFrame/Typed/Freeze.hs +++ b/src/DataFrame/Typed/Freeze.hs @@ -31,9 +31,7 @@ import DataFrame.Typed.Schema (KnownSchema (..)) import DataFrame.Typed.Types (TypedDataFrame (..)) -{- | Validate that an untyped 'DataFrame' matches the expected schema @cols@, -then wrap it. Returns 'Nothing' on mismatch. --} +-- | Validate an untyped 'DataFrame' against schema @cols@ and wrap it. freeze :: forall cols. (KnownSchema cols) => D.DataFrame -> Maybe (TypedDataFrame cols) freeze df = case validateSchema @cols df of @@ -49,17 +47,12 @@ Left err -> Left err Right _ -> Right (TDF df) -{- | Validate and wrap like 'freezeWithError', but throw a 'DataFrameException' -in 'IO' on mismatch. The throwing boundary used by the typed readers -(@readCsv@ \/ @readParquet@). --} +-- | Like 'freezeWithError' but throws a 'DataFrameException' on mismatch. freezeOrThrow :: forall cols. (KnownSchema cols) => D.DataFrame -> IO (TypedDataFrame cols) freezeOrThrow = either (throwIO . InternalException) pure . freezeWithError @cols -{- | Unwrap a typed DataFrame back to the untyped representation. -Always safe; discards type information. --} +-- | Unwrap a typed DataFrame back to the untyped representation. thaw :: TypedDataFrame cols -> D.DataFrame thaw (TDF df) = df @@ -72,9 +65,7 @@ instance ToDataFrame (TypedDataFrame cols) where toDataFrame = thaw -{- | Wrap an untyped DataFrame without any validation. -Used internally after delegation where the library guarantees schema correctness. --} +-- | Wrap an untyped DataFrame without any validation. unsafeFreeze :: D.DataFrame -> TypedDataFrame cols unsafeFreeze = TDF @@ -105,17 +96,13 @@ <> ", got " <> T.pack (C.columnTypeString col) -{- | Check if a Column's element type matches the expected SomeTypeRep. -For nullable columns (those with a bitmap), @Maybe a@ in the schema matches -a column whose inner type is @a@, since we store nullable data as -@BoxedColumn (Just bm) a@ or @UnboxedColumn (Just bm) a@ rather than -@Column (Maybe a)@. --} +-- | Check that a Column's element type matches the expected SomeTypeRep. matchesType :: SomeTypeRep -> C.Column -> Bool matchesType expected col = let expectedStr = show expected colTypeStr = C.columnTypeString col in expectedStr == colTypeStr - || ( C.hasMissing col -- nullable column: schema says "Maybe X", column stores "X" with a bitmap + -- A nullable column stores "X" under a "Maybe X" schema type. + || ( C.hasMissing col && Just colTypeStr == stripPrefix "Maybe " expectedStr ) diff --git a/src/DataFrame/Typed/Generic.hs b/src/DataFrame/Typed/Generic.hs --- a/src/DataFrame/Typed/Generic.hs +++ b/src/DataFrame/Typed/Generic.hs @@ -9,37 +9,7 @@ {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} -{- | -Module : DataFrame.Typed.Generic -License : MIT - -Generic-based opt-in for record-to-schema derivation. Mirrors the Template -Haskell splice in "DataFrame.Typed.TH" but builds the schema type from a -@GHC.Generics.Generic@ instance instead of @reify@. - -Use it like this: - -@ -data Order = Order - { orderId :: Int64 - , region :: Text - , amount :: Double - } deriving (Show, Eq, Generic) - -type OrderSchema = SchemaOf Order - -instance HasSchema Order OrderSchema where - toColumns = genericToColumns - fromColumns = genericFromColumns -@ - -Field names are translated with the @CamelCase -> snake_case@ rule -(matching 'DataFrame.Typed.TH.camelToSnake'); use 'SchemaOfRaw' if you -want the schema to keep the record selector names verbatim — in that -case you cannot use 'genericToColumns' \/ 'genericFromColumns' and must -either hand-roll the instance or use the TH splice with a custom name -transform. --} +-- | Generic-based opt-in for record-to-schema derivation. module DataFrame.Typed.Generic ( -- * Type-level schema derivation NameCase (..), @@ -88,17 +58,10 @@ import DataFrame.Typed.Schema (Append) import DataFrame.Typed.Util (camelToSnake) -{- | Field-name policy applied to record selectors when computing -'RepToSchema'. - -* 'SnakeCase' — translate @camelCaseField@ to @\"camel_case_field\"@. -* 'IdentityCase' — keep the selector name verbatim. --} +-- | Field-name policy applied to record selectors for 'RepToSchema'. data NameCase = SnakeCase | IdentityCase -{- | The schema type @'[ '(name, ty), ...]@ derived from the 'Rep' of a -record type, with the given 'NameCase' applied to each field name. --} +-- | The schema type derived from the 'Rep' of a record type. type family RepToSchema (nc :: NameCase) (r :: Type -> Type) :: [(Symbol, Type)] where RepToSchema nc (M1 D _ f) = RepToSchema nc f RepToSchema nc (M1 C _ f) = RepToSchema nc f @@ -110,7 +73,7 @@ TransformName 'SnakeCase s = CamelToSnake s TransformName 'IdentityCase s = s --- | Type-level camelCase -> snake_case. Matches 'camelToSnake' at the value level. +-- | Type-level camelCase to snake_case, matching 'camelToSnake'. type family CamelToSnake (s :: Symbol) :: Symbol where CamelToSnake s = SnakeStart (UnconsSymbol s) @@ -141,9 +104,7 @@ -- | Identity-cased schema derived from @a@'s 'Generic' representation. type SchemaOfRaw a = RepToSchema 'IdentityCase (Rep a) -{- | Walks the 'Rep' tree of a record, producing or consuming a list of -named columns. Used by 'genericToColumns' \/ 'genericFromColumns'. --} +-- | Walks the 'Rep' tree of a record, producing or consuming named columns. class GHasColumns (r :: Type -> Type) where gToColumns :: [r p] -> [(T.Text, C.Column)] gFromColumns :: D.DataFrame -> Either T.Text [r p] @@ -183,22 +144,12 @@ v <- requireColumn @a colName df pure (map (M1 . K1) (VB.toList v)) -{- | Default implementation of 'DataFrame.Typed.Record.toColumns' for any -@Generic@ record. Field names are translated with @camelCase -> snake_case@. - -@ -instance HasSchema Order (SchemaOf Order) where - toColumns = genericToColumns - fromColumns = genericFromColumns -@ --} +-- | Default 'DataFrame.Typed.Record.recordsToColumns' for any Generic record. genericToColumns :: forall a. (Generic a, GHasColumns (Rep a)) => [a] -> [(T.Text, C.Column)] genericToColumns = gToColumns . map from -{- | Default implementation of 'DataFrame.Typed.Record.fromColumns' for any -@Generic@ record. --} +-- | Default 'DataFrame.Typed.Record.recordsFromFrame' for any Generic record. genericFromColumns :: forall a. (Generic a, GHasColumns (Rep a)) => D.DataFrame -> Either T.Text [a] genericFromColumns df = fmap (map to) (gFromColumns df) diff --git a/src/DataFrame/Typed/Record.hs b/src/DataFrame/Typed/Record.hs --- a/src/DataFrame/Typed/Record.hs +++ b/src/DataFrame/Typed/Record.hs @@ -5,15 +5,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} -{- | -Module : DataFrame.Typed.Record -License : MIT - -Bridge a Haskell record type to a typed dataframe schema. Instances are -typically generated by the @deriveSchemaFromType@ Template Haskell splice -(see "DataFrame.Typed.TH"), but can be written by hand as well, or -plugged into the @Generic@-derived defaults from "DataFrame.Typed.Generic". --} +-- | Bridge a Haskell record type to a typed dataframe schema. module DataFrame.Typed.Record ( -- * Class HasSchema (..), @@ -35,49 +27,24 @@ import qualified Data.Vector as VB import qualified DataFrame.Internal.Column as C -import DataFrame.Internal.DataFrame (fromNamedColumns) +import DataFrame.Internal.DataFrame (fromColumns) import qualified DataFrame.Internal.DataFrame as D import DataFrame.Typed.Types (TypedDataFrame (..)) import GHC.TypeLits -{- | Bridge a Haskell record type @a@ to a typed-dataframe schema. - -The schema is exposed as an associated type family 'Schema' so that -instances can pick it up from a 'GHC.Generics.Rep' computation (see -'DataFrame.Typed.Generic.SchemaOf') or from an explicit list emitted by -'DataFrame.Typed.TH.deriveSchemaFromType'. - -@toColumns@ explodes a list of records into a list of named columns. -@fromColumns@ reconstructs the records from a 'D.DataFrame', returning -@Left err@ if a column is missing or has the wrong type. --} +-- | Bridge a Haskell record type to a typed-dataframe schema. class HasSchema a where type Schema a :: [(Symbol, Type)] - toColumns :: [a] -> [(T.Text, C.Column)] - fromColumns :: D.DataFrame -> Either T.Text [a] - -{- | Build an untyped 'D.DataFrame' from a list of records. - -@ -data Order = Order { orderId :: Int64, region :: Text, amount :: Double } -\$(deriveSchemaFromType ''Order) - -xs :: [Order] -xs = [Order 1 "us" 10.0, Order 2 "eu" 20.0] + recordsToColumns :: [a] -> [(T.Text, C.Column)] + recordsFromFrame :: D.DataFrame -> Either T.Text [a] -df :: DataFrame -df = fromRecords xs -@ --} +-- | Build an untyped 'D.DataFrame' from a list of records. fromRecords :: (HasSchema a) => [a] -> D.DataFrame -fromRecords = fromNamedColumns . toColumns - -{- | Parse a list of records out of an untyped 'D.DataFrame'. +fromRecords = fromColumns . recordsToColumns -Returns @Left err@ on schema mismatch (missing column, wrong type). --} +-- | Parse a list of records out of an untyped 'D.DataFrame'. toRecords :: (HasSchema a) => D.DataFrame -> Either T.Text [a] -toRecords = fromColumns +toRecords = recordsFromFrame -- | Like 'fromRecords' but returns a 'TypedDataFrame' tagged with the schema. fromRecordsTyped :: forall a. (HasSchema a) => [a] -> TypedDataFrame (Schema a) @@ -86,13 +53,9 @@ -- | Like 'toRecords' but accepts a 'TypedDataFrame'. toRecordsTyped :: forall a. (HasSchema a) => TypedDataFrame (Schema a) -> Either T.Text [a] -toRecordsTyped (TDF df) = fromColumns df - -{- | Extract a column as a boxed vector by name, returning a 'T.Text' error -on missing column or type mismatch. +toRecordsTyped (TDF df) = recordsFromFrame df -Used by code generated by 'DataFrame.Typed.TH.deriveSchemaFromType'. --} +-- | Extract a column as a boxed vector by name, or a 'T.Text' error. requireColumn :: forall a. (C.Columnable a) => T.Text -> D.DataFrame -> Either T.Text (VB.Vector a) diff --git a/src/DataFrame/Typed/Schema.hs b/src/DataFrame/Typed/Schema.hs --- a/src/DataFrame/Typed/Schema.hs +++ b/src/DataFrame/Typed/Schema.hs @@ -88,10 +88,7 @@ TypeError ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema") -{- | Like 'Lookup', but returns a harmless fallback ('Int') instead of -'TypeError' when the column is not found. Use together with -'AssertPresent' so the error fires exactly once. --} +-- | Like 'Lookup' but falls back to 'Int' instead of a 'TypeError'. type family SafeLookup (name :: Symbol) (cols :: [(Symbol, Type)]) :: Type where SafeLookup name ('(name, a) ': _) = a SafeLookup name (_ ': rest) = SafeLookup name rest @@ -253,12 +250,7 @@ TypeError ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema") -{- | Assert that each join key has the same element type in both schemas, -modulo 'Maybe'-wrapping on either side (the runtime join matches a nullable -key column against a plain one). Use together with 'AssertAllPresent', which -reports keys missing from either schema; absent keys are skipped here so the -error fires exactly once. --} +-- | Assert each join key has the same element type in both schemas. type family AssertKeyTypesMatch (keys :: [Symbol]) @@ -326,7 +318,7 @@ , AssertAllColumnsHaveType ns a cols ) --- | Is @a@ a real, unboxed numeric type — i.e. a valid numeric-column element? +-- | Is @a@ a real, unboxed numeric type? type family IsRealType (a :: Type) :: Bool where IsRealType Int = 'True IsRealType Int8 = 'True @@ -342,11 +334,7 @@ IsRealType Float = 'True IsRealType _ = 'False -{- | Emit a readable compile error when the column named @name@ is not a -real-number type, naming the calling function @fn@, the column, and the type it -actually has. Used by the numeric extractors so a wrong column type reads as a -repairable message rather than a bare @No instance for Real …@. --} +-- | Emit a readable compile error when a column is not a real-number type. type family AssertRealColumn (fn :: Symbol) (name :: Symbol) (a :: Type) :: Constraint where AssertRealColumn fn name a = AssertRealColumnGo fn name a (IsRealType a) @@ -364,11 +352,7 @@ ':<>: 'ShowType a ) -{- | Constraint that every column in the schema is a real (numeric), unboxed -type. Lets the whole-frame matrix extractors ('toDoubleMatrix' and friends) be -total — a non-numeric or nullable column is a compile error (with the offending -column named, via 'AssertRealColumn'), not a runtime 'Left'. --} +-- | Constraint that every column in the schema is a real unboxed type. type family AllColumnsReal (fn :: Symbol) (cols :: [(Symbol, Type)]) :: Constraint where AllColumnsReal fn '[] = () AllColumnsReal fn ('(n, a) ': rest) = @@ -387,21 +371,13 @@ ':$$: 'Text "Convert it (toDouble) or drop it before fitting." ) -{- | Strip 'Maybe' from all columns. Used by 'filterAllJust'. - -@'("x", (Maybe Double)@ becomes @'("x", Double@.)) -@'("y", Int@ stays @'("y", Int@.)) --} +-- | Strip 'Maybe' from all columns; used by 'filterAllJust'. type family StripAllMaybe (cols :: [(Symbol, Type)]) :: [(Symbol, Type)] where StripAllMaybe '[] = '[] StripAllMaybe ('(n, Maybe a) ': rest) = '(n, a) ': StripAllMaybe rest StripAllMaybe ('(n, a) ': rest) = '(n, a) ': StripAllMaybe rest -{- | Strip 'Maybe' from a single named column. Used by 'filterJust'. - -@StripMaybeAt "x" '[ '("x", Maybe Double), '("y", Int)]@ - = @'[ '("x", Double), '("y", Int)]@ --} +-- | Strip 'Maybe' from a single named column; used by 'filterJust'. type family StripMaybeAt (name :: Symbol) (cols :: [(Symbol, Type)]) :: [(Symbol, Type)] where StripMaybeAt name ('(name, Maybe a) ': rest) = '(name, a) ': rest StripMaybeAt name ('(name, a) ': rest) = '(name, a) ': rest @@ -635,19 +611,11 @@ (T.pack (symbolVal (Proxy @name)), someTypeRep (Proxy @a)) : schemaEvidence @rest -{- | The column names a schema declares, in schema order. Pass it to a reader's -options to fetch only those columns: - -@ -D.readCsvWithOpts - D.defaultReadOptions{D.readColumns = Just (schemaColumnNames \@(Schema Customer))} - "customers.csv" -@ --} +-- | The column names a schema declares, in schema order. schemaColumnNames :: forall cols. (KnownSchema cols) => [T.Text] schemaColumnNames = map fst (schemaEvidence @cols) --- | A class that provides a list of 'Text' values for a type-level list of Symbols. +-- | A class giving a list of 'Text' values for a type-level list of Symbols. class AllKnownSymbol (names :: [Symbol]) where symbolVals :: [T.Text] diff --git a/src/DataFrame/Typed/Types.hs b/src/DataFrame/Typed/Types.hs --- a/src/DataFrame/Typed/Types.hs +++ b/src/DataFrame/Typed/Types.hs @@ -42,14 +42,9 @@ import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.Column.Types (These (..)) import qualified DataFrame.Internal.DataFrame as D -import DataFrame.Internal.Expression (Expr, NamedExpr, UExpr (..)) - -{- | A phantom-typed wrapper over the untyped 'DataFrame'. +import DataFrame.Internal.Expression (Expr, NamedExpr, SomeExpr (..)) -The type parameter @cols@ is a type-level list of @'(name, ty)@ pairs -that tracks the schema at compile time. All operations delegate to the -untyped core at runtime and update the phantom type at compile time. --} +-- | A phantom-typed wrapper over the untyped 'DataFrame'. newtype TypedDataFrame (cols :: [(Symbol, Type)]) = TDF {unTDF :: D.DataFrame} instance Show (TypedDataFrame cols) where @@ -58,24 +53,14 @@ instance Eq (TypedDataFrame cols) where (TDF a) == (TDF b) = a == b -{- | A typed expression validated against schema @cols@, producing values of type @a@. - -Unlike the untyped 'Expr a', a 'TExpr' can only be constructed through -type-safe combinators ('col', 'lit', arithmetic operations) that verify -column references exist in the schema with the correct type. - -Use 'unTExpr' to extract the underlying 'Expr' for delegation to the untyped API. --} +-- | A typed expression validated against schema @cols@, yielding @a@. newtype TExpr (cols :: [(Symbol, Type)]) a = TExpr {unTExpr :: Expr a} -- | Shows the underlying expression; the schema phantom is type-level only. instance (Show a) => Show (TExpr cols a) where showsPrec d (TExpr e) = showsPrec d e -{- | The typed counterpart of an untyped expression type for schema @cols@: -@AsTExpr cols (Expr r) = TExpr cols r@. Lets a result type follow the frame — -an @Expr@ over a plain frame becomes a @TExpr@ over a typed one. --} +-- | The typed counterpart of an untyped expression type for schema @cols@. type family AsTExpr (cols :: [(Symbol, Type)]) (e :: Type) :: Type where AsTExpr cols (Expr r) = TExpr cols r @@ -95,17 +80,7 @@ newtype TypedGrouped (keys :: [Symbol]) (cols :: [(Symbol, Type)]) = TGD {unTGD :: D.GroupedDataFrame} -{- | Internal aggregation chain. Each cons prepends a 'Column' to the -@aggs@ phantom list. End users never construct this directly — they -compose 'DataFrame.Typed.Aggregate.as' entries with @(.)@ and let -'DataFrame.Typed.Aggregate.aggregate' apply the composition to -'TAggNil'. - -@ -as \@\"total\" (F.sum salary) - . as \@\"avg_age\" (F.mean age) -@ --} +-- | Internal aggregation chain; users compose it via the Aggregators API. data TAgg (keys :: [Symbol]) (cols :: [(Symbol, Type)]) (aggs :: [(Symbol, Type)]) where TAggNil :: TAgg keys cols '[] TAggCons :: @@ -118,12 +93,10 @@ TAgg keys cols aggs -> TAgg keys cols ('(name, a) ': aggs) -{- | Extract the runtime 'NamedExpr' list from a 'TAgg', in -declaration order (reversed from the cons-built order). --} +-- | Extract the runtime 'NamedExpr' list from a 'TAgg' in declaration order. taggToNamedExprs :: TAgg keys cols aggs -> [NamedExpr] taggToNamedExprs = reverse . go where go :: TAgg keys cols aggs -> [NamedExpr] go TAggNil = [] - go (TAggCons name (TExpr expr) rest) = (name, UExpr expr) : go rest + go (TAggCons name (TExpr expr) rest) = (name, SomeExpr expr) : go rest diff --git a/src/DataFrame/Typed/Util.hs b/src/DataFrame/Typed/Util.hs --- a/src/DataFrame/Typed/Util.hs +++ b/src/DataFrame/Typed/Util.hs @@ -4,15 +4,7 @@ import Data.Char (isUpper, toLower) -{- | Convert a camelCase identifier to snake_case. - ->>> camelToSnake "orderId" -"order_id" ->>> camelToSnake "amountUS" -"amount_u_s" ->>> camelToSnake "region" -"region" --} +-- | Convert a camelCase identifier to snake_case. camelToSnake :: String -> String camelToSnake [] = [] camelToSnake (c : cs) = toLower c : go cs