dataframe-operations 2.4.0.0 → 2.5.0.0

raw patch · 19 files changed

+945/−603 lines, 19 filesdep ~dataframe-corePVP ok

version bump matches the API change (PVP)

Dependency ranges changed: dataframe-core

API changes (from Hackage documentation)

Files

• dataframe-operations.cabal +4/−4
• src/DataFrame/Functions.hs +3/−3
• src/DataFrame/Monad.hs +1/−1
• src/DataFrame/Operations/AggregateScatter.hs +0/−224
• src/DataFrame/Operations/Aggregation.hs +152/−11
• src/DataFrame/Operations/Aggregation/Run.hs +389/−0
• src/DataFrame/Operations/Inference.hs +0/−32
• src/DataFrame/Operations/Join.hs +110/−47
• src/DataFrame/Operations/Join/Parallel.hs +222/−0
• src/DataFrame/Operations/JoinPar.hs +0/−215
• src/DataFrame/Operations/Merge.hs +1/−1
• src/DataFrame/Operations/Permutation.hs +6/−8
• src/DataFrame/Operations/SetOps.hs +3/−16
• src/DataFrame/Operations/Statistics.hs +1/−2
• src/DataFrame/Operations/Subset.hs +48/−33
• src/DataFrame/Operations/Transformations.hs +1/−1
• src/DataFrame/Operations/Typing.hs +1/−2
• src/DataFrame/Typed/Expr.hs +2/−2
• src/DataFrame/Typed/Statistics.hs +1/−1

@@ -1,6 +1,6 @@ cabal-version:      3.4 name:               dataframe-operations-version:            2.4.0.0+version:            2.5.0.0 synopsis:           Column operations, expression DSL, and statistics for the dataframe ecosystem. description:     Untyped column operations (select, filter, sort, join, groupBy,@@ -33,12 +33,12 @@                         DataFrame.Internal.Statistics                         DataFrame.Functions                         DataFrame.Monad-                        DataFrame.Operations.AggregateScatter+                        DataFrame.Operations.Aggregation.Run                         DataFrame.Operations.Aggregation                         DataFrame.Operations.Core                         DataFrame.Operations.Inference                         DataFrame.Operations.Join-                        DataFrame.Operations.JoinPar+                        DataFrame.Operations.Join.Parallel                         DataFrame.Operations.Merge                         DataFrame.Operations.Permutation                         DataFrame.Operations.SetOps@@ -58,7 +58,7 @@     build-depends:      base >= 4 && < 5,                         bytestring >= 0.11 && < 0.14,                         containers >= 0.6.7 && < 0.10,-                        dataframe-core >= 2.4 && < 2.5,+                        dataframe-core >= 2.5 && < 2.6,                         dataframe-parsing >= 2.2 && < 2.3,                         random >= 1.2 && < 2,                         regex-tdfa >= 1.3.0 && < 2,

@@ -13,7 +13,7 @@  module DataFrame.Functions (     module DataFrame.Functions,-    module DataFrame.Operators,+    module DataFrame.Expression.Operators,     add,     sub,     mult,@@ -39,14 +39,14 @@ import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU -import DataFrame.Internal.Nullable (+import DataFrame.Expression.Operators+import DataFrame.Internal.Expression.Operators.Nullable (     BaseType,     NullLift1Op (applyNull1),     NullLift1Result,     NullLift2Op (applyNull2),     NullLift2Result,  )-import DataFrame.Operators import Text.Regex.TDFA import Type.Reflection (typeRep) import Prelude hiding (maximum, minimum)

@@ -45,7 +45,7 @@ import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.DataFrame (DataFrame) import DataFrame.Internal.Expression (Expr (..), UExpr (..), prettyPrint)-import DataFrame.Internal.Nullable (BaseType)+import DataFrame.Internal.Expression.Operators.Nullable (BaseType) import qualified DataFrame.Operations.Core as D import DataFrame.Operations.Permutation (SortOrder) import qualified DataFrame.Operations.Permutation as D

@@ -1,224 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-}--{- | Execute a recognised aggregation plan ('AggPlan') through the vectorized-scatter kernel, producing one result column (length @nGroups@, canonical group-order). The scatter reductions live in 'DataFrame.Internal.AggKernel' (sequential)-and 'DataFrame.Internal.AggKernelPar' (parallel by disjoint group range); this-module handles the compound @max - min@ combine and the holistic grouped median.-A plan only reaches here once 'planAgg' verified the value columns are clean-unboxed Int/Double, so the @error@ branches are unreachable.--Every reduction takes the Round-5 grouping layout @(valueIndices, offsets)@ so-the parallel kernel can split the group-id range across capabilities with no-cross-worker merge. Each group's rows stay in original-row order within one-worker's range, so results are byte-identical to the sequential path at any @-N@.--}-module DataFrame.Operations.AggregateScatter (runPlan, runMomentPlan) where--import qualified Data.Text as T-import qualified Data.Vector.Algorithms.Intro as VA-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM--import Control.Concurrent (forkIO, getNumCapabilities)-import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)-import Control.Exception (SomeException, throwIO, try)-import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))-import DataFrame.Internal.AggKernel (Reduction (..), scatterColumnToDouble)-import DataFrame.Internal.AggKernelDirect (directReduce, directThreshold)-import DataFrame.Internal.AggKernelPar (momentScatterPar, scatterReducePar)-import DataFrame.Internal.AggPlan (AggPlan (..), MomentPlan (..), Moments (..))-import DataFrame.Internal.Column (Column (..), fromUnboxedVector)-import DataFrame.Internal.DataFrame (GroupedDataFrame (..), getColumn)-import System.IO.Unsafe (unsafePerformIO)-import Type.Reflection (typeRep)--runPlan :: GroupedDataFrame -> VU.Vector Int -> Int -> AggPlan -> Column-runPlan gdf rtg nGroups plan = case plan of-    PlanScatter red name -> scatterColumn red name-    PlanMaxMinusMin a b -> maxMinusMin vis offs nGroups (col a) (col b)-    PlanMedian name -> groupedMedian vis offs nGroups (col name)-  where-    vis = valueIndices gdf-    offs = offsets gdf-    {- The low-cardinality DIRECT-INDEXED fast path: for a small dense domain the-    grouping layer's @rowToGroup@ already maps row -> group, so we scatter-    straight off it (no @valueIndices@ gather). 'directReduce' only admits-    order-independent reductions (so the merged parallel result is byte-identical-    to -N1); anything it rejects keeps the order-preserving group-range kernel. -}-    scatterColumn red name =-        let c = col name-            direct-                | nGroups <= directThreshold = directReduce red rtg nGroups c-                | otherwise = Nothing-         in case direct of-                Just out -> out-                Nothing -> case scatterReducePar red vis offs nGroups c of-                    Just out -> out-                    Nothing -> error "runPlan: scatterReducePar rejected a planned column"-    col name = case getColumn name (fullDataframe gdf) of-        Just c -> c-        Nothing -> error ("runPlan: planned column missing: " ++ T.unpack name)--{- | Run a recognised moment (Q9 regression) plan as one fused scatter over the-two base columns, returning each output name bound to its moment field. The six-sufficient statistics (count, Sx, Sy, Sxx, Syy, Sxy) come out of a single pass,-replacing the three derive passes and six independent scatters of the-per-expression path. Byte-identical to the sequential kernel at any @-N@.--}-runMomentPlan ::-    GroupedDataFrame -> Int -> MomentPlan -> Maybe [(T.Text, Column)]-runMomentPlan gdf nGroups mp = do-    ms <- momentScatterPar vis offs nGroups (col (mpColX mp)) (col (mpColY mp))-    pure-        [ (mpNName mp, mN ms)-        , (mpSxName mp, mSx ms)-        , (mpSyName mp, mSy ms)-        , (mpSxxName mp, mSxx ms)-        , (mpSyyName mp, mSyy ms)-        , (mpSxyName mp, mSxy ms)-        ]-  where-    vis = valueIndices gdf-    offs = offsets gdf-    col name = case getColumn name (fullDataframe gdf) of-        Just c -> c-        Nothing -> error ("runMomentPlan: planned column missing: " ++ T.unpack name)--{- | @max a - min b@ on the small @nGroups@ arrays. Preserves the Int element-type of the source columns (matching the interpreter), falling back to a Double-combine otherwise.--}-maxMinusMin ::-    VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -> Column-maxMinusMin vis offs nGroups ca cb =-    case (ca, cb) of-        ( UnboxedColumn Nothing (_ :: VU.Vector x)-            , UnboxedColumn Nothing (_ :: VU.Vector y)-            )-                | Just Refl <- testEquality (typeRep @x) (typeRep @Int)-                , Just Refl <- testEquality (typeRep @y) (typeRep @Int) ->-                    let mx = scatterExtremaInt RMax vis offs nGroups ca-                        mn = scatterExtremaInt RMin vis offs nGroups cb-                     in fromUnboxedVector (VU.zipWith (-) mx mn)-        _ ->-            let mx = scatterExtremaDbl RMax vis offs nGroups ca-                mn = scatterExtremaDbl RMin vis offs nGroups cb-             in fromUnboxedVector (VU.zipWith (-) mx mn)--scatterExtremaInt ::-    Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> VU.Vector Int-scatterExtremaInt red vis offs nGroups c = case scatterReducePar red vis offs nGroups c of-    Just (UnboxedColumn _ (v :: VU.Vector a))-        | Just Refl <- testEquality (typeRep @a) (typeRep @Int) -> v-    _ -> error "scatterExtremaInt"--scatterExtremaDbl ::-    Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> VU.Vector Double-scatterExtremaDbl red vis offs nGroups c =-    case scatterReducePar red vis offs nGroups c of-        Just (UnboxedColumn _ (v :: VU.Vector a))-            | Just Refl <- testEquality (typeRep @a) (typeRep @Double) -> v-            | Just Refl <- testEquality (typeRep @a) (typeRep @Int) -> VU.map fromIntegral v-        _ -> error "scatterExtremaDbl"------------------------------------------------------------------------------------ Parallel holistic median----------------------------------------------------------------------------------{- | Holistic per-group median over a single unboxed Int/Double column. The-@valueIndices@/@offsets@ layout already places each group's rows in a contiguous-run, so we copy each group's values into a scratch buffer at its own offset and-sort that slice in place — each group's slice is independent, so the per-group-sorts split across capabilities by group range with no merge. Empty groups never-occur, so the result is total.--}-groupedMedian :: VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column-groupedMedian vis offs nGroups c = case scatterColumnToDouble c of-    Nothing -> error "groupedMedian: non-numeric planned column"-    Just vals -> fromUnboxedVector (medianByGroup vis offs nGroups vals)--medianByGroup ::-    VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double-medianByGroup vis offs nGroups vals = unsafePerformIO $ do-    let !n = VU.length vis-    buf <- VUM.new (max 1 n)-    out <- VUM.new (max 1 nGroups)-    caps <- getNumCapabilities-    let !bounds = groupRangeBounds offs nGroups caps-    -- Each worker fills+sorts the buffer slices of its own group range, then-    -- writes that range's medians. Disjoint ranges => safe to parallelise.-    forEachRange bounds caps $ \gs ge ->-        let grp !g-                | g >= ge = pure ()-                | otherwise = do-                    let !s = VU.unsafeIndex offs g-                        !e = VU.unsafeIndex offs (g + 1)-                        !len = e - s-                        fill !pos-                            | pos >= e = pure ()-                            | otherwise = do-                                VUM.unsafeWrite buf pos (VU.unsafeIndex vals (VU.unsafeIndex vis pos))-                                fill (pos + 1)-                    fill s-                    let slice = VUM.unsafeSlice s len buf-                    VA.sort slice-                    let mid = s + len `div` 2-                    med <--                        if odd len-                            then VUM.unsafeRead buf mid-                            else do-                                hi <- VUM.unsafeRead buf mid-                                lo <- VUM.unsafeRead buf (mid - 1)-                                pure ((hi + lo) / 2)-                    VUM.unsafeWrite out g med-                    grp (g + 1)-         in grp gs-    VU.unsafeFreeze (VUM.unsafeSlice 0 nGroups out)-{-# NOINLINE medianByGroup #-}------------------------------------------------------------------------------------ Group-range partitioning (shared with the median path)----------------------------------------------------------------------------------{- | Split @[0, nGroups)@ into @caps@ contiguous group ranges balanced by row-count. Identical policy to 'DataFrame.Internal.AggKernelPar.groupRangeBounds'.--}-groupRangeBounds :: VU.Vector Int -> Int -> Int -> VU.Vector Int-groupRangeBounds offs nGroups caps = VU.create $ do-    b <- VUM.new (caps + 1)-    let !nRows = VU.unsafeIndex offs nGroups-        !per = max 1 ((nRows + caps - 1) `div` caps)-        adv !target !gg-            | gg >= nGroups = nGroups-            | VU.unsafeIndex offs gg >= target = gg-            | otherwise = adv target (gg + 1)-        go !w !prev-            | w >= caps = VUM.unsafeWrite b caps nGroups-            | otherwise = do-                let !target = min nRows (w * per)-                    !g = adv target prev-                VUM.unsafeWrite b w g-                go (w + 1) g-    VUM.unsafeWrite b 0 0-    go 1 0-    pure b--forEachRange :: VU.Vector Int -> Int -> (Int -> Int -> IO ()) -> IO ()-forEachRange bounds caps act-    | caps <= 1 = act (VU.unsafeIndex bounds 0) (VU.unsafeIndex bounds caps)-    | otherwise = do-        vars <- mapM spawn [0 .. caps - 1]-        results <- mapM takeMVar vars-        mapM_ (either (throwIO :: SomeException -> IO ()) pure) results-  where-    spawn w = do-        var <- newEmptyMVar-        let !s = VU.unsafeIndex bounds w-            !e = VU.unsafeIndex bounds (w + 1)-        _ <- forkIO (try (act s e) >>= putMVar var)-        pure var

