diff --git a/dataframe-operations.cabal b/dataframe-operations.cabal --- a/dataframe-operations.cabal +++ b/dataframe-operations.cabal @@ -1,7 +1,6 @@ cabal-version: 2.4 name: dataframe-operations -version: 1.1.0.3 - +version: 1.1.1.0 synopsis: Column operations, expression DSL, and statistics for the dataframe ecosystem. description: Untyped column operations (select, filter, sort, join, groupBy, @@ -34,9 +33,12 @@ DataFrame.Functions DataFrame.Monad DataFrame.Internal.Statistics + DataFrame.Operations.AggregateScatter DataFrame.Operations.Aggregation DataFrame.Operations.Core + DataFrame.Operations.Inference DataFrame.Operations.Join + DataFrame.Operations.JoinPar DataFrame.Operations.Merge DataFrame.Operations.Permutation DataFrame.Operations.SetOps @@ -50,14 +52,14 @@ DataFrame.Typed.Join DataFrame.Typed.Operations build-depends: base >= 4 && < 5, + bytestring >= 0.11 && < 0.13, containers >= 0.6.7 && < 0.9, - dataframe-core >= 1.0.2 && < 1.1, - dataframe-parsing ^>= 1.0, + dataframe-core >= 1.1 && < 1.2, + dataframe-parsing ^>= 1.0.2, random >= 1.2 && < 2, regex-tdfa >= 1.3.0 && < 2, text >= 2.0 && < 3, time >= 1.12 && < 2, - unordered-containers >= 0.1 && < 1, vector ^>= 0.13, vector-algorithms ^>= 0.9 hs-source-dirs: src diff --git a/src/DataFrame/Operations/AggregateScatter.hs b/src/DataFrame/Operations/AggregateScatter.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/AggregateScatter.hs @@ -0,0 +1,224 @@ +{-# 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 @@ -1,7 +1,7 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE Strict #-} @@ -17,19 +17,14 @@ import qualified Data.Text as T import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Unboxed.Mutable as VUM import Control.Exception (throw) -import Control.Monad -import Control.Monad.ST (ST, runST) -import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) import DataFrame.Errors +import DataFrame.Internal.AggPlan (MomentPlan, planAgg, planMoments) import DataFrame.Internal.Column ( - Bitmap, Column (..), TypedColumn (..), atIndicesStable, - bitmapTestBit, ) import DataFrame.Internal.DataFrame ( DataFrame (..), @@ -39,123 +34,67 @@ ) import DataFrame.Internal.Expression import DataFrame.Internal.Grouping (buildRowToGroup, changingPoints, groupBy) -import DataFrame.Internal.Hash ( - fnvOffset, - mixDouble, - mixInt, - mixShow, - mixText, - nullSalt, - ) import DataFrame.Internal.Interpreter -import DataFrame.Internal.Types +import DataFrame.Internal.RowHash (computeRowHashesIO) +import DataFrame.Operations.AggregateScatter (runMomentPlan, runPlan) import DataFrame.Operations.Core import DataFrame.Operations.Subset -import Type.Reflection (typeRep) +import System.IO.Unsafe (unsafePerformIO) +{- | Per-row key hash over the selected key columns. Delegates to the shared +'computeRowHashesIO' kernel, which forks over contiguous row ranges for large +frames (the hashing of a wide 1e7-row text/factor join key dominates that join) +and is bit-for-bit identical to a single sequential pass at any capability count. +-} computeRowHashes :: [Int] -> DataFrame -> VU.Vector Int -computeRowHashes indices df = runST $ do +computeRowHashes indices df = let n = fst (dimensions df) - mv <- VUM.replicate n fnvOffset - - let selectedCols = map (columns df V.!) indices - - forM_ selectedCols $ \case - UnboxedColumn ubm (v :: VU.Vector a) -> - case testEquality (typeRep @a) (typeRep @Int) of - Just Refl -> hashKeyUnboxed mv ubm mixInt v - Nothing -> - case testEquality (typeRep @a) (typeRep @Double) of - Just Refl -> - hashKeyUnboxed mv ubm mixDouble v - Nothing -> - case sIntegral @a of - STrue -> - hashKeyUnboxed mv ubm (\h d -> mixInt h (fromIntegral @a @Int d)) v - SFalse -> - case sFloating @a of - STrue -> - hashKeyUnboxed mv ubm (\h d -> mixDouble h (realToFrac d :: Double)) v - SFalse -> - hashKeyUnboxed mv ubm mixShow v - BoxedColumn bm (v :: V.Vector a) -> - case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> - V.imapM_ - ( \i (t :: T.Text) -> do - h <- VUM.unsafeRead mv i - let h' = case bm of - Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt - _ -> mixText h t - VUM.unsafeWrite mv i h' - ) - v - Nothing -> - V.imapM_ - ( \i d -> do - h <- VUM.unsafeRead mv i - let h' = case bm of - Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt - _ -> mixShow h d - VUM.unsafeWrite mv i h' - ) - v - - VU.unsafeFreeze mv - -{- | Fold a value-mix over an unboxed key column into the running hash vector, -respecting the null bitmap (a null slot mixes 'nullSalt' instead of its -uninitialised stored value, so @Nothing@ never matches a present @Just x@). -The non-nullable case is the original tight loop. @INLINE@d to specialise per call. --} -hashKeyUnboxed :: - (VU.Unbox a) => - VUM.MVector s Int -> - Maybe Bitmap -> - (Int -> a -> Int) -> - VU.Vector a -> - ST s () -hashKeyUnboxed mv ubm mix v = case ubm of - Nothing -> - VU.imapM_ - ( \i x -> do - h <- VUM.unsafeRead mv i - VUM.unsafeWrite mv i (mix h x) - ) - v - Just bm -> - VU.imapM_ - ( \i x -> do - h <- VUM.unsafeRead mv i - VUM.unsafeWrite - mv - i - (if bitmapTestBit bm i then mix h x else mixInt h nullSalt) - ) - v -{-# INLINE hashKeyUnboxed #-} + selectedCols = map (columns df V.!) indices + in unsafePerformIO (computeRowHashesIO n selectedCols) +{-# NOINLINE computeRowHashes #-} {- | Aggregate a grouped dataframe using the expressions given. All ungrouped columns will be dropped. -} aggregate :: [NamedExpr] -> GroupedDataFrame -> DataFrame -aggregate aggs gdf@(Grouped df groupingColumns valIndices offs _rowToGroup) = +aggregate aggs gdf@(Grouped df groupingColumns valIndices offs rowToGroupV) = let df' = selectIndices (VU.map (valIndices VU.!) (VU.init offs)) (select groupingColumns df) - f (name, UExpr (expr :: Expr a)) d = - let - value = case interpretAggregation @a gdf expr of - Left e -> throw e - Right (UnAggregated _) -> throw $ UnaggregatedException (T.pack $ show expr) - Right (Aggregated (TColumn col)) -> col - in - insertColumn name value d + !nGroups = VU.length offs - 1 + + -- 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 + in insertColumn name value d + + -- Fused fast path: the Q9 regression family (count + five moment sums + -- of two base columns) becomes one scatter over the base columns, + -- dropping the derived product columns and the six separate folds. + -- 'planMoments' returns 'Nothing' on any other set, falling back below. + fusedMoments = do + mp <- planMoments gdf aggs :: Maybe MomentPlan + runMomentPlan gdf nGroups mp in - fold f aggs df' + case fusedMoments of + Just cols -> fold (uncurry insertColumn) cols df' + Nothing -> fold f aggs df' + +-- | The fall-back path: evaluate one named aggregation via the interpreter. +interpretNamed :: GroupedDataFrame -> NamedExpr -> Column +interpretNamed gdf (_, UExpr (expr :: Expr a)) = + case interpretAggregation @a gdf expr of + Left e -> throw e + Right (UnAggregated _) -> throw $ UnaggregatedException (T.pack $ show expr) + Right (Aggregated (TColumn col)) -> col selectIndices :: VU.Vector Int -> DataFrame -> DataFrame selectIndices xs df = diff --git a/src/DataFrame/Operations/Core.hs b/src/DataFrame/Operations/Core.hs --- a/src/DataFrame/Operations/Core.hs +++ b/src/DataFrame/Operations/Core.hs @@ -32,6 +32,7 @@ columnTypeString, fromList, fromVector, + materializePacked, toDoubleVector, toFloatVector, toIntVector, @@ -538,6 +539,7 @@ countNulls columnType : acc + go acc i col@(PackedText _ _) = go acc i (materializePacked col) nulls :: Column -> Int nulls (BoxedColumn (Just bm) xs) = @@ -552,6 +554,7 @@ nulls (UnboxedColumn (Just bm) xs) = let n = VG.length xs in n - VU.foldl' (\acc b -> acc + popCount b) 0 bm +nulls c@(PackedText _ _) = nulls (materializePacked c) nulls _ = 0 {- | Creates a dataframe from a list of tuples with name and column. @@ -748,12 +751,7 @@ pure $ V.generate (fst (dataframeDimensions df)) - ( \i -> - foldl - (\acc j -> acc `VU.snoc` ((m VG.! j) VG.! i)) - VU.empty - [0 .. (V.length m - 1)] - ) + (\i -> VU.generate (V.length m) (\j -> (m VG.! j) VG.! i)) {- | Returns a dataframe as a two dimensional vector of doubles. @@ -775,12 +773,7 @@ pure $ V.generate (fst (dataframeDimensions df)) - ( \i -> - foldl - (\acc j -> acc `VU.snoc` ((m VG.! j) VG.! i)) - VU.empty - [0 .. (V.length m - 1)] - ) + (\i -> VU.generate (V.length m) (\j -> (m VG.! j) VG.! i)) {- | Returns a dataframe as a two dimensional vector of ints. @@ -801,12 +794,7 @@ pure $ V.generate (fst (dataframeDimensions df)) - ( \i -> - foldl - (\acc j -> acc `VU.snoc` ((m VG.! j) VG.! i)) - VU.empty - [0 .. (V.length m - 1)] - ) + (\i -> VU.generate (V.length m) (\j -> (m VG.! j) VG.! i)) {- | Get a specific column as a vector. @@ -951,4 +939,6 @@ Just (BoxedColumn Nothing (_ :: V.Vector a)) -> UExpr (Col @a name) Just (UnboxedColumn (Just _) (_ :: VU.Vector a)) -> UExpr (Col @(Maybe a) name) Just (UnboxedColumn Nothing (_ :: VU.Vector a)) -> UExpr (Col @a name) + Just (PackedText (Just _) _) -> UExpr (Col @(Maybe T.Text) name) + Just (PackedText Nothing _) -> UExpr (Col @T.Text name) Nothing -> error $ "showDerivedExpressions: column not found: " ++ T.unpack name diff --git a/src/DataFrame/Operations/Inference.hs b/src/DataFrame/Operations/Inference.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Inference.hs @@ -0,0 +1,221 @@ +{-# 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 (..), + makeParsingAssumption, + makeParsingAssumptionBytes, + byteStringDateParser, + parseTimeOpt, + promoteIntColumn, + promoteIntColumnIndexed, + readIntStrict, +) where + +import qualified Data.ByteString as BS +import qualified Data.Text as T +import qualified Data.Text.Encoding as TE +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed.Mutable as VUM + +import Control.Monad.ST (runST) +import Data.Maybe (isJust) +import Data.Time (Day, defaultTimeLocale, parseTimeM) +import DataFrame.Internal.Column (Column (..), finalizeParseResult) +import DataFrame.Internal.Parsing (readByteStringDate, readInt) +import DataFrame.Internal.Parsing.Fast ( + parseBoolField, + parseDateField, + parseDoubleField, + parseIntField, + ) + +type DateFormat = String + +data ParsingAssumption + = BoolAssumption + | IntAssumption + | DoubleAssumption + | DateAssumption + | NoAssumption + | TextAssumption + deriving (Eq, Show) + +parseTimeOpt :: DateFormat -> T.Text -> Maybe Day +parseTimeOpt dateFormat s = + parseTimeM {- Accept leading/trailing whitespace -} + True + defaultTimeLocale + 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 + | not seen = NoAssumption + | b = BoolAssumption + | i && d = IntAssumption + | d = DoubleAssumption + | 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 + where + n = V.length cells + go !i !seen !b !int !d !dt + | i >= n || (seen && not (b || int || d || dt)) = + pickAssumption seen b int d dt + | otherwise = case V.unsafeIndex cells i of + Nothing -> go (i + 1) seen b int d dt + Just t -> + let bs = TE.encodeUtf8 t + in go + (i + 1) + True + (b && isJust (parseBoolField bs)) + (int && isJust (parseIntField bs)) + (d && isJust (parseDoubleField bs)) + (dt && isJust (parseTimeOpt dfmt t)) + +-- | Byte-level twin of 'makeParsingAssumption' for raw field bytes. +makeParsingAssumptionBytes :: + DateFormat -> V.Vector (Maybe BS.ByteString) -> ParsingAssumption +makeParsingAssumptionBytes dfmt cells = go 0 False True True True True + where + n = V.length cells + dateP = byteStringDateParser dfmt + go !i !seen !b !int !d !dt + | i >= n || (seen && not (b || int || d || dt)) = + pickAssumption seen b int d dt + | otherwise = case V.unsafeIndex cells i of + Nothing -> go (i + 1) seen b int d dt + Just bs -> + go + (i + 1) + True + (b && isJust (parseBoolField bs)) + (int && isJust (parseIntField bs)) + (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) -> + (src -> Maybe Int) -> + (src -> Maybe Double) -> + V.Vector src -> + Maybe Column +promoteIntColumn isNullAt parseI parseD vec = + promoteIntColumnIndexed + (V.length vec) + (\i -> isNullAt i (V.unsafeIndex vec i)) + (parseI . V.unsafeIndex vec) + (parseD . V.unsafeIndex vec) +{-# INLINE promoteIntColumn #-} + +{- | 'promoteIntColumn' over row indices instead of a materialized cell +vector (used by readers that keep cells as buffer offsets). +-} +promoteIntColumnIndexed :: + Int -> + (Int -> Bool) -> + (Int -> Maybe Int) -> + (Int -> Maybe Double) -> + Maybe Column +promoteIntColumnIndexed n isNullAt parseI parseD = runST $ do + values <- VUM.unsafeNew n + vmask <- VUM.unsafeNew n + let intLoop !i !anyNull + | i >= n = + fmap (uncurry UnboxedColumn) + <$> finalizeParseResult values vmask anyNull + | isNullAt i = do + VUM.unsafeWrite vmask i 0 + VUM.unsafeWrite values i 0 + intLoop (i + 1) True + | otherwise = case parseI i of + Just v -> do + VUM.unsafeWrite vmask i 1 + VUM.unsafeWrite values i v + intLoop (i + 1) anyNull + Nothing -> case parseD i of + Just dv -> promote i anyNull dv + Nothing -> return Nothing + + -- Switch to Double at row i: convert the Int prefix, then + -- continue parsing the remaining cells as Double. + promote !i !anyNull !dv = do + dvalues <- VUM.unsafeNew n + let copy !j + | j >= i = return () + | otherwise = do + x <- VUM.unsafeRead values j + VUM.unsafeWrite dvalues j (fromIntegral x) + copy (j + 1) + copy 0 + VUM.unsafeWrite vmask i 1 + VUM.unsafeWrite dvalues i dv + dblLoop dvalues (i + 1) anyNull + + dblLoop dvalues !i !anyNull + | i >= n = + fmap (uncurry UnboxedColumn) + <$> finalizeParseResult dvalues vmask anyNull + | isNullAt i = do + VUM.unsafeWrite vmask i 0 + VUM.unsafeWrite dvalues i 0 + dblLoop dvalues (i + 1) True + | otherwise = case parseD i of + Just dv -> do + VUM.unsafeWrite vmask i 1 + VUM.unsafeWrite dvalues i dv + dblLoop dvalues (i + 1) anyNull + Nothing -> return Nothing + intLoop 0 False +{-# INLINE promoteIntColumnIndexed #-} 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 @@ -12,12 +12,9 @@ import Control.Applicative ((<|>)) import Control.Exception (throw) -import Control.Monad (forM_, when) +import Control.Monad (when) import Control.Monad.ST (ST, runST) -import qualified Data.HashMap.Strict as HM -#if !MIN_VERSION_base(4,20,0) -import Data.List (foldl') -#endif +import Data.Bits ((.&.)) import qualified Data.Map.Strict as M import Data.Maybe (fromMaybe) import Data.STRef (newSTRef, readSTRef, writeSTRef) @@ -33,8 +30,16 @@ ) import DataFrame.Internal.Column as D 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 ( + parInnerProbe, + parLeftProbe, + shouldParallelizeJoin, + shouldParallelizeSmallBuildProbe, + ) +import System.IO.Unsafe (unsafePerformIO) import Type.Reflection -- | Equivalent to SQL join types. @@ -58,44 +63,103 @@ join FULL_OUTER xs right = fullOuterJoin xs right {- | Row-count threshold for the build side. -When the build side exceeds this, sort-merge join is used -instead of hash join to avoid L3 cache thrashing. +When the build side exceeds this, sort-merge join is used instead of the +single-threaded hash join. The parallel chunked-probe hash join +('parInnerKernel' \/ 'parLeftKernel') is preferred above this size when more +than one capability is available (see 'shouldParallelizeJoin'): partitioning +the probe over cores beats sort-merge for the large build sides measured here +(e.g. the 2M-row build in the join pipeline). Sort-merge remains the +single-threaded fallback. -} joinStrategyThreshold :: Int joinStrategyThreshold = 500_000 {- | A compact index mapping hash values to contiguous slices of -original row indices. All indices live in a single unboxed vector; -the HashMap stores @(offset, length)@ into that vector. +original row indices. All indices live in a single unboxed vector +(@ciSortedIndices@, sorted by hash). The lookup table is an open-addressing +linear-probe hash table held in three parallel unboxed vectors keyed by hash: +@ciKeys@ holds the hash at each slot, @ciStarts@ the run offset into +@ciSortedIndices@ (@-1@ marks an empty slot, since real offsets are @>= 0@), +and @ciLens@ the run length. @ciMask@ is @tableSize - 1@ (table size is a +power of two), used to map a hash to its home slot. -} data CompactIndex = CompactIndex { ciSortedIndices :: {-# UNPACK #-} !(VU.Vector Int) - , ciOffsets :: !(HM.HashMap Int (Int, Int)) + , ciKeys :: {-# UNPACK #-} !(VU.Vector Int) + , ciStarts :: {-# UNPACK #-} !(VU.Vector Int) + , ciLens :: {-# UNPACK #-} !(VU.Vector Int) + , ciMask :: {-# UNPACK #-} !Int } +{- | 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) + where + !mask = ciMask ci + !keys = ciKeys ci + !starts = ciStarts ci + !lens = ciLens ci + go !slot = + let !s = starts `VU.unsafeIndex` slot + in if s < 0 + then (-1, 0) + else + if keys `VU.unsafeIndex` slot == h + then (s, lens `VU.unsafeIndex` slot) + else go ((slot + 1) .&. mask) +{-# INLINE ciLookup #-} + +{- | Smallest power of two strictly greater than @n@, at least 2. +Used to size the open-addressing table so the load factor stays below ~0.5. +-} +nextPow2Above :: Int -> Int +nextPow2Above n = go 2 + where + go !p + | p > n = p + | otherwise = go (p * 2) + {- | Build a compact index from a vector of row hashes. -Sorts @(hash, originalIndex)@ pairs by hash, then scans for -contiguous runs to populate the offset map. +Sorts @(hash, originalIndex)@ pairs by hash, scans for contiguous runs, then +inserts each run into an open-addressing linear-probe table sized to keep the +load factor under ~0.5. -} buildCompactIndex :: VU.Vector Int -> CompactIndex buildCompactIndex hashes = let n = VU.length hashes - (sortedHashes, sortedIndices) = sortWithIndices hashes - !offs = buildOffsets sortedHashes n 0 HM.empty - in CompactIndex sortedIndices offs - where - buildOffsets :: - VU.Vector Int -> - Int -> - Int -> - HM.HashMap Int (Int, Int) -> - HM.HashMap Int (Int, Int) - buildOffsets !sh !n !i !acc - | i >= n = acc - | otherwise = - let !h = sh `VU.unsafeIndex` i - !end = findGroupEnd sh h (i + 1) n - in buildOffsets sh n end (HM.insert h (i, end - i) acc) + (sortedHashes, sortedIndices) = parSortByHash n hashes + -- Worst case every row is a distinct group; 2*n+1 keeps the table + -- sparse even then. Capacity is independent of the actual group count + -- so building stays a single pass with no resize. + !cap = nextPow2Above (2 * n) + !mask = cap - 1 + (keys, starts, lens) = runST $ do + mKeys <- VUM.unsafeNew cap + mStarts <- VUM.replicate cap (-1) + mLens <- VUM.unsafeNew cap + 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) + insert end + probe !h !slot !start !len = do + s <- VUM.unsafeRead mStarts slot + if s < 0 + then do + VUM.unsafeWrite mKeys slot h + VUM.unsafeWrite mStarts slot start + VUM.unsafeWrite mLens slot 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 -- | Find the end of a contiguous run of equal values starting at @j@. findGroupEnd :: VU.Vector Int -> Int -> Int -> Int -> Int @@ -203,19 +267,57 @@ rightHashes = D.computeRowHashes rightKeyIdxs right buildRows = min leftRows rightRows + -- Probe with the larger side, build on the smaller. The probe side + -- drives parallelism (it is partitioned across cores), so probing the + -- larger side maximises the win. + probeRows = max leftRows rightRows + -- Parallelize the probe either when the build is large enough to spill + -- cache (the original sort-merge regime) or when the build is small + -- (cache-resident, read-only) but the probe is huge: partitioning a + -- ~1e7-row probe over cores wins even against a tiny hot build (the + -- medium-factor lever). Both routes use the same bit-identical + -- 'parInnerKernel'. + useParallel = + shouldParallelizeJoin probeRows buildRows + || shouldParallelizeSmallBuildProbe probeRows (leftIxs, rightIxs) - | buildRows > joinStrategyThreshold = + | buildRows > joinStrategyThreshold && not useParallel = sortMergeInnerKernel leftHashes rightHashes | rightRows <= leftRows = -- Build on right (smaller or equal), probe with left - hashInnerKernel leftHashes rightHashes + innerKernel useParallel leftHashes rightHashes | otherwise = -- Build on left (smaller), probe with right, swap result - let (!rIxs, !lIxs) = hashInnerKernel rightHashes leftHashes + let (!rIxs, !lIxs) = innerKernel useParallel rightHashes leftHashes in (lIxs, rIxs) in assembleInner csSet left right leftIxs rightIxs +{- | Select the inner-join kernel: parallel chunked-probe hash join when the +caller decided the probe side is large enough and multi-core, otherwise the +sequential single-threaded hash kernel. +-} +innerKernel :: + Bool -> VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) +innerKernel True = parInnerKernel +innerKernel False = hashInnerKernel +{-# INLINE innerKernel #-} + +{- | Parallel inner-join kernel: build the 'CompactIndex' on @buildHashes@ once, +then probe @probeHashes@ in parallel. Bit-for-bit identical output to +'hashInnerKernel' (probe-row order preserved). Runs the IO probe via +'unsafePerformIO'; the computation is pure (no observable effects, fixed result +for fixed inputs). +-} +parInnerKernel :: + VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) +parInnerKernel probeHashes buildHashes = + let !ci = buildCompactIndex buildHashes + (pf, bf) = + unsafePerformIO (parInnerProbe (ciSortedIndices ci) (ciLookup ci) probeHashes) + in (pf, bf) +{-# NOINLINE parInnerKernel #-} + -- | Compute hashes for the given key column names in a DataFrame. buildHashColumn :: [T.Text] -> DataFrame -> VU.Vector Int buildHashColumn keys df = @@ -236,7 +338,6 @@ (VU.Vector Int, VU.Vector Int) hashProbeKernel ci probeHashes = let ciIxs = ciSortedIndices ci - ciOff = ciOffsets ci (pFrozen, bFrozen) = runST $ do let !probeN = VU.length probeHashes initCap = max 1 (min probeN 1_000_000) @@ -265,9 +366,10 @@ | i >= probeN = return () | otherwise = do let !h = probeHashes `VU.unsafeIndex` i - case HM.lookup h ciOff of - Nothing -> go (i + 1) - Just (!start, !len) -> do + (!start, !len) = ciLookup ci h + if start < 0 + then go (i + 1) + else do !p <- readSTRef posRef ensureCapacity (p + len) pv <- readSTRef pvRef @@ -311,7 +413,6 @@ let (pFrozen, bFrozen) = runST $ do let ci = buildCompactIndex buildHashes ciIxs = ciSortedIndices ci - ciOff = ciOffsets ci !probeN = VU.length probeHashes !buildN = VU.length buildHashes initCap = max 1 (min (probeN + buildN) 1_000_000) @@ -340,9 +441,10 @@ | i >= probeN = return () | otherwise = do let !h = probeHashes `VU.unsafeIndex` i - case HM.lookup h ciOff of - Nothing -> go (i + 1) - Just (!start, !len) -> do + (!start, !len) = ciLookup ci h + if start < 0 + then go (i + 1) + else do !p <- readSTRef posRef when (p + len > maxJoinOutputRows) $ error $ @@ -566,16 +668,36 @@ leftHashes = D.computeRowHashes leftKeyIdxs left rightHashes = D.computeRowHashes rightKeyIdxs right - -- Right is always the build side for left join + leftRows = fst (D.dimensions left) + -- Right is always the build side for left join; left is the probe side + -- (and drives parallelism). Parallelize either in the large-build regime + -- or when the build is small but the probe (left) is huge: the + -- read-only shared index is probed across cores with no synchronization. + useParallel = + shouldParallelizeJoin leftRows rightRows + || shouldParallelizeSmallBuildProbe leftRows (leftIxs, rightIxs) - | rightRows > joinStrategyThreshold = + | rightRows > joinStrategyThreshold && not useParallel = sortMergeLeftKernel leftHashes rightHashes + | useParallel = + parLeftKernel leftHashes rightHashes | otherwise = hashLeftKernel leftHashes rightHashes in -- rightIxs uses -1 as sentinel for "no match" assembleLeft csSet left right leftIxs rightIxs +{- | Parallel left-join kernel: build the 'CompactIndex' on @rightHashes@ once, +then probe @leftHashes@ in parallel. Bit-for-bit identical output to +'hashLeftKernel' (left-row order preserved, @-1@ sentinel for misses). +-} +parLeftKernel :: + 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) +{-# NOINLINE parLeftKernel #-} + {- | Hash-based left join kernel. Returns @(leftExpandedIndices, rightExpandedIndices)@ where right indices use @-1@ as sentinel for unmatched rows. @@ -587,7 +709,6 @@ hashLeftKernel leftHashes rightHashes = runST $ do let ci = buildCompactIndex rightHashes ciIxs = ciSortedIndices ci - ciOff = ciOffsets ci !leftN = VU.length leftHashes !rightN = VU.length rightHashes initCap = max 1 (min (leftN + rightN) 1_000_000) @@ -616,16 +737,17 @@ | i >= leftN = return () | otherwise = do let !h = leftHashes `VU.unsafeIndex` i + (!start, !len) = ciLookup ci h !p <- readSTRef posRef - case HM.lookup h ciOff of - Nothing -> do + if start < 0 + then do ensureCapacity (p + 1) lv <- readSTRef lvRef rv <- readSTRef rvRef VUM.unsafeWrite lv p i VUM.unsafeWrite rv p (-1) writeSTRef posRef (p + 1) - Just (!start, !len) -> do + else do ensureCapacity (p + len) lv <- readSTRef lvRef rv <- readSTRef rvRef @@ -864,89 +986,111 @@ hashFullOuterKernel leftHashes rightHashes = runST $ do let leftCI = buildCompactIndex leftHashes rightCI = buildCompactIndex rightHashes - leftOff = ciOffsets leftCI - rightOff = ciOffsets rightCI leftSI = ciSortedIndices leftCI rightSI = ciSortedIndices rightCI - - -- Count: matched + left-only + right-only - let leftEntries = HM.toList leftOff - rightEntries = HM.toList rightOff - - !matchedCount = - foldl' - ( \acc (h, (_, ll)) -> - case HM.lookup h rightOff of - Nothing -> acc - Just (_, rl) -> acc + ll * rl - ) - 0 - leftEntries - - !leftOnlyCount = - foldl' - ( \acc (h, (_, ll)) -> - if HM.member h rightOff then acc else acc + ll - ) - 0 - leftEntries - - !rightOnlyCount = - foldl' - ( \acc (h, (_, rl)) -> - if HM.member h leftOff then acc else acc + rl - ) - 0 - rightEntries + leftKeys = ciKeys leftCI + leftStarts = ciStarts leftCI + leftLens = ciLens leftCI + rightKeys = ciKeys rightCI + rightStarts = ciStarts rightCI + rightLens = ciLens rightCI + !leftCap = VU.length leftStarts + !rightCap = VU.length rightStarts - !totalCount = matchedCount + leftOnlyCount + rightOnlyCount + -- Count: matched + left-only + right-only. Iterate over occupied slots + -- (ciStarts /= -1) in each table, cross-referencing the other via ciLookup. + let countLeft !slot !acc + | slot >= leftCap = acc + | otherwise = + let !lStart = leftStarts `VU.unsafeIndex` slot + in if lStart < 0 + then countLeft (slot + 1) acc + else + let !h = leftKeys `VU.unsafeIndex` slot + !ll = leftLens `VU.unsafeIndex` slot + (!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 + then countRightOnly (slot + 1) acc + else + let !h = rightKeys `VU.unsafeIndex` slot + !rl = rightLens `VU.unsafeIndex` slot + (!ls, _) = ciLookup leftCI h + in countRightOnly (slot + 1) (acc + if ls < 0 then rl else 0) + !leftPlusMatched = countLeft 0 0 + !rightOnlyCount = countRightOnly 0 0 + !totalCount = leftPlusMatched + rightOnlyCount lv <- VUM.unsafeNew totalCount rv <- VUM.unsafeNew totalCount posRef <- newSTRef (0 :: Int) - -- Fill matched + left-only (iterate left keys) - forM_ leftEntries $ \(h, (lStart, lLen)) -> do - !p <- readSTRef posRef - case HM.lookup h rightOff of - Nothing -> do - -- Left-only rows - let goL !j !q - | j >= lLen = return () - | otherwise = do - VUM.unsafeWrite lv q (leftSI `VU.unsafeIndex` (lStart + j)) - VUM.unsafeWrite rv q (-1) - goL (j + 1) (q + 1) - goL 0 p - writeSTRef posRef (p + lLen) - Just (!rStart, !rLen) -> do - -- Cross product - fillCrossProduct - leftSI - rightSI - lStart - (lStart + lLen) - rStart - (rStart + rLen) - lv - rv - p - writeSTRef posRef (p + lLen * rLen) + -- Fill matched + left-only (iterate left slots) + let fillLeft !slot + | slot >= leftCap = return () + | otherwise = do + let !lStart = leftStarts `VU.unsafeIndex` slot + if lStart < 0 + then fillLeft (slot + 1) + else do + let !h = leftKeys `VU.unsafeIndex` slot + !lLen = leftLens `VU.unsafeIndex` slot + (!rStart, !rLen) = ciLookup rightCI h + !p <- readSTRef posRef + if rStart < 0 + then do + let goL !j !q + | j >= lLen = return () + | otherwise = do + VUM.unsafeWrite lv q (leftSI `VU.unsafeIndex` (lStart + j)) + VUM.unsafeWrite rv q (-1) + goL (j + 1) (q + 1) + goL 0 p + writeSTRef posRef (p + lLen) + else do + fillCrossProduct + leftSI + rightSI + lStart + (lStart + lLen) + rStart + (rStart + rLen) + lv + rv + p + writeSTRef posRef (p + lLen * rLen) + fillLeft (slot + 1) + fillLeft 0 - -- Fill right-only (iterate right keys not in left) - forM_ rightEntries $ \(h, (rStart, rLen)) -> - case HM.lookup h leftOff of - Just _ -> return () - Nothing -> do - !p <- readSTRef posRef - let goR !j !q - | j >= rLen = return () - | otherwise = do - VUM.unsafeWrite lv q (-1) - VUM.unsafeWrite rv q (rightSI `VU.unsafeIndex` (rStart + j)) - goR (j + 1) (q + 1) - goR 0 p - writeSTRef posRef (p + rLen) + -- Fill right-only (iterate right slots not in left) + let fillRightOnly !slot + | slot >= rightCap = return () + | otherwise = do + let !rStart = rightStarts `VU.unsafeIndex` slot + if rStart < 0 + then fillRightOnly (slot + 1) + else do + let !h = rightKeys `VU.unsafeIndex` slot + !rLen = rightLens `VU.unsafeIndex` slot + (!ls, _) = ciLookup leftCI h + if ls >= 0 + then fillRightOnly (slot + 1) + else do + !p <- readSTRef posRef + let goR !j !q + | j >= rLen = return () + | otherwise = do + VUM.unsafeWrite lv q (-1) + VUM.unsafeWrite rv q (rightSI `VU.unsafeIndex` (rStart + j)) + goR (j + 1) (q + 1) + goR 0 p + writeSTRef posRef (p + rLen) + fillRightOnly (slot + 1) + fillRightOnly 0 (,) <$> VU.unsafeFreeze lv <*> VU.unsafeFreeze rv @@ -1060,6 +1204,9 @@ -- Coalesce two nullable columns: take first non-Nothing per row, -- producing a non-optional column. coalesceKeyColumn :: Column -> Column -> Column + coalesceKeyColumn l r + | D.isPackedText l || D.isPackedText r = + coalesceKeyColumn (D.materializePacked l) (D.materializePacked r) coalesceKeyColumn (BoxedColumn lBm (lCol :: VB.Vector a)) (BoxedColumn rBm (rCol :: VB.Vector b)) = diff --git a/src/DataFrame/Operations/JoinPar.hs b/src/DataFrame/Operations/JoinPar.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/JoinPar.hs @@ -0,0 +1,215 @@ +{-# 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/Permutation.hs b/src/DataFrame/Operations/Permutation.hs --- a/src/DataFrame/Operations/Permutation.hs +++ b/src/DataFrame/Operations/Permutation.hs @@ -27,6 +27,7 @@ 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) @@ -117,6 +118,12 @@ UnboxedColumn _ (v :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of Just Refl -> \i j -> compare (v `VU.unsafeIndex` i) (v `VU.unsafeIndex` j) Nothing -> \_ _ -> EQ + PackedText _ p -> case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> \i j -> + let (ai, oi, li) = packedSlice p i + (aj, oj, lj) = packedSlice p j + in sliceCmpBytes ai oi li aj oj lj + Nothing -> \_ _ -> EQ sortOrderComparator (Desc (Col name :: Expr a)) df = case unsafeGetColumn name df of BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of @@ -124,6 +131,12 @@ Nothing -> \_ _ -> EQ UnboxedColumn _ (v :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of Just Refl -> \i j -> compare (v `VU.unsafeIndex` j) (v `VU.unsafeIndex` i) + Nothing -> \_ _ -> EQ + PackedText _ p -> case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> \i j -> + let (ai, oi, li) = packedSlice p i + (aj, oj, lj) = packedSlice p j + in sliceCmpBytes aj oj lj ai oi li Nothing -> \_ _ -> EQ sortOrderComparator _ _ = error "Sorting on compound column" 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 @@ -248,6 +248,7 @@ Just ((BoxedColumn _ (column :: V.Vector a'))) -> case testEquality (typeRep @a') (typeRep @a) of Just Refl -> VG.sum column Nothing -> 0 + Just (PackedText _ _) -> 0 sum expr df = case interpret df expr of Left e -> throw e Right (TColumn xs) -> case toVector @a @V.Vector xs of 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 @@ -45,6 +45,7 @@ ) import DataFrame.Internal.Expression import DataFrame.Internal.Interpreter +import DataFrame.Internal.PackedText (packedIndexText, packedLength) import DataFrame.Operations.Core () import DataFrame.Operations.Merge () import DataFrame.Operations.Transformations (apply) @@ -130,10 +131,12 @@ -- | Dataframe to filter DataFrame -> DataFrame -filter (Col filterColumnName) condition df = case getColumn filterColumnName df of +filter e@(Col filterColumnName) condition df = case getColumn filterColumnName df of Nothing -> throw $ ColumnsNotFoundException [filterColumnName] "filter" (M.keys $ columnIndices df) + Just c@(PackedText _ _) -> + filter e condition (insertColumn filterColumnName (materializePacked c) df) Just _col@(BoxedColumn bm (column :: V.Vector b)) -> -- Check direct type match first, then try Maybe b match for nullable columns case testEquality (typeRep @a) (typeRep @b) of @@ -240,6 +243,8 @@ filter (Col @(Maybe a) colName) isJust df & apply @(Maybe a) fromJust colName UnboxedColumn (Just _) (_col :: VU.Vector a) -> filter (Col @(Maybe a) colName) isJust df & apply @(Maybe a) fromJust colName + c@(PackedText (Just _) _) -> + filterJust colName (insertColumn colName (materializePacked c) df) _ -> df Just _ -> df @@ -255,6 +260,8 @@ Just column | hasMissing column -> case column of BoxedColumn (Just _) (_col :: V.Vector a) -> filter (Col @(Maybe a) colName) isNothing df UnboxedColumn (Just _) (_col :: VU.Vector a) -> filter (Col @(Maybe a) colName) isNothing df + c@(PackedText (Just _) _) -> + filterNothing colName (insertColumn colName (materializePacked c) df) _ -> df _ -> df @@ -484,6 +491,10 @@ 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) 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 @@ -6,7 +6,13 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -module DataFrame.Operations.Typing where +-- The inference lattice (sample classification, 'ParsingAssumption', +-- Int -> Double promotion) lives in "DataFrame.Operations.Inference" +-- and is re-exported here for backwards compatibility. +module DataFrame.Operations.Typing ( + module DataFrame.Operations.Typing, + module DataFrame.Operations.Inference, +) where import qualified Data.Map as M import qualified Data.Text as T @@ -30,6 +36,7 @@ ensureOptional, finalizeParseResult, fromVector, + materializePacked, ) import DataFrame.Internal.DataFrame ( DataFrame (..), @@ -39,11 +46,10 @@ import DataFrame.Internal.Parsing import DataFrame.Internal.Schema import DataFrame.Operations.Core () +import DataFrame.Operations.Inference import Text.Read import Type.Reflection -type DateFormat = String - {- | How parse failures are surfaced in the resulting column. * 'NoSafeRead' — strict parsing: failures throw (via 'read'). @@ -228,14 +234,16 @@ (parseUnboxedColumnWithPred False isNull readBool cols) (handleTextAssumption isNull cols) +{- | Int columns: one fused pass with in-place Int -> Double promotion +('promoteIntColumn'); a cell parsing as neither demotes to Text over +the retained raw cells. 'readIntStrict' rejects overflow so a huge +integer promotes to its true 'Double' value instead of wrapping. +-} handleIntAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column handleIntAssumption isNull cols = - case parseUnboxedColumnWithPred 0 isNull readInt cols of - Just (mbm, vec) -> UnboxedColumn mbm vec - Nothing -> - unboxedOrFallback - (parseUnboxedColumnWithPred 0 isNull readDouble cols) - (handleTextAssumption isNull cols) + case promoteIntColumn (\_ t -> isNull t) readIntStrict readDouble cols of + Just col -> col + Nothing -> handleTextAssumption isNull cols handleDoubleAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column handleDoubleAssumption isNull cols = @@ -334,14 +342,6 @@ convertOnlyEmpty :: T.Text -> Maybe T.Text convertOnlyEmpty v = if v == "" then Nothing else Just v -parseTimeOpt :: DateFormat -> T.Text -> Maybe Day -parseTimeOpt dateFormat s = - parseTimeM {- Accept leading/trailing whitespace -} - True - defaultTimeLocale - dateFormat - (T.unpack s) - unsafeParseTime :: DateFormat -> T.Text -> Day unsafeParseTime dateFormat s = parseTimeOrError {- Accept leading/trailing whitespace -} @@ -361,40 +361,6 @@ hasSameConstructor Nothing Nothing = True hasSameConstructor _ _ = False -makeParsingAssumption :: - DateFormat -> V.Vector (Maybe T.Text) -> ParsingAssumption -makeParsingAssumption dateFormat asMaybeText - -- All the examples are "NA", "Null", "", so we can't make any shortcut - -- assumptions and just have to go the long way. - | V.all (== Nothing) asMaybeText = NoAssumption - -- After accounting for nulls, parsing for Ints and Doubles results in the - -- same corresponding positions of Justs and Nothings, so we assume - -- that the best way to parse is Int - | vecSameConstructor asMaybeText asMaybeBool = BoolAssumption - | vecSameConstructor asMaybeText asMaybeInt - && vecSameConstructor asMaybeText asMaybeDouble = - IntAssumption - -- After accounting for nulls, the previous condition fails, so some (or none) can be parsed as Ints - -- and some can be parsed as Doubles, so we make the assumpotion of doubles. - | vecSameConstructor asMaybeText asMaybeDouble = DoubleAssumption - -- After accounting for nulls, parsing for Dates results in the same corresponding - -- positions of Justs and Nothings, so we assume that the best way to parse is Date. - | vecSameConstructor asMaybeText asMaybeDate = DateAssumption - | otherwise = TextAssumption - where - asMaybeBool = V.map (>>= readBool) asMaybeText - asMaybeInt = V.map (>>= readInt) asMaybeText - asMaybeDouble = V.map (>>= readDouble) asMaybeText - asMaybeDate = V.map (>>= parseTimeOpt dateFormat) asMaybeText - -data ParsingAssumption - = BoolAssumption - | IntAssumption - | DoubleAssumption - | DateAssumption - | NoAssumption - | TextAssumption - {- | Re-type columns of a 'DataFrame' according to the supplied schema map. The caller provides a @resolveMode@ function that maps a column name to its 'SafeReadMode' — typically built from a global default plus an overrides map @@ -426,6 +392,9 @@ EitherRead -> fromVector (V.map ((readEitherRaw @a) . toStr) col) asType :: SafeReadMode -> SchemaType -> Column -> Column + -- A raw CSV string column may arrive as PackedText; decode to boxed Text + -- so the re-parse arms below can read the cells. + asType mode st c@(PackedText _ _) = asType mode st (materializePacked c) asType mode (SType (_ :: P.Proxy a)) c@(BoxedColumn _ (col :: V.Vector b)) = case typeRep @a of App t1 _t2 -> case eqTypeRep t1 (typeRep @Maybe) of Just HRefl -> case testEquality (typeRep @a) (typeRep @b) of diff --git a/src/DataFrame/Typed/Join.hs b/src/DataFrame/Typed/Join.hs --- a/src/DataFrame/Typed/Join.hs +++ b/src/DataFrame/Typed/Join.hs @@ -26,7 +26,11 @@ -- | Typed inner join on one or more key columns. innerJoin :: forall (keys :: [Symbol]) left right. - (AllKnownSymbol keys) => + ( AllKnownSymbol keys + , AssertAllPresent keys left + , AssertAllPresent keys right + , AssertKeyTypesMatch keys left right + ) => TypedDataFrame left -> TypedDataFrame right -> TypedDataFrame (InnerJoinSchema keys left right) @@ -38,7 +42,11 @@ -- | Typed left join. leftJoin :: forall (keys :: [Symbol]) left right. - (AllKnownSymbol keys) => + ( AllKnownSymbol keys + , AssertAllPresent keys left + , AssertAllPresent keys right + , AssertKeyTypesMatch keys left right + ) => TypedDataFrame left -> TypedDataFrame right -> TypedDataFrame (LeftJoinSchema keys left right) @@ -50,7 +58,11 @@ -- | Typed right join. rightJoin :: forall (keys :: [Symbol]) left right. - (AllKnownSymbol keys) => + ( AllKnownSymbol keys + , AssertAllPresent keys left + , AssertAllPresent keys right + , AssertKeyTypesMatch keys left right + ) => TypedDataFrame left -> TypedDataFrame right -> TypedDataFrame (RightJoinSchema keys left right) @@ -62,7 +74,11 @@ -- | Typed full outer join. fullOuterJoin :: forall (keys :: [Symbol]) left right. - (AllKnownSymbol keys) => + ( AllKnownSymbol keys + , AssertAllPresent keys left + , AssertAllPresent keys right + , AssertKeyTypesMatch keys left right + ) => TypedDataFrame left -> TypedDataFrame right -> TypedDataFrame (FullOuterJoinSchema keys left right)