diff --git a/dataframe-operations.cabal b/dataframe-operations.cabal --- a/dataframe-operations.cabal +++ b/dataframe-operations.cabal @@ -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, diff --git a/src/DataFrame/Functions.hs b/src/DataFrame/Functions.hs --- a/src/DataFrame/Functions.hs +++ b/src/DataFrame/Functions.hs @@ -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) diff --git a/src/DataFrame/Monad.hs b/src/DataFrame/Monad.hs --- a/src/DataFrame/Monad.hs +++ b/src/DataFrame/Monad.hs @@ -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 diff --git a/src/DataFrame/Operations/AggregateScatter.hs b/src/DataFrame/Operations/AggregateScatter.hs deleted file mode 100644 --- a/src/DataFrame/Operations/AggregateScatter.hs +++ /dev/null @@ -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 diff --git a/src/DataFrame/Operations/Aggregation.hs b/src/DataFrame/Operations/Aggregation.hs --- a/src/DataFrame/Operations/Aggregation.hs +++ b/src/DataFrame/Operations/Aggregation.hs @@ -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 diff --git a/src/DataFrame/Operations/Aggregation/Run.hs b/src/DataFrame/Operations/Aggregation/Run.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Aggregation/Run.hs @@ -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 diff --git a/src/DataFrame/Operations/Inference.hs b/src/DataFrame/Operations/Inference.hs --- a/src/DataFrame/Operations/Inference.hs +++ b/src/DataFrame/Operations/Inference.hs @@ -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) -> diff --git a/src/DataFrame/Operations/Join.hs b/src/DataFrame/Operations/Join.hs --- a/src/DataFrame/Operations/Join.hs +++ b/src/DataFrame/Operations/Join.hs @@ -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 ) diff --git a/src/DataFrame/Operations/Join/Parallel.hs b/src/DataFrame/Operations/Join/Parallel.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Join/Parallel.hs @@ -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]) diff --git a/src/DataFrame/Operations/JoinPar.hs b/src/DataFrame/Operations/JoinPar.hs deleted file mode 100644 --- a/src/DataFrame/Operations/JoinPar.hs +++ /dev/null @@ -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 diff --git a/src/DataFrame/Operations/Merge.hs b/src/DataFrame/Operations/Merge.hs --- a/src/DataFrame/Operations/Merge.hs +++ b/src/DataFrame/Operations/Merge.hs @@ -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 diff --git a/src/DataFrame/Operations/Permutation.hs b/src/DataFrame/Operations/Permutation.hs --- a/src/DataFrame/Operations/Permutation.hs +++ b/src/DataFrame/Operations/Permutation.hs @@ -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 diff --git a/src/DataFrame/Operations/SetOps.hs b/src/DataFrame/Operations/SetOps.hs --- a/src/DataFrame/Operations/SetOps.hs +++ b/src/DataFrame/Operations/SetOps.hs @@ -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 diff --git a/src/DataFrame/Operations/Statistics.hs b/src/DataFrame/Operations/Statistics.hs --- a/src/DataFrame/Operations/Statistics.hs +++ b/src/DataFrame/Operations/Statistics.hs @@ -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) diff --git a/src/DataFrame/Operations/Subset.hs b/src/DataFrame/Operations/Subset.hs --- a/src/DataFrame/Operations/Subset.hs +++ b/src/DataFrame/Operations/Subset.hs @@ -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) = diff --git a/src/DataFrame/Operations/Transformations.hs b/src/DataFrame/Operations/Transformations.hs --- a/src/DataFrame/Operations/Transformations.hs +++ b/src/DataFrame/Operations/Transformations.hs @@ -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) diff --git a/src/DataFrame/Operations/Typing.hs b/src/DataFrame/Operations/Typing.hs --- a/src/DataFrame/Operations/Typing.hs +++ b/src/DataFrame/Operations/Typing.hs @@ -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, diff --git a/src/DataFrame/Typed/Expr.hs b/src/DataFrame/Typed/Expr.hs --- a/src/DataFrame/Typed/Expr.hs +++ b/src/DataFrame/Typed/Expr.hs @@ -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 diff --git a/src/DataFrame/Typed/Statistics.hs b/src/DataFrame/Typed/Statistics.hs --- a/src/DataFrame/Typed/Statistics.hs +++ b/src/DataFrame/Typed/Statistics.hs @@ -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)