@@ -14,29 +14,50 @@     changingPoints, ) where +import qualified Data.List as L+import qualified Data.Map.Strict as MS import qualified Data.Text as T import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU  import Control.Exception (throw) import DataFrame.Errors-import DataFrame.Internal.AggPlan (MomentPlan, planAgg, planMoments)+import DataFrame.Internal.Aggregation.Kernel.Fused (+    mkFusedAgg,+    mkGatherAgg,+    runFusedAggs,+    runGatherAggs,+ )+import DataFrame.Internal.Aggregation.Kernel.Scatter (streamGroupCap)+import DataFrame.Internal.Aggregation.Plan (+    AggPlan (..),+    MomentPlan,+    planAgg,+    planMoments,+ )+import DataFrame.Internal.Aggregation.Reduction (Reduction (..)) import DataFrame.Internal.Column (     Column (..),     TypedColumn (..),     atIndicesStable,+    atIndicesStableMulti,  ) import DataFrame.Internal.DataFrame (     DataFrame (..),     GroupedDataFrame (..),     columnNames,+    getColumn,     insertColumn,  ) import DataFrame.Internal.Expression import DataFrame.Internal.Grouping (buildRowToGroup, changingPoints, groupBy) import DataFrame.Internal.Interpreter-import DataFrame.Internal.RowHash (computeRowHashesIO)-import DataFrame.Operations.AggregateScatter (runMomentPlan, runPlan)+import DataFrame.Internal.Row.RowHash (computeRowHashesIO)+import DataFrame.Operations.Aggregation.Run (+    runMedianVarFused,+    runMomentPlan,+    runPlan,+ ) import DataFrame.Operations.Core import DataFrame.Operations.Subset import System.IO.Unsafe (unsafePerformIO)@@ -55,25 +76,143 @@  {- | Aggregate a grouped dataframe using the expressions given. All ungrouped columns will be dropped.++NOTE: this function deliberately never pattern-matches or strictly binds the+'Grouped' per-row fields (this module is compiled with @-XStrict@, whose strict+patterns and bindings would force them): on direct-grouped frames BOTH+'valueIndices' (the placement permutation) and 'rowToGroup' are deferred+thunks, and each aggregation path needs at most one of them — always passed as+un-forced argument expressions. Key columns materialize through 'groupRepRows'+(one representative row per group) instead of gathering+@valueIndices[offsets[g]]@. -} aggregate :: [NamedExpr] -> GroupedDataFrame -> DataFrame-aggregate aggs gdf@(Grouped df groupingColumns valIndices offs rowToGroupV) =+aggregate aggs gdf =     let+        df = fullDataframe gdf+        offs = offsets gdf++        {- Key columns materialize through ONE fused parallel gather over the+        representative rows ('atIndicesStableMulti'): the 1e8-group Q10 result+        was six sequential latency-bound random-gather passes as per-column+        'selectIndices'. Each result column is still identical to (and as+        deferred as) the per-column gather; the row-count field uses the eager+        @offsets@ so nothing here forces 'groupRepRows' early. -}         df' =-            selectIndices-                (VU.map (valIndices VU.!) (VU.init offs))-                (select groupingColumns df)+            let sub = select (groupedColumns gdf) df+             in sub+                    { columns =+                        V.fromList+                            ( atIndicesStableMulti+                                (groupRepRows gdf)+                                (V.toList (columns sub))+                            )+                    , dataframeDimensions = (nGroups, snd (dataframeDimensions sub))+                    }          !nGroups = VU.length offs - 1+        !nRows' = fst (dataframeDimensions df) +        {- Fused multi-reduction fast path (Q3/Q4/Q5-shaped aggregates): every+        recognised simple scatter reduction (sum/mean/count/min/max over a clean+        unboxed Int/Double column) in this aggregate runs in ONE pass instead of+        one full pass per expression. At or below 'streamGroupCap' groups that+        pass streams over (rowToGroup, columns) with per-worker accumulators —+        never touching the (possibly lazy) valueIndices; above it (necessarily a+        hash-path grouping, whose valueIndices is already eager) the accumulator+        arrays would thrash, so the pass gathers by disjoint group range instead+        (register accumulators, bit-identical to the unfused gather kernels,+        one traversal instead of one per expression). Only taken when at+        least two reductions fuse; non-fusable expressions (median, var/std,+        max-min, arbitrary DSL) keep their per-expression path below. -}+        {- This binding is strict (-XStrict), so it must stay empty-and-cheap+        whenever the moment path below already covers the aggregate — otherwise+        the pass would run redundantly before the moment result is consulted. -}+        fusedScatterCols :: MS.Map T.Text Column+        fusedScatterCols = case fusedMoments of+            Just _ -> MS.empty+            Nothing+                | nGroups <= streamGroupCap ->+                    let cands =+                            [ (name, fa)+                            | (name, ue) <- aggs+                            , Just (PlanScatter red cname) <- [planAgg gdf ue]+                            , Just c <- [getColumn cname df]+                            , Just fa <- [mkFusedAgg nGroups (rowToGroup gdf) red c]+                            ]+                     in if length cands >= 2+                            then+                                MS.fromList+                                    (zip (map fst cands) (runFusedAggs nRows' nGroups (map snd cands)))+                            else MS.empty+                | otherwise ->+                    let cands =+                            [ (name, ga)+                            | (name, ue) <- aggs+                            , Just (PlanScatter red cname) <- [planAgg gdf ue]+                            , Just c <- [getColumn cname df]+                            , Just ga <-+                                [mkGatherAgg nGroups (valueIndices gdf) offs red c]+                            ]+                     in {- Unlike the stream branch, a SINGLE candidate also+                        takes this path: the per-expression fallback is a+                        scatter over rowToGroup, which on a hash-path grouping+                        is now a deferred thunk — the gather kernel (documented+                        bit-identical to the unfused kernels) works off the+                        already-eager valueIndices instead and skips that whole+                        random-write pass. -}+                        if not (null cands)+                            then+                                MS.fromList+                                    ( zip+                                        (map fst cands)+                                        ( runGatherAggs+                                            (valueIndices gdf)+                                            offs+                                            nGroups+                                            (map snd cands)+                                        )+                                    )+                            else MS.empty++        {- Fused median + std/var over one column (the Q6 shape): both are+        holistic gathers over the same values, so one shared gather serves the+        Welford fold and the median selection ('runMedianVarFused',+        bit-identical to the separate kernels). Only built when a median and a+        std/var on the same column appear together; empty-and-cheap otherwise+        (same strictness caveat as 'fusedScatterCols'). -}+        medianVarCols :: MS.Map T.Text Column+        medianVarCols = case fusedMoments of+            Just _ -> MS.empty+            Nothing ->+                let plans = [(name, plan) | (name, ue) <- aggs, Just plan <- [planAgg gdf ue]]+                    medCols = L.nub [c | (_, PlanMedian c) <- plans]+                 in MS.fromList+                        [ kv+                        | cname <- medCols+                        , let stds = [nm | (nm, PlanScatter RStd c) <- plans, c == cname]+                        , let vars = [nm | (nm, PlanScatter RVar c) <- plans, c == cname]+                        , not (null stds && null vars)+                        , Just c <- [getColumn cname df]+                        , Just (medC, varC, stdC) <- [runMedianVarFused gdf nGroups c]+                        , kv <-+                            [(nm, medC) | (nm, PlanMedian c') <- plans, c' == cname]+                                ++ [(nm, stdC) | nm <- stds]+                                ++ [(nm, varC) | nm <- vars]+                        ]+         -- Fast path: a recognised reduction scatters in one unboxed pass.         -- Anything 'planAgg' rejects keeps the existing interpreter, so the         -- general typed + DSL aggregate API stays correct for arbitrary         -- expressions.         f ne@(name, uexpr) d =-            let value = case planAgg gdf uexpr of-                    Just plan -> runPlan gdf rowToGroupV nGroups plan-                    Nothing -> interpretNamed gdf ne+            let value = case MS.lookup name medianVarCols of+                    Just c -> c+                    Nothing -> case MS.lookup name fusedScatterCols of+                        Just c -> c+                        Nothing -> case planAgg gdf uexpr of+                            Just plan -> runPlan gdf (rowToGroup gdf) nGroups plan+                            Nothing -> interpretNamed gdf ne              in insertColumn name value d          -- Fused fast path: the Q9 regression family (count + five moment sums@@ -107,4 +246,6 @@ distinct :: DataFrame -> DataFrame distinct df = selectIndices (VU.map (indices VU.!) (VU.init os)) df   where-    (Grouped _ _ indices os _rtg) = groupBy (columnNames df) df+    -- The trailing field stays a wildcard: under -XStrict a named pattern+    -- variable would force the (possibly deferred) rowToGroup thunk.+    (Grouped _ _ indices os _) = groupBy (columnNames df) df

@@ -0,0 +1,389 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | Execute a recognised aggregation plan ('AggPlan'), producing one result+column (length @nGroups@, canonical group order).++This is the dispatch layer: it owns the policy of WHICH kernel runs — the dense+direct-indexed one ('DataFrame.Internal.Aggregation.Kernel.Dense') when the+group domain is small enough, otherwise the group-range scatter+('DataFrame.Internal.Aggregation.Kernel.Scatter') — and handles the compound+@max - min@ combine and the holistic grouped median itself. The kernels carry no+policy of their own. A plan only reaches here once 'planAgg' verified the value+columns are clean unboxed Int/Double, so the @error@ branches are unreachable.++Every reduction takes the Round-5 grouping layout @(valueIndices, offsets)@ so+the parallel kernel can split the group-id range across capabilities with no+cross-worker merge. Each group's rows stay in original-row order within one+worker's range, so results are byte-identical to the sequential path at any @-N@.+(Two exceptions, both deterministic at a fixed @-N@: the direct streaming+Double sum/mean above the small-group cutoff, whose chunked partials change the+float summation order, and the direct var/std, which finalize from+(count, sum, sumsq) partials rather than the gather kernel's row-order Welford+recurrence; see 'DataFrame.Internal.AggKernelDirect'.)+-}+module DataFrame.Operations.Aggregation.Run (+    runPlan,+    runMomentPlan,+    runMedianVarFused,+) where++import qualified Data.Text as T+import qualified Data.Vector.Algorithms.Intro as VA+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM++import Control.Concurrent (getNumCapabilities)+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))+import DataFrame.Internal.Aggregation.Kernel.Dense (+    denseMaxMinusMin,+    denseReduce,+ )+import DataFrame.Internal.Aggregation.Kernel.Moments (+    Moments (..),+    momentScatterPar,+    momentStreamPar,+ )+import DataFrame.Internal.Aggregation.Kernel.Scatter (+    maxMinusMinScatterPar,+    scatterReducePar,+    streamGroupCap,+ )+import DataFrame.Internal.Aggregation.Plan (+    AggPlan (..),+    MomentPlan (..),+ )+import DataFrame.Internal.Aggregation.Reduction (+    Reduction (..),+    cleanDoubleVector,+ )+import DataFrame.Internal.Column (Column (..), fromUnboxedVector)+import DataFrame.Internal.Control.Concurrent (parallelBounds_)+import DataFrame.Internal.DataFrame (GroupedDataFrame (..), getColumn)+import System.IO.Unsafe (unsafePerformIO)+import Type.Reflection (typeRep)++{- | Group-domain size at or below which the dense direct-indexed kernel is+chosen; wider domains go to the group-range scatter kernel. A @2^18@-slot+accumulator is replicated per worker there, which is what bounds this. Dispatch+policy, so it lives with the dispatcher rather than inside the kernel.+-}+denseThreshold :: Int+denseThreshold = 262144++runPlan :: GroupedDataFrame -> VU.Vector Int -> Int -> AggPlan -> Column+runPlan gdf rtg nGroups plan = case plan of+    PlanScatter red name -> scatterColumn red name+    PlanMaxMinusMin a b ->+        {- min/max are order-independent, so both fused single-pass kernels+        (the direct streaming one up to 'streamGroupCap', the group-range+        gather one above it) are exactly the two gather extrema they replace;+        anything they reject (mixed/unclean columns, small inputs) keeps the+        two-pass gather path. The streaming cap extends past 'denseThreshold'+        for the same reason as the fused multi-reduction pass: on a+        direct-grouped frame it works off the eager @rowToGroup@ and skips the+        deferred @valueIndices@ placement entirely (measured at 1e6 groups /+        1e8 rows on -N16: stream 1.1s against placement 0.7s + gather 0.7s). -}+        let ca = col a+            cb = col b+            direct+                | nGroups <= streamGroupCap = denseMaxMinusMin rtg nGroups ca cb+                | otherwise = maxMinusMinScatterPar vis offs nGroups ca cb+         in case direct of+                Just out -> out+                Nothing -> maxMinusMin vis offs nGroups ca cb+    PlanMedian name -> groupedMedian vis offs nGroups (col name)+  where+    vis = valueIndices gdf+    offs = offsets gdf+    {- The low-cardinality DENSE fast path: for a small dense domain the grouping+    layer's @rowToGroup@ already maps row -> group, so we scatter straight off it+    (no @valueIndices@ gather). 'denseReduce' admits the order-independent+    reductions (exact partial merge, byte-identical to -N1) plus the streaming+    Double sum/mean/var/std variants (byte-identical sequential row order at+    small group counts; deterministic chunked partials for the large-domain+    sum/mean — see "DataFrame.Internal.Aggregation.Kernel.Dense"); anything it+    rejects keeps the order-preserving group-range kernel. -}+    scatterColumn red name =+        let c = col name+            dense+                | nGroups <= denseThreshold = denseReduce red rtg nGroups c+                {- Top-2 selection merges exactly (a multiset selection, no+                float adds until finalize), so like the fused passes it streams+                off @rowToGroup@ up to 'streamGroupCap': on a direct-grouped+                frame that skips the deferred @valueIndices@ placement, which+                costs more than the accumulator cache misses it saves+                (measured at 1e6 groups / 1e8 rows on -N16: stream 1.0s+                against placement 0.7s + gather 0.4s). RTop2Snd shares the+                same accumulator machinery and merge-exactness. -}+                | red == RTop2Sum || red == RTop2Snd+                , nGroups <= streamGroupCap =+                    denseReduce red rtg nGroups c+                | otherwise = Nothing+         in case dense of+                Just out -> out+                Nothing -> case scatterReducePar red vis offs nGroups c of+                    Just out -> out+                    Nothing -> error "runPlan: scatterReducePar rejected a planned column"+    col name = case getColumn name (fullDataframe gdf) of+        Just c -> c+        Nothing -> error ("runPlan: planned column missing: " ++ T.unpack name)++{- | Run a recognised moment (Q9 regression) plan as one fused scatter over the+two base columns, returning each output name bound to its moment field. The six+sufficient statistics (count, Sx, Sy, Sxx, Syy, Sxy) come out of a single pass,+replacing the three derive passes and six independent scatters of the+per-expression path. The streaming kernel's count is exact; its five Double+sums accumulate per worker chunk in original row order and merge in fixed+worker order — deterministic at a fixed @-N@, float summation order chunk-major+rather than per-group (see 'momentStreamPar'). The gather fallback remains+byte-identical to the sequential kernel at any @-N@.+-}+runMomentPlan ::+    GroupedDataFrame -> Int -> MomentPlan -> Maybe [(T.Text, Column)]+runMomentPlan gdf nGroups mp = do+    let cx = col (mpColX mp)+        cy = col (mpColY mp)+        {- Preferred: the streaming kernel — one fused pass over rowToGroup and+        the TYPED base columns (no sequential Int->Double materialization, no+        valueIndices gather, so a direct-grouped frame never runs its placement+        pass). Falls back to the gather kernel above 'streamGroupCap' or on+        unclean columns. -}+        streamed = momentStreamPar (rowToGroup gdf) nGroups cx cy+    ms <- case streamed of+        Just m -> Just m+        Nothing -> momentScatterPar vis offs nGroups cx cy+    pure+        [ (mpNName mp, mN ms)+        , (mpSxName mp, mSx ms)+        , (mpSyName mp, mSy ms)+        , (mpSxxName mp, mSxx ms)+        , (mpSyyName mp, mSyy ms)+        , (mpSxyName mp, mSxy ms)+        ]+  where+    vis = valueIndices gdf+    offs = offsets gdf+    col name = case getColumn name (fullDataframe gdf) of+        Just c -> c+        Nothing -> error ("runMomentPlan: planned column missing: " ++ T.unpack name)++{- | @max a - min b@ on the small @nGroups@ arrays. Preserves the Int element+type of the source columns (matching the interpreter), falling back to a Double+combine otherwise.+-}+maxMinusMin ::+    VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -> Column+maxMinusMin vis offs nGroups ca cb =+    case (ca, cb) of+        ( UnboxedColumn Nothing (_ :: VU.Vector x)+            , UnboxedColumn Nothing (_ :: VU.Vector y)+            )+                | Just Refl <- testEquality (typeRep @x) (typeRep @Int)+                , Just Refl <- testEquality (typeRep @y) (typeRep @Int) ->+                    let mx = scatterExtremaInt RMax vis offs nGroups ca+                        mn = scatterExtremaInt RMin vis offs nGroups cb+                     in fromUnboxedVector (VU.zipWith (-) mx mn)+        _ ->+            let mx = scatterExtremaDbl RMax vis offs nGroups ca+                mn = scatterExtremaDbl RMin vis offs nGroups cb+             in fromUnboxedVector (VU.zipWith (-) mx mn)++scatterExtremaInt ::+    Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> VU.Vector Int+scatterExtremaInt red vis offs nGroups c = case scatterReducePar red vis offs nGroups c of+    Just (UnboxedColumn _ (v :: VU.Vector a))+        | Just Refl <- testEquality (typeRep @a) (typeRep @Int) -> v+    _ -> error "scatterExtremaInt"++scatterExtremaDbl ::+    Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> VU.Vector Double+scatterExtremaDbl red vis offs nGroups c =+    case scatterReducePar red vis offs nGroups c of+        Just (UnboxedColumn _ (v :: VU.Vector a))+            | Just Refl <- testEquality (typeRep @a) (typeRep @Double) -> v+            | Just Refl <- testEquality (typeRep @a) (typeRep @Int) -> VU.map fromIntegral v+        _ -> error "scatterExtremaDbl"++-------------------------------------------------------------------------------+-- Fused holistic median + var/std over one shared gather+-------------------------------------------------------------------------------++{- | Fused grouped median and variance family over the SAME column: one gather+into the shared scratch buffer serves both. Returns+@(median, variance, stddev)@ columns, or 'Nothing' on a non-numeric column+(the caller keeps the separate per-expression kernels).++The Welford fold runs over each gathered slice in ascending original-row order+— exactly the recurrence, order and finalize of the var/std kernels+('DataFrame.Internal.AggKernelPar.varPar' and the sequential @varScatter@,+which agree bit-for-bit) — BEFORE the in-place median selection permutes the+slice, and the selection then proceeds exactly as 'groupedMedian'. Both+outputs are therefore bit-identical to the unfused paths; the second full+gather pass is what the fusion saves (measured ~35% off the median+sd pair at+1e4 groups / 1e8 rows on -N16).+-}+runMedianVarFused ::+    GroupedDataFrame -> Int -> Column -> Maybe (Column, Column, Column)+runMedianVarFused gdf nGroups c = do+    vals <- cleanDoubleVector c+    let (med, var) =+            medianVarByGroup (valueIndices gdf) (offsets gdf) nGroups vals+    pure+        ( fromUnboxedVector med+        , fromUnboxedVector var+        , fromUnboxedVector (VU.map sqrt var)+        )++medianVarByGroup ::+    VU.Vector Int ->+    VU.Vector Int ->+    Int ->+    VU.Vector Double ->+    (VU.Vector Double, VU.Vector Double)+medianVarByGroup vis offs nGroups vals = unsafePerformIO $ do+    let !n = VU.length vis+    buf <- VUM.new (max 1 n)+    medOut <- VUM.new (max 1 nGroups)+    varOut <- VUM.new (max 1 nGroups)+    caps <- getNumCapabilities+    let !bounds = groupRangeBounds offs nGroups caps+    parallelBounds_ caps bounds $ \gs ge ->+        let grp !g+                | g >= ge = pure ()+                | otherwise = do+                    let !s = VU.unsafeIndex offs g+                        !e = VU.unsafeIndex offs (g + 1)+                        !len = e - s+                        fill !pos+                            | pos >= e = pure ()+                            | otherwise = do+                                VUM.unsafeWrite buf pos (VU.unsafeIndex vals (VU.unsafeIndex vis pos))+                                fill (pos + 1)+                    fill s+                    -- Welford over the gathered slice (still in row order).+                    let welford !pos !c !mu !mm+                            | pos >= e =+                                pure (if c < 2 then 0 else mm / fromIntegral (c - 1))+                            | otherwise = do+                                x <- VUM.unsafeRead buf pos+                                let !c' = c + 1+                                    !delta = x - mu+                                    !mu' = mu + delta / fromIntegral c'+                                    !mm' = mm + delta * (x - mu')+                                welford (pos + 1) c' mu' mm'+                    var <- welford s (0 :: Int) 0 0+                    VUM.unsafeWrite varOut g var+                    -- Median selection, as in 'medianByGroup' (permutes the slice).+                    let slice = VUM.unsafeSlice s len buf+                        !mid = len `div` 2+                    VA.select slice (mid + 1)+                    let scan !i !hi !lo+                            | i > mid = pure (hi, lo)+                            | otherwise = do+                                x <- VUM.unsafeRead slice i+                                if x > hi+                                    then scan (i + 1) x hi+                                    else scan (i + 1) hi (max lo x)+                    (hi, lo) <- scan 0 (negate (1 / 0)) (negate (1 / 0))+                    let med = if odd len then hi else (hi + lo) / 2+                    VUM.unsafeWrite medOut g med+                    grp (g + 1)+         in grp gs+    med <- VU.unsafeFreeze (VUM.unsafeSlice 0 nGroups medOut)+    var <- VU.unsafeFreeze (VUM.unsafeSlice 0 nGroups varOut)+    pure (med, var)+{-# NOINLINE medianVarByGroup #-}++-------------------------------------------------------------------------------+-- Parallel holistic median+-------------------------------------------------------------------------------++{- | Holistic per-group median over a single unboxed Int/Double column. The+@valueIndices@/@offsets@ layout already places each group's rows in a contiguous+run, so we copy each group's values into a scratch buffer at its own offset and+select the median-rank order statistics in that slice in place (O(len) per+group rather than the O(len log len) full sort) — each group's slice is+independent, so the per-group selections split across capabilities by group+range with no merge. Order statistics are value-determined, so the result is+identical to the sorting variant. Empty groups never occur, so the result is+total.+-}+groupedMedian :: VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column+groupedMedian vis offs nGroups c = case cleanDoubleVector c of+    Nothing -> error "groupedMedian: non-numeric planned column"+    Just vals -> fromUnboxedVector (medianByGroup vis offs nGroups vals)++medianByGroup ::+    VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Double -> VU.Vector Double+medianByGroup vis offs nGroups vals = unsafePerformIO $ do+    let !n = VU.length vis+    buf <- VUM.new (max 1 n)+    out <- VUM.new (max 1 nGroups)+    caps <- getNumCapabilities+    let !bounds = groupRangeBounds offs nGroups caps+    -- Each worker fills+sorts the buffer slices of its own group range, then+    -- writes that range's medians. Disjoint ranges => safe to parallelise.+    parallelBounds_ caps bounds $ \gs ge ->+        let grp !g+                | g >= ge = pure ()+                | otherwise = do+                    let !s = VU.unsafeIndex offs g+                        !e = VU.unsafeIndex offs (g + 1)+                        !len = e - s+                        fill !pos+                            | pos >= e = pure ()+                            | otherwise = do+                                VUM.unsafeWrite buf pos (VU.unsafeIndex vals (VU.unsafeIndex vis pos))+                                fill (pos + 1)+                    fill s+                    let slice = VUM.unsafeSlice s len buf+                        !mid = len `div` 2+                    {- Move the least mid+1 values to the front (in no+                    particular order); the two largest of those are the order+                    statistics at sorted positions mid and mid-1. -}+                    VA.select slice (mid + 1)+                    let scan !i !hi !lo+                            | i > mid = pure (hi, lo)+                            | otherwise = do+                                x <- VUM.unsafeRead slice i+                                if x > hi+                                    then scan (i + 1) x hi+                                    else scan (i + 1) hi (max lo x)+                    (hi, lo) <- scan 0 (negate (1 / 0)) (negate (1 / 0))+                    let med = if odd len then hi else (hi + lo) / 2+                    VUM.unsafeWrite out g med+                    grp (g + 1)+         in grp gs+    VU.unsafeFreeze (VUM.unsafeSlice 0 nGroups out)+{-# NOINLINE medianByGroup #-}++-------------------------------------------------------------------------------+-- Group-range partitioning (shared with the median path)+-------------------------------------------------------------------------------++{- | Split @[0, nGroups)@ into @caps@ contiguous group ranges balanced by row+count. Identical policy to 'DataFrame.Internal.Aggregation.Kernel.Scatter.groupRangeBounds'.+-}+groupRangeBounds :: VU.Vector Int -> Int -> Int -> VU.Vector Int+groupRangeBounds offs nGroups caps = VU.create $ do+    b <- VUM.new (caps + 1)+    let !nRows = VU.unsafeIndex offs nGroups+        !per = max 1 ((nRows + caps - 1) `div` caps)+        adv !target !gg+            | gg >= nGroups = nGroups+            | VU.unsafeIndex offs gg >= target = gg+            | otherwise = adv target (gg + 1)+        go !w !prev+            | w >= caps = VUM.unsafeWrite b caps nGroups+            | otherwise = do+                let !target = min nRows (w * per)+                    !g = adv target prev+                VUM.unsafeWrite b w g+                go (w + 1) g+    VUM.unsafeWrite b 0 0+    go 1 0+    pure b

