{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE TypeApplications #-}

module DataFrame.Internal.Grouping (
    groupBy,
    groupBySeq,
    groupByPar,
    buildRowToGroup,
    changingPoints,
) where

import qualified Data.List as L
import qualified Data.Map as M
import qualified Data.Text as T
import qualified Data.Vector as V
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.Column (
    Bitmap,
    Column (..),
    bitmapTestBit,
    materializeMerged,
 )
import DataFrame.Internal.DataFrame (DataFrame (..), GroupedDataFrame (..))
import DataFrame.Internal.DictEncode (dictEncodeColumnUpTo)
import DataFrame.Internal.GroupingDirect (
    DirectGrouping (..),
    directGroupThreshold,
    tryDirectGroupColumn,
 )
import DataFrame.Internal.GroupingPar (parallelAssignGroups, shouldParallelize)
import DataFrame.Internal.Hash
import DataFrame.Internal.HashTable (htInsert, newHashTable)
import DataFrame.Internal.PackedText (
    PackedSel,
    PackedTextData (..),
    offAt,
    offCount,
    packedLength,
    packedSlice,
    selAt,
    selLength,
    sliceEqBytes,
 )
import DataFrame.Internal.RadixRank (rankByHash)
import DataFrame.Internal.Types
import System.IO.Unsafe (unsafePerformIO)
import Type.Reflection (typeRep)

{- | O(k * n) group the dataframe by the given key columns, bucketing rows with an
open-addressing hash table that re-verifies keys on each hash hit. Groups are
numbered in first-appearance order; 'valueIndices'/'offsets' follow by counting sort.
-}
groupBy ::
    [T.Text] ->
    DataFrame ->
    GroupedDataFrame
groupBy :: [Text] -> DataFrame -> GroupedDataFrame
groupBy [Text]
names DataFrame
df
    | (Text -> Bool) -> [Text] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Text -> [Text] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` DataFrame -> [Text]
columnNames DataFrame
df) [Text]
names =
        DataFrameException -> GroupedDataFrame
forall a e. Exception e => e -> a
throw (DataFrameException -> GroupedDataFrame)
-> DataFrameException -> GroupedDataFrame
forall a b. (a -> b) -> a -> b
$
            [Text] -> Text -> [Text] -> DataFrameException
ColumnsNotFoundException
                ([Text]
names [Text] -> [Text] -> [Text]
forall a. Eq a => [a] -> [a] -> [a]
L.\\ DataFrame -> [Text]
columnNames DataFrame
df)
                Text
"groupBy"
                (DataFrame -> [Text]
columnNames DataFrame
df)
    | DataFrame -> Int
nRows DataFrame
df Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 =
        DataFrame
-> [Text]
-> Vector Int
-> Vector Int
-> Vector Int
-> GroupedDataFrame
Grouped
            DataFrame
df
            [Text]
names
            Vector Int
forall a. Unbox a => Vector a
VU.empty
            ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList [Int
0])
            Vector Int
forall a. Unbox a => Vector a
VU.empty
    | Just GroupedDataFrame
dg <- [Text] -> DataFrame -> Maybe GroupedDataFrame
tryDirectGroup [Text]
names DataFrame
df = GroupedDataFrame
dg
    | Int -> Bool
shouldParallelize Int
n = [Text] -> DataFrame -> GroupedDataFrame
groupByPar [Text]
names DataFrame
df
    | Bool
otherwise = [Text] -> DataFrame -> GroupedDataFrame
groupBySeq [Text]
names DataFrame
df
  where
    !n :: Int
n = DataFrame -> Int
nRows DataFrame
df

{- | Low-cardinality direct-indexed grouping fast path
('DataFrame.Internal.GroupingDirect'): fires only for a single clean small-range
@Int@ key. Returns 'Nothing' on any other key shape, falling back to the hash path.
-}
tryDirectGroup :: [T.Text] -> DataFrame -> Maybe GroupedDataFrame
tryDirectGroup :: [Text] -> DataFrame -> Maybe GroupedDataFrame
tryDirectGroup [Text
name] DataFrame
df = do
    Column
col <- Text -> Map Text Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Text
name (DataFrame -> Map Text Int
columnIndices DataFrame
df) Maybe Int -> (Int -> Maybe Column) -> Maybe Column
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Int
i -> DataFrame -> Vector Column
columns DataFrame
df Vector Column -> Int -> Maybe Column
forall a. Vector a -> Int -> Maybe a
V.!? Int
i
    case Column -> Maybe DirectGrouping
tryDirectGroupColumn Column
col of
        Just DirectGrouping
dg ->
            GroupedDataFrame -> Maybe GroupedDataFrame
forall a. a -> Maybe a
Just (DataFrame
-> [Text]
-> Vector Int
-> Vector Int
-> Vector Int
-> GroupedDataFrame
Grouped DataFrame
df [Text
name] (DirectGrouping -> Vector Int
dgValueIndices DirectGrouping
dg) (DirectGrouping -> Vector Int
dgOffsets DirectGrouping
dg) (DirectGrouping -> Vector Int
dgRowToGroup DirectGrouping
dg))
        Maybe DirectGrouping
Nothing -> case Column
col of
            PackedText Maybe Bitmap
Nothing PackedTextData
p -> DataFrame -> [Text] -> PackedTextData -> Maybe GroupedDataFrame
dictCodesGroup DataFrame
df [Text
name] PackedTextData
p
            Column
_ -> Int -> DataFrame -> [Text] -> Column -> Maybe GroupedDataFrame
tryDictGroup (DataFrame -> Int
nRows DataFrame
df) DataFrame
df [Text
name] Column
col
tryDirectGroup [Text]
_ DataFrame
_ = Maybe GroupedDataFrame
forall a. Maybe a
Nothing

dictCodesGroup ::
    DataFrame -> [T.Text] -> PackedTextData -> Maybe GroupedDataFrame
dictCodesGroup :: DataFrame -> [Text] -> PackedTextData -> Maybe GroupedDataFrame
dictCodesGroup DataFrame
df [Text]
names PackedTextData
p = do
    PackedSel
sel <- PackedTextData -> Maybe PackedSel
ptSel PackedTextData
p
    Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (PackedTextData -> Bool
ptCanonicalSel PackedTextData
p)
    let offs :: PackedOffsets
offs = PackedTextData -> PackedOffsets
ptOffsets PackedTextData
p
        card :: Int
card = PackedOffsets -> Int
offCount PackedOffsets
offs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        n :: Int
n = PackedSel -> Int
selLength PackedSel
sel
    Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
card Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
card Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
directGroupThreshold Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0)
    Vector Int
counts <- Int -> PackedSel -> Maybe (Vector Int)
codeHistogram Int
card PackedSel
sel
    let occupied :: Vector Int
occupied = (Int -> Bool) -> Vector Int -> Vector Int
forall a. Unbox a => (a -> Bool) -> Vector a -> Vector a
VU.filter (\Int
g -> Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
counts Int
g Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (Int -> Int -> Vector Int
forall a. (Unbox a, Num a) => a -> Int -> Vector a
VU.enumFromN Int
0 Int
card)
        nGroups :: Int
nGroups = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
occupied
        entryHash :: Int -> Int
entryHash Int
g =
            let o :: Int
o = PackedOffsets -> Int -> Int
offAt PackedOffsets
offs Int
g
             in Int -> Array -> Int -> Int -> Int
mixBytes Int
fnvOffset (PackedTextData -> Array
ptBytes PackedTextData
p) Int
o (PackedOffsets -> Int -> Int
offAt PackedOffsets
offs (Int
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
o)
        rank :: Vector Int
rank =
            (forall s. ST s (Vector Int)) -> Vector Int
forall a. (forall s. ST s a) -> a
runST ((Int -> ST s Int) -> Int -> ST s (Vector Int)
forall (m :: * -> *).
PrimMonad m =>
(Int -> m Int) -> Int -> m (Vector Int)
rankByHash (Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> (Int -> Int) -> Int -> ST s Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int
entryHash (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
occupied) Int
nGroups)
        remap :: Vector Int
remap = (forall s. ST s (Vector Int)) -> Vector Int
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int)) -> Vector Int)
-> (forall s. ST s (Vector Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
            MVector s Int
m <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
card
            (Int -> Int -> ST s ()) -> Vector Int -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
(Int -> a -> m b) -> Vector a -> m ()
VU.imapM_ (\Int
j Int
g -> MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
m Int
g (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
rank Int
j)) Vector Int
occupied
            MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
m
        rtg :: Vector Int
rtg = Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
remap (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PackedSel -> Int -> Int
selAt PackedSel
sel)
        (Vector Int
vis, Vector Int
os) = Vector Int -> Int -> (Vector Int, Vector Int)
indicesFromGroups Vector Int
rtg Int
nGroups
    GroupedDataFrame -> Maybe GroupedDataFrame
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DataFrame
-> [Text]
-> Vector Int
-> Vector Int
-> Vector Int
-> GroupedDataFrame
Grouped DataFrame
df [Text]
names Vector Int
vis Vector Int
os Vector Int
rtg)

-- | Per-code occupancy counts; 'Nothing' as soon as any code is negative.
codeHistogram :: Int -> PackedSel -> Maybe (VU.Vector Int)
codeHistogram :: Int -> PackedSel -> Maybe (Vector Int)
codeHistogram Int
card PackedSel
sel = (forall s. ST s (Maybe (Vector Int))) -> Maybe (Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Maybe (Vector Int))) -> Maybe (Vector Int))
-> (forall s. ST s (Maybe (Vector Int))) -> Maybe (Vector Int)
forall a b. (a -> b) -> a -> b
$ do
    MVector s Int
counts <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
card Int
0
    let n :: Int
n = PackedSel -> Int
selLength PackedSel
sel
        go :: Int -> ST s (Maybe (Vector Int))
go Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Vector Int -> Maybe (Vector Int)
forall a. a -> Maybe a
Just (Vector Int -> Maybe (Vector Int))
-> ST s (Vector Int) -> ST s (Maybe (Vector Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
counts
            | Bool
otherwise = do
                let c :: Int
c = PackedSel -> Int -> Int
selAt PackedSel
sel Int
i
                if Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
card
                    then Maybe (Vector Int) -> ST s (Maybe (Vector Int))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Vector Int)
forall a. Maybe a
Nothing
                    else do
                        MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
VUM.unsafeModify MVector s Int
MVector (PrimState (ST s)) Int
counts (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
c
                        Int -> ST s (Maybe (Vector Int))
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> ST s (Maybe (Vector Int))
go Int
0

{- | Dictionary-encode a single text key to dense int codes, then derive
@valueIndices@/@offsets@ by counting sort. Profiled slower than the fused hash
group-by on every db-benchmark question, so it always falls back ('dictGroupEnabled').
-}
tryDictGroup ::
    Int -> DataFrame -> [T.Text] -> Column -> Maybe GroupedDataFrame
tryDictGroup :: Int -> DataFrame -> [Text] -> Column -> Maybe GroupedDataFrame
tryDictGroup Int
n DataFrame
df [Text]
names Column
col
    | Bool
dictGroupEnabled Bool -> Bool -> Bool
&& Bool -> Bool
not (Int -> Bool
shouldParallelize Int
n) = do
        (Vector Int
codes, Int
card) <- Int -> Column -> Maybe (Vector Int, Int)
dictEncodeColumnUpTo Int
dictSingleThreshold Column
col
        let (Vector Int
vis, Vector Int
os) = Vector Int -> Int -> (Vector Int, Vector Int)
indicesFromGroups Vector Int
codes Int
card
        GroupedDataFrame -> Maybe GroupedDataFrame
forall a. a -> Maybe a
Just (DataFrame
-> [Text]
-> Vector Int
-> Vector Int
-> Vector Int
-> GroupedDataFrame
Grouped DataFrame
df [Text]
names Vector Int
vis Vector Int
os Vector Int
codes)
    | Bool
otherwise = Maybe GroupedDataFrame
forall a. Maybe a
Nothing

{- | Master switch for the single-key dict-encode grouping path. 'False' because
it profiled slower than the hash group-by on every db-benchmark group-by question
(see 'tryDictGroup'); the path is kept compiled and tested but not taken.
-}
dictGroupEnabled :: Bool
dictGroupEnabled :: Bool
dictGroupEnabled = Bool
False

{- | Cardinality ceiling for the single-key dict-encode probe: it bails to 'Nothing'
once the distinct count passes this. Only consulted when 'dictGroupEnabled' is 'True'.
-}
dictSingleThreshold :: Int
dictSingleThreshold :: Int
dictSingleThreshold = Int
4096

{- | The sequential grouping path: a single open-addressing table over all rows,
canonically remapped. Always available regardless of capabilities; the parallel
path is verified equal to it by a property test.
-}
groupBySeq :: [T.Text] -> DataFrame -> GroupedDataFrame
groupBySeq :: [Text] -> DataFrame -> GroupedDataFrame
groupBySeq [Text]
names DataFrame
df =
    let !n :: Int
n = DataFrame -> Int
nRows DataFrame
df
        indicesToGroup :: [Int]
indicesToGroup = [Text] -> DataFrame -> [Int]
keyColIndices [Text]
names DataFrame
df
        (Vector Int
rtg0, Vector Int
repHash, Vector Int
repRow) = DataFrame -> [Int] -> Int -> (Vector Int, Vector Int, Vector Int)
assignGroups DataFrame
df [Int]
indicesToGroup Int
n
        !nGroups :: Int
nGroups = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
repHash
        !remap :: Vector Int
remap = Vector Int -> Vector Int -> Vector Int
canonicalRemap Vector Int
repHash Vector Int
repRow
        !rtg :: Vector Int
rtg = (Int -> Int) -> Vector Int -> Vector Int
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
remap) Vector Int
rtg0
        (Vector Int
vis, Vector Int
os) = Vector Int -> Int -> (Vector Int, Vector Int)
indicesFromGroups Vector Int
rtg Int
nGroups
     in DataFrame
-> [Text]
-> Vector Int
-> Vector Int
-> Vector Int
-> GroupedDataFrame
Grouped DataFrame
df [Text]
names Vector Int
vis Vector Int
os Vector Int
rtg

{- | The parallel partitioned grouping path (see 'DataFrame.Internal.GroupingPar'):
forks one task per capability, producing output bit-for-bit identical to
'groupBySeq'. Pure via 'unsafePerformIO' (deterministic thread fan-out only).
-}
groupByPar :: [T.Text] -> DataFrame -> GroupedDataFrame
groupByPar :: [Text] -> DataFrame -> GroupedDataFrame
groupByPar [Text]
names DataFrame
df =
    let !n :: Int
n = DataFrame -> Int
nRows DataFrame
df
        indicesToGroup :: [Int]
indicesToGroup = [Text] -> DataFrame -> [Int]
keyColIndices [Text]
names DataFrame
df
        !hashes :: Vector Int
hashes = (forall s. ST s (Vector Int)) -> Vector Int
forall a. (forall s. ST s a) -> a
runST (DataFrame -> [Int] -> Int -> ST s (Vector Int)
forall s. DataFrame -> [Int] -> Int -> ST s (Vector Int)
computeHashes DataFrame
df [Int]
indicesToGroup Int
n)
        !eqRow :: Int -> Int -> Bool
eqRow = DataFrame -> [Int] -> Int -> Int -> Bool
eqKeyRow DataFrame
df [Int]
indicesToGroup
        (Vector Int
rtg, Vector Int
vis, Vector Int
os) = IO (Vector Int, Vector Int, Vector Int)
-> (Vector Int, Vector Int, Vector Int)
forall a. IO a -> a
unsafePerformIO (Int
-> Vector Int
-> (Int -> Int -> Bool)
-> IO (Vector Int, Vector Int, Vector Int)
parallelAssignGroups Int
n Vector Int
hashes Int -> Int -> Bool
eqRow)
     in DataFrame
-> [Text]
-> Vector Int
-> Vector Int
-> Vector Int
-> GroupedDataFrame
Grouped DataFrame
df [Text]
names Vector Int
vis Vector Int
os Vector Int
rtg
{-# NOINLINE groupByPar #-}

-- | Column indices of the requested key columns, in column order.
keyColIndices :: [T.Text] -> DataFrame -> [Int]
keyColIndices :: [Text] -> DataFrame -> [Int]
keyColIndices [Text]
names DataFrame
df =
    Map Text Int -> [Int]
forall k a. Map k a -> [a]
M.elems (Map Text Int -> [Int]) -> Map Text Int -> [Int]
forall a b. (a -> b) -> a -> b
$ (Text -> Int -> Bool) -> Map Text Int -> Map Text Int
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\Text
k Int
_ -> Text
k Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text]
names) (DataFrame -> Map Text Int
columnIndices DataFrame
df)

{- | Assign every row to a dense group id in first-appearance order. Returns
@(rowToGroup, repHash, repRow)@ — the hash and representative row of each group.
'eqKeyRow' re-verifies the real key on each hash hit so colliding keys stay apart.
-}
assignGroups ::
    DataFrame -> [Int] -> Int -> (VU.Vector Int, VU.Vector Int, VU.Vector Int)
assignGroups :: DataFrame -> [Int] -> Int -> (Vector Int, Vector Int, Vector Int)
assignGroups DataFrame
df [Int]
indicesToGroup Int
n = (forall s. ST s (Vector Int, Vector Int, Vector Int))
-> (Vector Int, Vector Int, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int, Vector Int, Vector Int))
 -> (Vector Int, Vector Int, Vector Int))
-> (forall s. ST s (Vector Int, Vector Int, Vector Int))
-> (Vector Int, Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
    Vector Int
hashes <- DataFrame -> [Int] -> Int -> ST s (Vector Int)
forall s. DataFrame -> [Int] -> Int -> ST s (Vector Int)
computeHashes DataFrame
df [Int]
indicesToGroup Int
n
    let !eqRow :: Int -> Int -> Bool
eqRow = DataFrame -> [Int] -> Int -> Int -> Bool
eqKeyRow DataFrame
df [Int]
indicesToGroup
    HashTable s
ht <- Int -> ST s (HashTable (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (HashTable (PrimState m))
newHashTable Int
n
    MVector s Int
rtg <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
n
    MVector s Int
repHashM <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
n
    MVector s Int
repRowM <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
n
    let go :: Int -> Int -> ST s Int
go !Int
i !Int
next
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
next
            | Bool
otherwise = do
                let !h :: Int
h = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
hashes Int
i
                (Int
gid, Bool
isNew) <- HashTable (PrimState (ST s))
-> (Int -> Int -> Bool) -> Int -> Int -> Int -> ST s (Int, Bool)
forall (m :: * -> *).
PrimMonad m =>
HashTable (PrimState m)
-> (Int -> Int -> Bool) -> Int -> Int -> Int -> m (Int, Bool)
htInsert HashTable s
HashTable (PrimState (ST s))
ht Int -> Int -> Bool
eqRow Int
next Int
i Int
h
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
rtg Int
i Int
gid
                Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
isNew (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                    MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
repHashM Int
next Int
h
                    MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
repRowM Int
next Int
i
                Int -> Int -> ST s Int
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (if Bool
isNew then Int
next Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 else Int
next)
    !Int
nGroups <- Int -> Int -> ST s Int
go Int
0 Int
0
    Vector Int
frozen <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
rtg
    Vector Int
repHash <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze (Int -> Int -> MVector s Int -> MVector s Int
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
VUM.slice Int
0 Int
nGroups MVector s Int
repHashM)
    Vector Int
repRow <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze (Int -> Int -> MVector s Int -> MVector s Int
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
VUM.slice Int
0 Int
nGroups MVector s Int
repRowM)
    (Vector Int, Vector Int, Vector Int)
-> ST s (Vector Int, Vector Int, Vector Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector Int
frozen, Vector Int
repHash, Vector Int
repRow)

{- | Map each first-appearance group id to its canonical id: groups ordered by
ascending representative hash (tie-broken by representative row), making group
order a deterministic function of the key set so set ops commute. O(g), no sort.
-}
canonicalRemap :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
canonicalRemap :: Vector Int -> Vector Int -> Vector Int
canonicalRemap Vector Int
repHash Vector Int
_repRow =
    (forall s. ST s (Vector Int)) -> Vector Int
forall a. (forall s. ST s a) -> a
runST ((Int -> ST s Int) -> Int -> ST s (Vector Int)
forall (m :: * -> *).
PrimMonad m =>
(Int -> m Int) -> Int -> m (Vector Int)
rankByHash (Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> (Int -> Int) -> Int -> ST s Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
repHash) (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
repHash))

{- | Compute the FNV row-hash of the key columns into a fresh unboxed vector,
mixing 'nullSalt' for null slots so a missing value never collides with a
present one of the same bits.
-}
computeHashes :: DataFrame -> [Int] -> Int -> ST s (VU.Vector Int)
computeHashes :: forall s. DataFrame -> [Int] -> Int -> ST s (Vector Int)
computeHashes DataFrame
df [Int]
indicesToGroup Int
n = do
    MVector s Int
mh <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
n Int
fnvOffset
    -- Merged key columns are exotic; hash their eager form.
    let selectedCols :: [Column]
selectedCols = (Int -> Column) -> [Int] -> [Column]
forall a b. (a -> b) -> [a] -> [b]
map (Column -> Column
materializeMerged (Column -> Column) -> (Int -> Column) -> Int -> Column
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (DataFrame -> Vector Column
columns DataFrame
df Vector Column -> Int -> Column
forall a. Vector a -> Int -> a
V.!)) [Int]
indicesToGroup
    [Column] -> (Column -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Column]
selectedCols ((Column -> ST s ()) -> ST s ()) -> (Column -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \case
        UnboxedColumn Maybe Bitmap
ubm (Vector a
v :: VU.Vector a) ->
            case TypeRep a -> TypeRep Int -> Maybe (a :~: Int)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Int) of
                Just a :~: Int
Refl -> MVector s Int
-> Maybe Bitmap -> (Int -> Int -> Int) -> Vector Int -> ST s ()
forall a s.
Unbox a =>
MVector s Int
-> Maybe Bitmap -> (Int -> a -> Int) -> Vector a -> ST s ()
hashUnboxed MVector s Int
mh Maybe Bitmap
ubm Int -> Int -> Int
mixInt Vector a
Vector Int
v
                Maybe (a :~: Int)
Nothing ->
                    case TypeRep a -> TypeRep Double -> Maybe (a :~: Double)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Double) of
                        Just a :~: Double
Refl -> MVector s Int
-> Maybe Bitmap
-> (Int -> Double -> Int)
-> Vector Double
-> ST s ()
forall a s.
Unbox a =>
MVector s Int
-> Maybe Bitmap -> (Int -> a -> Int) -> Vector a -> ST s ()
hashUnboxed MVector s Int
mh Maybe Bitmap
ubm Int -> Double -> Int
mixDouble Vector a
Vector Double
v
                        Maybe (a :~: Double)
Nothing ->
                            case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @a of
                                SBool (IntegralTypes a)
STrue ->
                                    MVector s Int
-> Maybe Bitmap -> (Int -> a -> Int) -> Vector a -> ST s ()
forall a s.
Unbox a =>
MVector s Int
-> Maybe Bitmap -> (Int -> a -> Int) -> Vector a -> ST s ()
hashUnboxed MVector s Int
mh Maybe Bitmap
ubm (\Int
h a
d -> Int -> Int -> Int
mixInt Int
h (forall a b. (Integral a, Num b) => a -> b
fromIntegral @a @Int a
d)) Vector a
v
                                SBool (IntegralTypes a)
SFalse ->
                                    case forall a. SBoolI (FloatingTypes a) => SBool (FloatingTypes a)
sFloating @a of
                                        SBool (FloatingTypes a)
STrue ->
                                            MVector s Int
-> Maybe Bitmap -> (Int -> a -> Int) -> Vector a -> ST s ()
forall a s.
Unbox a =>
MVector s Int
-> Maybe Bitmap -> (Int -> a -> Int) -> Vector a -> ST s ()
hashUnboxed MVector s Int
mh Maybe Bitmap
ubm (\Int
h a
d -> Int -> Double -> Int
mixDouble Int
h (a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac a
d :: Double)) Vector a
v
                                        SBool (FloatingTypes a)
SFalse ->
                                            MVector s Int
-> Maybe Bitmap -> (Int -> a -> Int) -> Vector a -> ST s ()
forall a s.
Unbox a =>
MVector s Int
-> Maybe Bitmap -> (Int -> a -> Int) -> Vector a -> ST s ()
hashUnboxed MVector s Int
mh Maybe Bitmap
ubm Int -> a -> Int
forall a. Show a => Int -> a -> Int
mixShow Vector a
v
        BoxedColumn Maybe Bitmap
bm (Vector a
v :: V.Vector a) ->
            case TypeRep a -> TypeRep Text -> Maybe (a :~: Text)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @T.Text) of
                Just a :~: Text
Refl ->
                    (Int -> Text -> ST s ()) -> Vector Text -> ST s ()
forall (m :: * -> *) a b.
Monad m =>
(Int -> a -> m b) -> Vector a -> m ()
V.imapM_
                        ( \Int
i Text
t -> do
                            !Int
h <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
mh Int
i
                            let h' :: Int
h' = case Maybe Bitmap
bm of
                                    Just Bitmap
bm' | Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm' Int
i) -> Int -> Int -> Int
mixInt Int
h Int
nullSalt
                                    Maybe Bitmap
_ -> Int -> Text -> Int
mixText Int
h Text
t
                            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
mh Int
i Int
h'
                        )
                        Vector a
Vector Text
v
                Maybe (a :~: Text)
Nothing ->
                    (Int -> a -> ST s ()) -> Vector a -> ST s ()
forall (m :: * -> *) a b.
Monad m =>
(Int -> a -> m b) -> Vector a -> m ()
V.imapM_
                        ( \Int
i a
d -> do
                            !Int
h <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
mh Int
i
                            let h' :: Int
h' = case Maybe Bitmap
bm of
                                    Just Bitmap
bm' | Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm' Int
i) -> Int -> Int -> Int
mixInt Int
h Int
nullSalt
                                    Maybe Bitmap
_ -> Int -> a -> Int
forall a. Show a => Int -> a -> Int
mixShow Int
h a
d
                            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
mh Int
i Int
h'
                        )
                        Vector a
v
        PackedText Maybe Bitmap
bm PackedTextData
p -> MVector s Int -> Maybe Bitmap -> PackedTextData -> ST s ()
forall s.
MVector s Int -> Maybe Bitmap -> PackedTextData -> ST s ()
hashPacked MVector s Int
mh Maybe Bitmap
bm PackedTextData
p
        MergedColumn Column
_ Column
_ ->
            [Char] -> ST s ()
forall a. HasCallStack => [Char] -> a
error [Char]
"computeHashes: MergedColumn is normalized before hashing"
    MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
mh

{- | Build the row-key equality predicate over the selected key columns.
@eqKeyRow df idxs a b@ is 'True' iff rows @a@ and @b@ agree on all key columns
(validity first, a null equals only a null). Used to reject hash collisions.
-}
eqKeyRow :: DataFrame -> [Int] -> Int -> Int -> Bool
eqKeyRow :: DataFrame -> [Int] -> Int -> Int -> Bool
eqKeyRow DataFrame
df [Int]
indicesToGroup =
    let !preds :: [Int -> Int -> Bool]
preds = (Int -> Int -> Int -> Bool) -> [Int] -> [Int -> Int -> Bool]
forall a b. (a -> b) -> [a] -> [b]
map (Column -> Int -> Int -> Bool
colEqRow (Column -> Int -> Int -> Bool)
-> (Int -> Column) -> Int -> Int -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (DataFrame -> Vector Column
columns DataFrame
df Vector Column -> Int -> Column
forall a. Vector a -> Int -> a
V.!)) [Int]
indicesToGroup
        go :: [t -> t -> Bool] -> t -> t -> Bool
go [] t
_ t
_ = Bool
True
        go (t -> t -> Bool
p : [t -> t -> Bool]
ps) t
a t
b = t -> t -> Bool
p t
a t
b Bool -> Bool -> Bool
&& [t -> t -> Bool] -> t -> t -> Bool
go [t -> t -> Bool]
ps t
a t
b
     in [Int -> Int -> Bool] -> Int -> Int -> Bool
forall {t} {t}. [t -> t -> Bool] -> t -> t -> Bool
go [Int -> Int -> Bool]
preds

{- | Per-column row equality respecting nulls. Two rows are equal at a column
when both are null, or both are valid and their values compare equal.
-}
colEqRow :: Column -> (Int -> Int -> Bool)
colEqRow :: Column -> Int -> Int -> Bool
colEqRow c :: Column
c@(MergedColumn Column
_ Column
_) = Column -> Int -> Int -> Bool
colEqRow (Column -> Column
materializeMerged Column
c)
colEqRow (UnboxedColumn Maybe Bitmap
bm Vector a
v) =
    let eqV :: Int -> Int -> Bool
eqV Int
a Int
b = Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
b
     in Maybe Bitmap -> (Int -> Int -> Bool) -> Int -> Int -> Bool
withNulls Maybe Bitmap
bm Int -> Int -> Bool
eqV
colEqRow (BoxedColumn Maybe Bitmap
bm Vector a
v) =
    let eqV :: Int -> Int -> Bool
eqV Int
a Int
b = Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndex Vector a
v Int
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndex Vector a
v Int
b
     in Maybe Bitmap -> (Int -> Int -> Bool) -> Int -> Int -> Bool
withNulls Maybe Bitmap
bm Int -> Int -> Bool
eqV
colEqRow (PackedText Maybe Bitmap
bm PackedTextData
p) =
    let eqV :: Int -> Int -> Bool
eqV Int
a Int
b =
            let (Array
arrA, Int
oA, Int
lA) = PackedTextData -> Int -> (Array, Int, Int)
packedSlice PackedTextData
p Int
a
                (Array
arrB, Int
oB, Int
lB) = PackedTextData -> Int -> (Array, Int, Int)
packedSlice PackedTextData
p Int
b
             in Array -> Int -> Int -> Array -> Int -> Int -> Bool
sliceEqBytes Array
arrA Int
oA Int
lA Array
arrB Int
oB Int
lB
     in Maybe Bitmap -> (Int -> Int -> Bool) -> Int -> Int -> Bool
withNulls Maybe Bitmap
bm Int -> Int -> Bool
eqV
{-# INLINE colEqRow #-}

{- | Wrap a value-equality with null handling: equal iff both valid and the
values agree, or both null.
-}
withNulls :: Maybe Bitmap -> (Int -> Int -> Bool) -> (Int -> Int -> Bool)
withNulls :: Maybe Bitmap -> (Int -> Int -> Bool) -> Int -> Int -> Bool
withNulls Maybe Bitmap
Nothing Int -> Int -> Bool
eqV = Int -> Int -> Bool
eqV
withNulls (Just Bitmap
bm) Int -> Int -> Bool
eqV = \Int
a Int
b ->
    case (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
a, Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
b) of
        (Bool
True, Bool
True) -> Int -> Int -> Bool
eqV Int
a Int
b
        (Bool
False, Bool
False) -> Bool
True
        (Bool, Bool)
_ -> Bool
False
{-# INLINE withNulls #-}

{- | Derive @(valueIndices, offsets)@ from @rowToGroup@ via a stable counting
sort on the group id: a per-group count, a prefix-sum into group offsets, then a
single placement pass keeps rows in original order within each group.
-}
indicesFromGroups :: VU.Vector Int -> Int -> (VU.Vector Int, VU.Vector Int)
indicesFromGroups :: Vector Int -> Int -> (Vector Int, Vector Int)
indicesFromGroups Vector Int
rtg Int
nGroups = (forall s. ST s (Vector Int, Vector Int))
-> (Vector Int, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int, Vector Int))
 -> (Vector Int, Vector Int))
-> (forall s. ST s (Vector Int, Vector Int))
-> (Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
    let !n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
rtg
    MVector s Int
counts <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int
nGroups Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
0
    let countLoop :: Int -> ST s ()
countLoop !Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                let !g :: Int
g = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
rtg Int
i
                Int
c <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
counts Int
g
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
counts Int
g (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                Int -> ST s ()
countLoop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> ST s ()
countLoop Int
0
    MVector s Int
offsM <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new (Int
nGroups Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    let scan :: Int -> Int -> ST s ()
scan !Int
k !Int
acc
            | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
nGroups = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
offsM Int
k Int
acc
                Int
c <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
counts Int
k
                Int -> Int -> ST s ()
scan (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
c)
    Int -> Int -> ST s ()
scan Int
0 Int
0
    let seed :: Int -> ST s ()
seed !Int
k
            | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
nGroups = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                Int
s <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
offsM Int
k
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
counts Int
k Int
s
                Int -> ST s ()
seed (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> ST s ()
seed Int
0
    MVector s Int
vis <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
n
    let place :: Int -> ST s ()
place !Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                let !g :: Int
g = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
rtg Int
i
                Int
pos <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
counts Int
g
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
vis Int
pos Int
i
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
counts Int
g (Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                Int -> ST s ()
place (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> ST s ()
place Int
0
    Vector Int
offs <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
offsM
    Vector Int
frozenVis <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
vis
    (Vector Int, Vector Int) -> ST s (Vector Int, Vector Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector Int
frozenVis, Vector Int
offs)

{- | Fold a value-mix over an unboxed column into the running hash vector,
respecting the null bitmap: a null slot mixes a fixed 'nullSalt' sentinel.
-}
hashUnboxed ::
    (VU.Unbox a) =>
    VUM.MVector s Int ->
    Maybe Bitmap ->
    (Int -> a -> Int) ->
    VU.Vector a ->
    ST s ()
hashUnboxed :: forall a s.
Unbox a =>
MVector s Int
-> Maybe Bitmap -> (Int -> a -> Int) -> Vector a -> ST s ()
hashUnboxed MVector s Int
mh Maybe Bitmap
ubm Int -> a -> Int
mix Vector a
v = case Maybe Bitmap
ubm of
    Maybe Bitmap
Nothing ->
        (Int -> a -> ST s ()) -> Vector a -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
(Int -> a -> m b) -> Vector a -> m ()
VU.imapM_
            ( \Int
i a
x -> do
                !Int
h <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
mh Int
i
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
mh Int
i (Int -> a -> Int
mix Int
h a
x)
            )
            Vector a
v
    Just Bitmap
bm ->
        (Int -> a -> ST s ()) -> Vector a -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
(Int -> a -> m b) -> Vector a -> m ()
VU.imapM_
            ( \Int
i a
x -> do
                !Int
h <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
mh Int
i
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite
                    MVector s Int
MVector (PrimState (ST s)) Int
mh
                    Int
i
                    (if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i then Int -> a -> Int
mix Int
h a
x else Int -> Int -> Int
mixInt Int
h Int
nullSalt)
            )
            Vector a
v
{-# INLINE hashUnboxed #-}

{- | Hash a packed-text column over its raw UTF-8 byte slices (no per-row
'Data.Text.Text'), mixing 'nullSalt' for null rows. Shares 'mixBytes' with
'mixText' so packed and boxed Text columns hash identically.
-}
hashPacked ::
    VUM.MVector s Int -> Maybe Bitmap -> PackedTextData -> ST s ()
hashPacked :: forall s.
MVector s Int -> Maybe Bitmap -> PackedTextData -> ST s ()
hashPacked MVector s Int
mh Maybe Bitmap
bm PackedTextData
p = Int -> ST s ()
go Int
0
  where
    !n :: Int
n = PackedTextData -> Int
packedLength PackedTextData
p
    go :: Int -> ST s ()
go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            !Int
h <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
mh Int
i
            let h' :: Int
h' = case Maybe Bitmap
bm of
                    Just Bitmap
bm' | Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm' Int
i) -> Int -> Int -> Int
mixInt Int
h Int
nullSalt
                    Maybe Bitmap
_ -> let (Array
arr, Int
o, Int
l) = PackedTextData -> Int -> (Array, Int, Int)
packedSlice PackedTextData
p Int
i in Int -> Array -> Int -> Int -> Int
mixBytes Int
h Array
arr Int
o Int
l
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
mh Int
i Int
h'
            Int -> ST s ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE hashPacked #-}

-- Inline accessors to avoid depending on Operations.Core

columnNames :: DataFrame -> [T.Text]
columnNames :: DataFrame -> [Text]
columnNames = Map Text Int -> [Text]
forall k a. Map k a -> [k]
M.keys (Map Text Int -> [Text])
-> (DataFrame -> Map Text Int) -> DataFrame -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame -> Map Text Int
columnIndices

nRows :: DataFrame -> Int
nRows :: DataFrame -> Int
nRows = (Int, Int) -> Int
forall a b. (a, b) -> a
fst ((Int, Int) -> Int)
-> (DataFrame -> (Int, Int)) -> DataFrame -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame -> (Int, Int)
dataframeDimensions

{- | Build the rowToGroup lookup vector from valueIndices and offsets.
rowToGroup[i] = k means row i belongs to group k.
-}
buildRowToGroup :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int
buildRowToGroup :: Int -> Vector Int -> Vector Int -> Vector Int
buildRowToGroup Int
n Vector Int
vis Vector Int
os = (forall s. ST s (Vector Int)) -> Vector Int
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int)) -> Vector Int)
-> (forall s. ST s (Vector Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
    MVector s Int
rtg <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
n
    let nGroups :: Int
nGroups = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
os Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
nGroups Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
k ->
        let s :: Int
s = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
os Int
k
            e :: Int
e = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
os (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
         in [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
s .. Int
e Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
rtg (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
vis Int
i) Int
k
    MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
rtg
{-# NOINLINE buildRowToGroup #-}

changingPoints :: VU.Vector (Int, Int) -> VU.Vector Int
changingPoints :: Vector (Int, Int) -> Vector Int
changingPoints Vector (Int, Int)
vs =
    Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector a
VU.reverse
        ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList (Vector (Int, Int) -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector (Int, Int)
vs Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: ([Int], Int) -> [Int]
forall a b. (a, b) -> a
fst ((([Int], Int) -> Int -> (Int, Int) -> ([Int], Int))
-> ([Int], Int) -> Vector (Int, Int) -> ([Int], Int)
forall b a. Unbox b => (a -> Int -> b -> a) -> a -> Vector b -> a
VU.ifoldl' ([Int], Int) -> Int -> (Int, Int) -> ([Int], Int)
forall {b} {a} {a}. Eq b => ([a], b) -> a -> (a, b) -> ([a], b)
findChangePoints ([Int], Int)
initialState Vector (Int, Int)
vs)))
  where
    initialState :: ([Int], Int)
initialState = ([Int
0], (Int, Int) -> Int
forall a b. (a, b) -> b
snd (Vector (Int, Int) -> (Int, Int)
forall a. Unbox a => Vector a -> a
VU.head Vector (Int, Int)
vs))
    findChangePoints :: ([a], b) -> a -> (a, b) -> ([a], b)
findChangePoints (![a]
offs, !b
currentVal) a
index (a
_, !b
newVal)
        | b
currentVal b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
newVal = ([a]
offs, b
currentVal)
        | Bool
otherwise = (a
index a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
offs, b
newVal)