@@ -1,13 +1,6 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE ScopedTypeVariables #-} -{- | Single-pass sampled inference lattice (Round-2 S2). One walk over-the sampled cells maintains a candidate mask {Bool, Int, Double, Date}-cleared by attempting the WS-B byte parsers; the assumption is the-highest-priority surviving candidate. Shared by every reader (audit T1),-together with the Int -> Double prefix promotion that replaces the-full-column retry chain (audit T2).--} module DataFrame.Operations.Inference (     DateFormat,     ParsingAssumption (..),@@ -57,29 +50,17 @@         dateFormat         (T.unpack s) -{- | The default @%Y-%m-%d@ format takes the WS-B byte-level fast path;-custom formats keep the reference 'readByteStringDate' parser.--} byteStringDateParser :: DateFormat -> BS.ByteString -> Maybe Day byteStringDateParser "%Y-%m-%d" = parseDateField byteStringDateParser fmt = readByteStringDate fmt {-# INLINE byteStringDateParser #-} -{- | 'DataFrame.Internal.Parsing.readInt' that rejects overflow instead of-wrapping. Fields of <= 18 chars cannot overflow and keep the Text-level-parse; longer (rare) fields take the exact byte-level parser, so an-overflowing cell demotes\/promotes instead of silently wrapping.--} readIntStrict :: T.Text -> Maybe Int readIntStrict t     | T.length t <= 18 = readInt t     | otherwise = parseIntField (TE.encodeUtf8 t) {-# INLINE readIntStrict #-} -{- | Candidate-mask priority, reproducing the documented fallback order:-an all-null sample makes no assumption; Int wins only when the Double-mask agrees (so mixed Int\/Double samples classify as Double).--} pickAssumption ::     Bool -> Bool -> Bool -> Bool -> Bool -> ParsingAssumption pickAssumption seen b i d dt@@ -90,12 +71,6 @@     | dt = DateAssumption     | otherwise = TextAssumption -{- | Classify a sample of decoded 'T.Text' cells ('Nothing' = null).-Bool\/Int\/Double candidates are tested with the WS-B byte parsers on-the UTF-8 bytes (strip-tolerant, overflow-rejecting); the Date-candidate keeps 'parseTimeOpt' so custom formats behave exactly as-before. The walk exits early once every candidate is cleared.--} makeParsingAssumption ::     DateFormat -> V.Vector (Maybe T.Text) -> ParsingAssumption makeParsingAssumption dfmt cells = go 0 False True True True True@@ -137,13 +112,6 @@                     (d && isJust (parseDoubleField bs))                     (dt && isJust (dateP bs)) -{- | Fused Int pass with in-place promotion (audit T2): on the first-non-null cell that fails Int but parses as Double, the built Int prefix-is converted by a vector map and the pass continues as Double over the-retained raw cells. @Nothing@ = some cell parses as neither (the caller-demotes the column to Text). Null slots hold sentinels (0 \/ 0.0)-guarded by the bitmap, exactly like the unpromoted passes.--} promoteIntColumn ::     forall src.     (Int -> src -> Bool) ->

@@ -40,7 +40,7 @@ import Control.Exception (throw) import Control.Monad (when) import Control.Monad.ST (ST, runST)-import Data.Bits ((.&.))+import Data.Bits (popCount, unsafeShiftL, unsafeShiftR, (.&.), (.|.)) import qualified Data.Map.Strict as M import Data.Maybe (fromMaybe) import Data.STRef (newSTRef, readSTRef, writeSTRef)@@ -51,15 +51,28 @@ import qualified Data.Vector.Algorithms.Merge as VA import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.Word (Word64) import DataFrame.Errors (     DataFrameException (ColumnsNotFoundException),  )-import DataFrame.Internal.Column as D+import DataFrame.Internal.Algorithms.Sort.Radix.Parallel (parSortByHash)+import DataFrame.Internal.Column as D (+    Column (BoxedColumn, UnboxedColumn),+    atIndicesStable,+    columnTypeString,+    fromUnboxedVector,+    fromVector,+    gatherWithSentinel,+    isPackedText,+    materializePacked,+    mkMergedColumns,+ )+import DataFrame.Internal.Column.Bitmap (bitmapTestBit) import DataFrame.Internal.DataFrame as D-import DataFrame.Internal.ParRadixSort (parSortByHash) import DataFrame.Operations.Aggregation as D import DataFrame.Operations.Core as D-import DataFrame.Operations.JoinPar (+import DataFrame.Operations.Join.Parallel (+    ProbeTable (..),     parInnerProbe,     parLeftProbe,     shouldParallelizeJoin,@@ -102,28 +115,64 @@ data CompactIndex = CompactIndex     { ciSortedIndices :: {-# UNPACK #-} !(VU.Vector Int)     , ciKeys :: {-# UNPACK #-} !(VU.Vector Int)-    , ciStarts :: {-# UNPACK #-} !(VU.Vector Int)-    , ciLens :: {-# UNPACK #-} !(VU.Vector Int)+    , ciRuns :: {-# UNPACK #-} !(VU.Vector Int)+    {- ^ @(start, len)@ of each run packed 32/32 into one 'Int'; @-1@ = empty+    slot. One vector instead of separate starts\/lens halves the table's+    memory (a 1e8-row build side is a 2^28-slot table: 2.1GB instead of+    4.3GB) and saves a random cache-line read per lookup hit.+    -}     , ciMask :: {-# UNPACK #-} !Int     } +-- | Pack a run's @(start, len)@ 32\/32 into one non-negative 'Int'.+ciPackRun :: Int -> Int -> Int+ciPackRun !start !len = (start `unsafeShiftL` 32) .|. len+{-# INLINE ciPackRun #-}++-- | Start field of a packed run.+ciRunStart :: Int -> Int+ciRunStart !w = w `unsafeShiftR` 32+{-# INLINE ciRunStart #-}++-- | Length field of a packed run.+ciRunLen :: Int -> Int+ciRunLen !w = w .&. 0xFFFF_FFFF+{-# INLINE ciRunLen #-}++{- | 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 (text keys cluster catastrophically: contiguous-pileup chains+in the hundreds), so the table must never index by @h .&. mask@ directly.+@shift@ is @64 - log2 cap@.+-}+ciSlot :: Int -> Int -> Int+ciSlot !shift !h =+    fromIntegral+        ((fromIntegral h * (0x9E37_79B9_7F4A_7C15 :: Word64)) `unsafeShiftR` shift)+{-# INLINE ciSlot #-}++-- | @64 - log2 cap@ for a table with slot mask @mask@ (@cap@ a power of two).+ciShiftFor :: Int -> Int+ciShiftFor !mask = 64 - popCount mask+{-# INLINE ciShiftFor #-}+ {- | Look up a hash in the open-addressing table. Returns @(start, len)@ of the matching run, or @(-1, 0)@ on a miss. -} ciLookup :: CompactIndex -> Int -> (Int, Int)-ciLookup ci !h = go (h .&. mask)+ciLookup ci !h = go (ciSlot shift h)   where     !mask = ciMask ci+    !shift = ciShiftFor mask     !keys = ciKeys ci-    !starts = ciStarts ci-    !lens = ciLens ci+    !runs = ciRuns ci     go !slot =-        let !s = starts `VU.unsafeIndex` slot-         in if s < 0+        let !w = runs `VU.unsafeIndex` slot+         in if w < 0                 then (-1, 0)                 else                     if keys `VU.unsafeIndex` slot == h-                        then (s, lens `VU.unsafeIndex` slot)+                        then (ciRunStart w, ciRunLen w)                         else go ((slot + 1) .&. mask) {-# INLINE ciLookup #-} @@ -142,36 +191,38 @@ sized for the worst case (every row distinct) so building never resizes. -} buildCompactIndex :: VU.Vector Int -> CompactIndex+buildCompactIndex hashes+    | VU.length hashes > 0x7FFF_FFFF =+        error+            "buildCompactIndex: build side exceeds 2^31 rows (packed run fields are 32-bit)" buildCompactIndex hashes =     let n = VU.length hashes         (sortedHashes, sortedIndices) = parSortByHash n hashes         !cap = nextPow2Above (2 * n)         !mask = cap - 1-        (keys, starts, lens) = runST $ do+        !shift = ciShiftFor mask+        (keys, runs) = runST $ do             mKeys <- VUM.unsafeNew cap-            mStarts <- VUM.replicate cap (-1)-            mLens <- VUM.unsafeNew cap+            mRuns <- VUM.replicate cap (-1)             let insert !i                     | i >= n = return ()                     | otherwise = do                         let !h = sortedHashes `VU.unsafeIndex` i                             !end = findGroupEnd sortedHashes h (i + 1) n-                        probe h (h .&. mask) i (end - i)+                        probe h (ciSlot shift h) i (end - i)                         insert end                 probe !h !slot !start !len = do-                    s <- VUM.unsafeRead mStarts slot-                    if s < 0+                    w <- VUM.unsafeRead mRuns slot+                    if w < 0                         then do                             VUM.unsafeWrite mKeys slot h-                            VUM.unsafeWrite mStarts slot start-                            VUM.unsafeWrite mLens slot len+                            VUM.unsafeWrite mRuns slot (ciPackRun start len)                         else probe h ((slot + 1) .&. mask) start len             insert 0-            (,,)+            (,)                 <$> VU.unsafeFreeze mKeys-                <*> VU.unsafeFreeze mStarts-                <*> VU.unsafeFreeze mLens-     in CompactIndex sortedIndices keys starts lens mask+                <*> VU.unsafeFreeze mRuns+     in CompactIndex sortedIndices keys runs mask  -- | Find the end of a contiguous run of equal values starting at @j@. findGroupEnd :: VU.Vector Int -> Int -> Int -> Int -> Int@@ -315,10 +366,22 @@ parInnerKernel probeHashes buildHashes =     let !ci = buildCompactIndex buildHashes         (pf, bf) =-            unsafePerformIO (parInnerProbe (ciSortedIndices ci) (ciLookup ci) probeHashes)+            unsafePerformIO (parInnerProbe (ciProbeTable ci) probeHashes)      in (pf, bf) {-# NOINLINE parInnerKernel #-} +{- | The raw table fields of a 'CompactIndex', in the closure-free shape the+parallel probe kernels consume.+-}+ciProbeTable :: CompactIndex -> ProbeTable+ciProbeTable ci =+    ProbeTable+        { ptSorted = ciSortedIndices ci+        , ptKeys = ciKeys ci+        , ptRuns = ciRuns ci+        , ptMask = ciMask ci+        }+ -- | Compute hashes for the given key column names in a DataFrame. buildHashColumn :: [T.Text] -> DataFrame -> VU.Vector Int buildHashColumn keys df =@@ -602,7 +665,7 @@                             then                                 insertIfPresent                                     name-                                    (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name)+                                    (D.mkMergedColumns <$> getExpandedLeft name <*> getExpandedRight name)                                     df                             else                                 insertIfPresent name (getExpandedRight name) df@@ -681,7 +744,7 @@     VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) parLeftKernel leftHashes rightHashes =     let !ci = buildCompactIndex rightHashes-     in unsafePerformIO (parLeftProbe (ciSortedIndices ci) (ciLookup ci) leftHashes)+     in unsafePerformIO (parLeftProbe (ciProbeTable ci) leftHashes) {-# NOINLINE parLeftKernel #-}  {- | Hash-based left join kernel, returning @(leftExpandedIndices,@@ -900,7 +963,7 @@                             then                                 insertIfPresent                                     name-                                    (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name)+                                    (D.mkMergedColumns <$> getExpandedLeft name <*> getExpandedRight name)                                     df                             else insertIfPresent name (getExpandedRight name) df             )@@ -970,34 +1033,32 @@         leftSI = ciSortedIndices leftCI         rightSI = ciSortedIndices rightCI         leftKeys = ciKeys leftCI-        leftStarts = ciStarts leftCI-        leftLens = ciLens leftCI+        leftRuns = ciRuns leftCI         rightKeys = ciKeys rightCI-        rightStarts = ciStarts rightCI-        rightLens = ciLens rightCI-        !leftCap = VU.length leftStarts-        !rightCap = VU.length rightStarts+        rightRuns = ciRuns rightCI+        !leftCap = VU.length leftRuns+        !rightCap = VU.length rightRuns      let countLeft !slot !acc             | slot >= leftCap = acc             | otherwise =-                let !lStart = leftStarts `VU.unsafeIndex` slot-                 in if lStart < 0+                let !lw = leftRuns `VU.unsafeIndex` slot+                 in if lw < 0                         then countLeft (slot + 1) acc                         else                             let !h = leftKeys `VU.unsafeIndex` slot-                                !ll = leftLens `VU.unsafeIndex` slot+                                !ll = ciRunLen lw                                 (!rs, !rl) = ciLookup rightCI h                              in countLeft (slot + 1) (acc + if rs < 0 then ll else ll * rl)         countRightOnly !slot !acc             | slot >= rightCap = acc             | otherwise =-                let !rStart = rightStarts `VU.unsafeIndex` slot-                 in if rStart < 0+                let !rw = rightRuns `VU.unsafeIndex` slot+                 in if rw < 0                         then countRightOnly (slot + 1) acc                         else                             let !h = rightKeys `VU.unsafeIndex` slot-                                !rl = rightLens `VU.unsafeIndex` slot+                                !rl = ciRunLen rw                                 (!ls, _) = ciLookup leftCI h                              in countRightOnly (slot + 1) (acc + if ls < 0 then rl else 0)         !leftPlusMatched = countLeft 0 0@@ -1011,12 +1072,13 @@     let fillLeft !slot             | slot >= leftCap = return ()             | otherwise = do-                let !lStart = leftStarts `VU.unsafeIndex` slot-                if lStart < 0+                let !lw = leftRuns `VU.unsafeIndex` slot+                if lw < 0                     then fillLeft (slot + 1)                     else do                         let !h = leftKeys `VU.unsafeIndex` slot-                            !lLen = leftLens `VU.unsafeIndex` slot+                            !lStart = ciRunStart lw+                            !lLen = ciRunLen lw                             (!rStart, !rLen) = ciLookup rightCI h                         !p <- readSTRef posRef                         if rStart < 0@@ -1047,12 +1109,13 @@     let fillRightOnly !slot             | slot >= rightCap = return ()             | otherwise = do-                let !rStart = rightStarts `VU.unsafeIndex` slot-                if rStart < 0+                let !rw = rightRuns `VU.unsafeIndex` slot+                if rw < 0                     then fillRightOnly (slot + 1)                     else do                         let !h = rightKeys `VU.unsafeIndex` slot-                            !rLen = rightLens `VU.unsafeIndex` slot+                            !rStart = ciRunStart rw+                            !rLen = ciRunLen rw                             (!ls, _) = ciLookup leftCI h                         if ls >= 0                             then fillRightOnly (slot + 1)@@ -1232,7 +1295,7 @@                             then                                 insertIfPresent                                     name-                                    (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name)+                                    (D.mkMergedColumns <$> getExpandedLeft name <*> getExpandedRight name)                                     df                             else insertIfPresent name (getExpandedRight name) df             )

@@ -0,0 +1,222 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-}++{- |+Parallel chunked-probe join kernels. The build side is indexed once into a+shared, read-only 'CompactIndex' (open-addressing, from+"DataFrame.Operations.Join"); the probe side is split into @caps@ /contiguous/+row ranges and probed in parallel by 'forkFinally' workers (no sparks). Each worker+makes two passes over its range — a count pass to size its slice, then a fill+pass — and writes into the single shared output buffers at a precomputed+prefix-sum offset. Because ranges are contiguous and laid out in range order,+the produced @(probeIxs, buildIxs)@ vectors are /bit-for-bit identical/ to the+sequential 'hashInnerKernel' \/ 'hashLeftKernel': probe rows appear in original+order and, within a probe row, build matches in @ciSortedIndices@ order.+-}+module DataFrame.Operations.Join.Parallel (+    ProbeTable (..),+    parInnerProbe,+    parLeftProbe,+    shouldParallelizeJoin,+    shouldParallelizeSmallBuildProbe,+    parJoinThreshold,+    parBuildThreshold,+    parProbeThreshold,+) where++import Control.Concurrent (getNumCapabilities)+import Data.Bits (popCount, unsafeShiftR, (.&.))+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.Word (Word64)+import DataFrame.Internal.Control.Concurrent (capabilities, forkJoin_)++{- | Below this many probe rows the fork/coordination overhead is not worth it;+the caller uses its sequential 'ST' kernel instead.+-}+parJoinThreshold :: Int+parJoinThreshold = 200000++{- | Below this many build rows the shared 'CompactIndex' is small and hot, so+the sequential hash probe is already memory-bound-fast and the fork overhead+loses (measured: a 1e4-row Text-key build probed by 1e7 rows is /slower/ in+parallel). Parallelism only pays once the build index is large enough to spill+cache — exactly the regime where sort-merge used to be chosen.+-}+parBuildThreshold :: Int+parBuildThreshold = 500000++{- | Whether a join should take the parallel probe path: more than one+capability, a probe side of at least 'parJoinThreshold' rows, and a build side+of at least 'parBuildThreshold' rows (a small/hot index is faster probed+sequentially).+-}+shouldParallelizeJoin :: Int -> Int -> Bool+shouldParallelizeJoin probeRows buildRows =+    probeRows >= parJoinThreshold+        && buildRows >= parBuildThreshold+        && capabilities > 1+{-# NOINLINE shouldParallelizeJoin #-}++{- | Above this many probe rows the probe-side row hashing and table lookups+dominate the join, so partitioning the probe across cores wins even when the+build side is small and cache-resident (the regime 'shouldParallelizeJoin'+deliberately leaves sequential). Sized at 1e6: below it the per-question gain is+swamped by 'forkFinally'/coordination overhead, so small/medium-inner joins stay+sequential (measured). This is the small-build large-probe lever closing the+medium-factor 1e7 join (1e7 probe x ~1e4 build).+-}+parProbeThreshold :: Int+parProbeThreshold = 1000000++{- | Whether a /small-build/ join (build below 'parBuildThreshold', so radix+partitioning / sort-merge is not used) should take the parallel probe path: a+very large probe side (at least 'parProbeThreshold') and more than one+capability. The shared build index is read-only across threads, so probing it in+parallel needs no synchronization. Independent of build size on purpose: the+build is already tiny; the cost is the 1e7-row probe hashing, which parallelizes+cleanly.+-}+shouldParallelizeSmallBuildProbe :: Int -> Bool+shouldParallelizeSmallBuildProbe probeRows =+    probeRows >= parProbeThreshold+        && capabilities > 1+{-# NOINLINE shouldParallelizeSmallBuildProbe #-}++{- | A read-only view of the build-side index needed by the probe: the raw+open-addressing table vectors of the @CompactIndex@ (which lives in+"DataFrame.Operations.Join"; passing the fields avoids an import cycle).+Passing concrete vectors instead of a lookup closure keeps the per-row probe+loop free of unknown calls and boxed-tuple allocation — the lookup is inlined+into the count and fill loops.+-}+data ProbeTable = ProbeTable+    { ptSorted :: !(VU.Vector Int)+    , ptKeys :: !(VU.Vector Int)+    , ptRuns :: !(VU.Vector Int)+    -- ^ @(start, len)@ packed 32\/32 per slot; @-1@ = empty (see @ciRuns@).+    , ptMask :: {-# UNPACK #-} !Int+    }++{- | Parallel inner-join probe. @parInnerProbe table probeHashes@ returns+@(probeIxs, buildIxs)@ identical to a sequential probe of the same index. The+build index must already be constructed from the build side.+-}+parInnerProbe ::+    ProbeTable ->+    VU.Vector Int ->+    IO (VU.Vector Int, VU.Vector Int)+parInnerProbe = runProbe False++{- | Parallel left-join probe. Like 'parInnerProbe' but every probe row emits at+least one output row; unmatched rows carry a @-1@ sentinel in the build column.+-}+parLeftProbe ::+    ProbeTable ->+    VU.Vector Int ->+    IO (VU.Vector Int, VU.Vector Int)+parLeftProbe = runProbe True++{- | Shared two-pass parallel probe. @keepUnmatched@ selects left- vs+inner-join semantics. Splits @[0, probeN)@ into @caps@ contiguous ranges, counts+each range's output, prefix-sums to global offsets, then fills the single output+buffers in parallel.+-}+runProbe ::+    Bool ->+    ProbeTable ->+    VU.Vector Int ->+    IO (VU.Vector Int, VU.Vector Int)+runProbe keepUnmatched pt probeHashes = do+    caps <- getNumCapabilities+    let !probeN = VU.length probeHashes+        !nChunks = max 1 (min caps probeN)+        !sorted = ptSorted pt+        !keys = ptKeys pt+        !runs = ptRuns pt+        !mask = ptMask pt+        !shift = 64 - popCount mask+        -- Packed (start,len) run for hash @h@, or -1 on a miss. Home slot is+        -- the top log2(cap) bits of a Fibonacci multiply (the row hash's low+        -- bits are poorly diffused); must match the build-side ciSlot exactly.+        findRun !h =+            go+                ( fromIntegral+                    ((fromIntegral h * (0x9E3779B97F4A7C15 :: Word64)) `unsafeShiftR` shift)+                )+          where+            go !slot =+                let !w = runs `VU.unsafeIndex` slot+                 in if w < 0+                        then -1+                        else+                            if keys `VU.unsafeIndex` slot == h+                                then w+                                else go ((slot + 1) .&. mask)+        chunkBounds k = (lo, hi)+          where+            !lo = (probeN * k) `div` nChunks+            !hi = (probeN * (k + 1)) `div` nChunks+        -- Count pass: output rows produced by probe range [lo, hi).+        countRange !lo !hi =+            let go !i !acc+                    | i >= hi = acc+                    | otherwise =+                        let !w = findRun (VU.unsafeIndex probeHashes i)+                         in if w < 0+                                then go (i + 1) (if keepUnmatched then acc + 1 else acc)+                                else go (i + 1) (acc + (w .&. 0xFFFFFFFF))+             in go lo 0+    chunkCounts <- VUM.new (nChunks + 1)+    forkRanges nChunks $ \k ->+        let (lo, hi) = chunkBounds k+         in VUM.unsafeWrite chunkCounts k (countRange lo hi)+    -- Exclusive prefix sum -> per-chunk global start offsets; total at [nChunks].+    let scan !k !acc+            | k > nChunks = pure acc+            | otherwise = do+                c <- if k < nChunks then VUM.unsafeRead chunkCounts k else pure 0+                VUM.unsafeWrite chunkCounts k acc+                scan (k + 1) (acc + c)+    !total <- scan 0 0+    pv <- VUM.unsafeNew (max 1 total)+    bv <- VUM.unsafeNew (max 1 total)+    -- Fill pass: each chunk writes from its prefix-sum offset.+    offs <- VU.unsafeFreeze chunkCounts+    forkRanges nChunks $ \k -> do+        let (lo, hi) = chunkBounds k+            !base = VU.unsafeIndex offs k+            fill !i !p+                | i >= hi = pure ()+                | otherwise = do+                    let !w = findRun (VU.unsafeIndex probeHashes i)+                    if w < 0+                        then+                            if keepUnmatched+                                then do+                                    VUM.unsafeWrite pv p i+                                    VUM.unsafeWrite bv p (-1)+                                    fill (i + 1) (p + 1)+                                else fill (i + 1) p+                        else do+                            let !start = w `unsafeShiftR` 32+                                !len = w .&. 0xFFFFFFFF+                                writeMatch !j !q+                                    | j >= len = pure ()+                                    | otherwise = do+                                        VUM.unsafeWrite pv q i+                                        VUM.unsafeWrite bv q (VU.unsafeIndex sorted (start + j))+                                        writeMatch (j + 1) (q + 1)+                            writeMatch 0 p+                            fill (i + 1) (p + len)+        fill lo base+    pf <- VU.unsafeFreeze (VUM.slice 0 total pv)+    bf <- VU.unsafeFreeze (VUM.slice 0 total bv)+    pure (pf, bf)++{- | Run @body k@ for @k@ in @[0, nChunks)@, one chunk per task, on @nChunks@+forked threads; rethrow the first failure. Chunk @k@ is owned by exactly one+thread, so concurrent writes to disjoint output regions are race-free.+-}+forkRanges :: Int -> (Int -> IO ()) -> IO ()+forkRanges nChunks body = forkJoin_ (map body [0 .. nChunks - 1])

@@ -1,215 +0,0 @@-{-# LANGUAGE BangPatterns #-}-{-# LANGUAGE ScopedTypeVariables #-}--{- |-Parallel chunked-probe join kernels. The build side is indexed once into a-shared, read-only 'CompactIndex' (open-addressing, from-"DataFrame.Operations.Join"); the probe side is split into @caps@ /contiguous/-row ranges and probed in parallel by 'forkIO' workers (no sparks). Each worker-makes two passes over its range — a count pass to size its slice, then a fill-pass — and writes into the single shared output buffers at a precomputed-prefix-sum offset. Because ranges are contiguous and laid out in range order,-the produced @(probeIxs, buildIxs)@ vectors are /bit-for-bit identical/ to the-sequential 'hashInnerKernel' \/ 'hashLeftKernel': probe rows appear in original-order and, within a probe row, build matches in @ciSortedIndices@ order.--This is the parallel==sequential correctness gate (see-@tests/Operations/ParallelJoin.hs@). A sequential fallback is used when there is-a single capability or the probe side is below 'parJoinThreshold'; the caller-('innerJoin' \/ 'leftJoin') decides via 'shouldParallelizeJoin'.--}-module DataFrame.Operations.JoinPar (-    parInnerProbe,-    parLeftProbe,-    shouldParallelizeJoin,-    shouldParallelizeSmallBuildProbe,-    parJoinThreshold,-    parBuildThreshold,-    parProbeThreshold,-) where--import Control.Concurrent (forkIO, getNumCapabilities)-import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)-import Control.Exception (SomeException, throwIO, try)-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as VUM-import System.IO.Unsafe (unsafePerformIO)--{- | Below this many probe rows the fork/coordination overhead is not worth it;-the caller uses its sequential 'ST' kernel instead.--}-parJoinThreshold :: Int-parJoinThreshold = 200000--{- | Below this many build rows the shared 'CompactIndex' is small and hot, so-the sequential hash probe is already memory-bound-fast and the fork overhead-loses (measured: a 1e4-row Text-key build probed by 1e7 rows is /slower/ in-parallel). Parallelism only pays once the build index is large enough to spill-cache — exactly the regime where sort-merge used to be chosen.--}-parBuildThreshold :: Int-parBuildThreshold = 500000--{- | Whether a join should take the parallel probe path: more than one-capability, a probe side of at least 'parJoinThreshold' rows, and a build side-of at least 'parBuildThreshold' rows (a small/hot index is faster probed-sequentially).--}-shouldParallelizeJoin :: Int -> Int -> Bool-shouldParallelizeJoin probeRows buildRows =-    probeRows >= parJoinThreshold-        && buildRows >= parBuildThreshold-        && capabilities > 1-{-# NOINLINE shouldParallelizeJoin #-}--{- | Above this many probe rows the probe-side row hashing and table lookups-dominate the join, so partitioning the probe across cores wins even when the-build side is small and cache-resident (the regime 'shouldParallelizeJoin'-deliberately leaves sequential). Sized at 1e6: below it the per-question gain is-swamped by 'forkIO'/coordination overhead, so small/medium-inner joins stay-sequential (measured). This is the small-build large-probe lever closing the-medium-factor 1e7 join (1e7 probe x ~1e4 build).--}-parProbeThreshold :: Int-parProbeThreshold = 1000000--{- | Whether a /small-build/ join (build below 'parBuildThreshold', so radix-partitioning / sort-merge is not used) should take the parallel probe path: a-very large probe side (at least 'parProbeThreshold') and more than one-capability. The shared build index is read-only across threads, so probing it in-parallel needs no synchronization. Independent of build size on purpose: the-build is already tiny; the cost is the 1e7-row probe hashing, which parallelizes-cleanly.--}-shouldParallelizeSmallBuildProbe :: Int -> Bool-shouldParallelizeSmallBuildProbe probeRows =-    probeRows >= parProbeThreshold-        && capabilities > 1-{-# NOINLINE shouldParallelizeSmallBuildProbe #-}--capabilities :: Int-capabilities = unsafePerformIO getNumCapabilities-{-# NOINLINE capabilities #-}--{- | A read-only view of the build-side index needed by the probe: the lookup-returns @(start, len)@ of the matching run in @sortedIndices@, or @(-1, 0)@ on a-miss. Passed in by the caller so this module need not depend on the-'CompactIndex' record directly.--}-data ProbeIndex = ProbeIndex-    { piSorted :: !(VU.Vector Int)-    , piLookup :: !(Int -> (Int, Int))-    }--{- | Parallel inner-join probe. @parInnerProbe sortedIdxs lookup probeHashes@-returns @(probeIxs, buildIxs)@ identical to a sequential probe of the same-index. The build index must already be constructed from the build side.--}-parInnerProbe ::-    VU.Vector Int ->-    (Int -> (Int, Int)) ->-    VU.Vector Int ->-    IO (VU.Vector Int, VU.Vector Int)-parInnerProbe sortedIdxs lookupFn =-    runProbe False (ProbeIndex sortedIdxs lookupFn)--{- | Parallel left-join probe. Like 'parInnerProbe' but every probe row emits at-least one output row; unmatched rows carry a @-1@ sentinel in the build column.--}-parLeftProbe ::-    VU.Vector Int ->-    (Int -> (Int, Int)) ->-    VU.Vector Int ->-    IO (VU.Vector Int, VU.Vector Int)-parLeftProbe sortedIdxs lookupFn =-    runProbe True (ProbeIndex sortedIdxs lookupFn)--{- | Shared two-pass parallel probe. @keepUnmatched@ selects left- vs-inner-join semantics. Splits @[0, probeN)@ into @caps@ contiguous ranges, counts-each range's output, prefix-sums to global offsets, then fills the single output-buffers in parallel.--}-runProbe ::-    Bool ->-    ProbeIndex ->-    VU.Vector Int ->-    IO (VU.Vector Int, VU.Vector Int)-runProbe keepUnmatched pidx probeHashes = do-    caps <- getNumCapabilities-    let !probeN = VU.length probeHashes-        !nChunks = max 1 (min caps probeN)-        !sorted = piSorted pidx-        !lookupFn = piLookup pidx-        chunkBounds k = (lo, hi)-          where-            !lo = (probeN * k) `div` nChunks-            !hi = (probeN * (k + 1)) `div` nChunks-        -- Count pass: output rows produced by probe range [lo, hi).-        countRange !lo !hi =-            let go !i !acc-                    | i >= hi = acc-                    | otherwise =-                        let (!start, !len) = lookupFn (VU.unsafeIndex probeHashes i)-                         in if start < 0-                                then go (i + 1) (if keepUnmatched then acc + 1 else acc)-                                else go (i + 1) (acc + len)-             in go lo 0-    chunkCounts <- VUM.new (nChunks + 1)-    forkRanges nChunks $ \k ->-        let (lo, hi) = chunkBounds k-         in VUM.unsafeWrite chunkCounts k (countRange lo hi)-    -- Exclusive prefix sum -> per-chunk global start offsets; total at [nChunks].-    let scan !k !acc-            | k > nChunks = pure acc-            | otherwise = do-                c <- if k < nChunks then VUM.unsafeRead chunkCounts k else pure 0-                VUM.unsafeWrite chunkCounts k acc-                scan (k + 1) (acc + c)-    !total <- scan 0 0-    pv <- VUM.unsafeNew (max 1 total)-    bv <- VUM.unsafeNew (max 1 total)-    -- Fill pass: each chunk writes from its prefix-sum offset.-    offs <- VU.unsafeFreeze chunkCounts-    forkRanges nChunks $ \k -> do-        let (lo, hi) = chunkBounds k-            !base = VU.unsafeIndex offs k-            fill !i !p-                | i >= hi = pure ()-                | otherwise = do-                    let (!start, !len) = lookupFn (VU.unsafeIndex probeHashes i)-                    if start < 0-                        then-                            if keepUnmatched-                                then do-                                    VUM.unsafeWrite pv p i-                                    VUM.unsafeWrite bv p (-1)-                                    fill (i + 1) (p + 1)-                                else fill (i + 1) p-                        else do-                            let writeMatch !j !q-                                    | j >= len = pure ()-                                    | otherwise = do-                                        VUM.unsafeWrite pv q i-                                        VUM.unsafeWrite bv q (VU.unsafeIndex sorted (start + j))-                                        writeMatch (j + 1) (q + 1)-                            writeMatch 0 p-                            fill (i + 1) (p + len)-        fill lo base-    pf <- VU.unsafeFreeze (VUM.slice 0 total pv)-    bf <- VU.unsafeFreeze (VUM.slice 0 total bv)-    pure (pf, bf)--{- | Run @body k@ for @k@ in @[0, nChunks)@, one chunk per task, on @nChunks@-forked threads; rethrow the first failure. Chunk @k@ is owned by exactly one-thread, so concurrent writes to disjoint output regions are race-free.--}-forkRanges :: Int -> (Int -> IO ()) -> IO ()-forkRanges nChunks body = do-    vars <- mapM spawn [0 .. nChunks - 1]-    results <- mapM takeMVar vars-    mapM_ (either (throwIO :: SomeException -> IO ()) pure) results-  where-    spawn k = do-        var <- newEmptyMVar-        _ <- forkIO (try (body k) >>= putMVar var)-        pure var

@@ -49,7 +49,7 @@                                 Nothing ->                                     D.insertColumn name (D.leftExpandColumn sumRows b'') df                                 Just a'' ->-                                    let concatedColumns = D.concatColumnsEither a'' b''+                                    let concatedColumns = D.mappendColumnsEither a'' b''                                      in D.insertColumn name concatedColumns df             result = L.foldl' (addColumns a b) D.empty (D.columnNames a `L.union` D.columnNames b)          in

@@ -34,13 +34,13 @@     atIndicesStable,     materializeMerged,  )+import DataFrame.Internal.Data.PackedText (packedSlice, sliceCmpBytes) import DataFrame.Internal.DataFrame (     DataFrame (..),     columnNames,     unsafeGetColumn,  ) import DataFrame.Internal.Expression (Expr (Col), getColumns)-import DataFrame.Internal.PackedText (packedSlice, sliceCmpBytes) import DataFrame.Operations.Core (dimensions) import DataFrame.Operations.Transformations (derive) import System.Random (Random (randomR), RandomGen)@@ -191,14 +191,12 @@     shuffleVec :: (RandomGen g) => g -> VU.Vector Int     shuffleVec g = runST $ do         vm <- VUM.generate k id-        let (n, nGen) = randomR (1, k - 1) g-        go vm n nGen+        go vm (k - 1) g         VU.unsafeFreeze vm -    go _v (-1) _ = pure ()-    go _v 0 _ = pure ()-    go v maxInd gen =+    go _v i _ | i <= 0 = pure ()+    go v i gen =         let-            (n, nextGen) = randomR (1, maxInd) gen+            (j, nextGen) = randomR (0, i) gen          in-            VUM.swap v 0 n *> go (VUM.tail v) (maxInd - 1) nextGen+            VUM.swap v i j *> go v (i - 1) nextGen

@@ -1,19 +1,6 @@+{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ScopedTypeVariables #-} -{- |-Module      : DataFrame.Operations.SetOps-Description : Set-theoretic ("topos") row operations.--These treat a 'DataFrame' as a /set/ of rows and implement the subobject-lattice from relational algebra: 'union', 'intersect', 'difference', and-'symmetricDifference'. Every result is deduplicated, so each operation has the-schema-preserving shape @DataFrame -> DataFrame -> DataFrame@.--Row equality is the same hash-based notion used by 'distinct' (see-"DataFrame.Operations.Aggregation"), so these operations and 'distinct' agree-on what "the same row" means. Both inputs are expected to share a schema; the-typed layer ('DataFrame.Typed') enforces that statically.--} module DataFrame.Operations.SetOps (     union,     intersect,@@ -79,8 +66,8 @@     chosen =         [ VU.head members         | k <- [0 .. nGroups - 1]-        , let s = VU.unsafeIndex offs k-              e = VU.unsafeIndex offs (k + 1)+        , let !s = VU.unsafeIndex offs k+              !e = VU.unsafeIndex offs (k + 1)               members = VU.slice s (e - s) vis               inLeft = VU.any (< leftRows) members               inRight = VU.any (>= leftRows) members

@@ -57,11 +57,10 @@     getColumn,  ) import DataFrame.Internal.Expression+import DataFrame.Internal.Expression.Operators.Nullable (BaseType) import DataFrame.Internal.Interpreter-import DataFrame.Internal.Nullable (BaseType) import DataFrame.Internal.Row (showValue, toAny) import DataFrame.Internal.Statistics-import DataFrame.Internal.Types import DataFrame.Operations.Core import DataFrame.Operations.Subset (filterJust) import DataFrame.Operations.Transformations (ImputeOp (..), imputeCore)

@@ -5,6 +5,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-}@@ -48,7 +49,6 @@     stratifiedSplit,      -- * Label helpers (exported for satellite packages)-    columnToTextVec,     rowsAtIndices, ) where @@ -74,25 +74,53 @@     DataFrameException (..),     TypeErrorContext (..),  )-import DataFrame.Internal.Column+import DataFrame.Expression.Operators (+    col,+    name,+    (.<.),+    (.<=.),+    (.>.),+    (.>=.),+ )+import DataFrame.Internal.Column (+    Column (..),+    Columnable,+    TypedColumn (TColumn),+    atIndicesStable,+    columnToTextVec,+    findIndices,+    hasMissing,+    materializeMerged,+    materializePacked,+    mkRandom,+    sliceColumn,+    takeColumn,+    takeLastColumn,+ )+import DataFrame.Internal.Column.Bitmap (bitmapTestBit) import DataFrame.Internal.DataFrame (     DataFrame (..),     columnNames,+    dataframeDimensions,     derivingExpressions,     empty,     getColumn,     insertColumn,     unsafeGetColumn,  )-import DataFrame.Internal.Expression-import DataFrame.Internal.Interpreter-import DataFrame.Internal.PackedText (packedIndexText, packedLength)+import DataFrame.Internal.Expression (Expr (Col, Lit), normalize)+import DataFrame.Internal.Interpreter (Ctx (..), eval, interpret, materialize) import DataFrame.Operations.Core () import DataFrame.Operations.Merge () import DataFrame.Operations.Transformations (apply)-import DataFrame.Operators-import System.Random-import Type.Reflection+import System.Random (RandomGen, SplitGen (..))+import Type.Reflection (+    eqTypeRep,+    typeRep,+    pattern App,+    type (:~:) (Refl),+    type (:~~:) (HRefl),+ ) import Prelude hiding (drop, filter, take)  #if MIN_VERSION_random(1,3,0)@@ -148,12 +176,16 @@ range :: (Int, Int) -> DataFrame -> DataFrame range (start, end) d =     d-        { columns = V.map (sliceColumn (clip start 0 r) n') (columns d)+        { columns = V.map (sliceColumn start' n') (columns d)         , dataframeDimensions = (n', c)         }   where     (r, c) = dataframeDimensions d-    n' = clip (end - start) 0 r+    start' = clip start 0 r+    -- Clamp both endpoints before subtracting: end - start' on an unclamped+    -- end wraps for very negative values and reopens the range.+    end' = clip end start' r+    n' = end' - start'  clip :: Int -> Int -> Int -> Int clip n left right = min right $ max n left@@ -518,27 +550,6 @@      in         map (exclude [name cRand]) (go (folds - 1) withRand) --- | Convert any Column to a vector of Text labels (one per row).-columnToTextVec :: Column -> V.Vector T.Text-columnToTextVec c@(MergedColumn _ _) = columnToTextVec (materializeMerged c)-columnToTextVec (BoxedColumn bm (col' :: V.Vector a)) =-    case bm of-        Nothing -> case testEquality (typeRep @a) (typeRep @T.Text) of-            Just Refl -> col'-            Nothing -> V.map (T.pack . show) col'-        Just bitmap ->-            V.imap (\i x -> if bitmapTestBit bitmap i then T.pack (show x) else "null") col'-columnToTextVec (UnboxedColumn bm col') =-    case bm of-        Nothing -> V.map (T.pack . show) (V.convert col')-        Just bitmap ->-            V.generate (VU.length col') $ \i ->-                if bitmapTestBit bitmap i then T.pack (show (col' VU.! i)) else "null"-columnToTextVec (PackedText bm p) =-    V.generate (packedLength p) $ \i -> case bm of-        Just bitmap | not (bitmapTestBit bitmap i) -> "null"-        _ -> packedIndexText p i- -- | Build a map from stringified label to row indices. groupByIndices :: Column -> M.Map T.Text (VU.Vector Int) groupByIndices col' =@@ -573,7 +584,9 @@ stratifiedSample gen p strataCol df =     let col' = case strataCol of             Col colName -> unsafeGetColumn colName df-            _ -> unwrapTypedColumn (either throw id (interpret @a df strataCol))+            _ -> either throw id $ do+                v <- eval (FlatCtx df) strataCol+                pure $ materialize @a (fst (dataframeDimensions df)) v         groups = M.elems (groupByIndices col')         go _ [] = mempty         go g (ixs : rest) =@@ -597,7 +610,9 @@ stratifiedSplit gen p strataCol df =     let col' = case strataCol of             Col colName -> unsafeGetColumn colName df-            _ -> unwrapTypedColumn (either throw id (interpret @a df strataCol))+            _ -> either throw id $ do+                v <- eval (FlatCtx df) strataCol+                pure $ materialize @a (fst (dataframeDimensions df)) v         groups = M.elems (groupByIndices col')         go _ [] = (mempty, mempty)         go g (ixs : rest) =

@@ -52,8 +52,8 @@  ) import DataFrame.Internal.DataFrame (DataFrame (..), getColumn, insertColumn) import DataFrame.Internal.Expression+import DataFrame.Internal.Expression.Operators.Nullable (BaseType) import DataFrame.Internal.Interpreter-import DataFrame.Internal.Nullable (BaseType) import DataFrame.Operations.Core import GHC.TypeLits (ErrorMessage (..), TypeError) import Type.Reflection (typeRep)

@@ -29,15 +29,14 @@ import Data.Time import Data.Type.Equality (TestEquality (..)) import DataFrame.Internal.Column (-    Bitmap,     Column (..),     Columnable,-    bitmapTestBit,     ensureOptional,     finalizeParseResult,     fromVector,     materializePacked,  )+import DataFrame.Internal.Column.Bitmap (Bitmap, bitmapTestBit) import DataFrame.Internal.DataFrame (     DataFrame (..),     insertColumn,

@@ -137,12 +137,13 @@  import qualified DataFrame.Functions as F import DataFrame.Internal.Column (Columnable)+import DataFrame.Internal.Column.Types (Promote, PromoteDiv) import DataFrame.Internal.Expression (     BinUDF (..),     Expr (..),     UnUDF (..),  )-import DataFrame.Internal.Nullable (+import DataFrame.Internal.Expression.Operators.Nullable (     BaseType,     DivWidenOp,     NullCmpResult,@@ -158,7 +159,6 @@     widenArithOp,     widenCmpOp,  )-import DataFrame.Internal.Types (Promote, PromoteDiv)  import qualified Data.Vector.Unboxed as VU import DataFrame.Typed.Expr.Extra

@@ -46,7 +46,7 @@  import DataFrame.Internal.Column (Columnable) import qualified DataFrame.Internal.DataFrame as D-import DataFrame.Internal.Nullable (BaseType)+import DataFrame.Internal.Expression.Operators.Nullable (BaseType) import qualified DataFrame.Operations.Core as Core import qualified DataFrame.Operations.Statistics as Stats import DataFrame.Operations.Transformations (ImputeOp)