DictEncode.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

{- | Dictionary-encode a text (or factor) group key to dense @Int@ codes: each row
gets a first-appearance code @0..card-1@ (NULL reserved) plus the cardinality. A
tested building block; profiled slower than the hash group-by, so unused for now.
-}
module DataFrame.Internal.DictEncode (
    dictEncodeColumn,
    dictEncodeColumnUpTo,
    dictMaxCardinality,
) where

import Control.Monad.ST (runST)
import qualified Data.Text as T
import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import Type.Reflection (typeRep)

import DataFrame.Internal.Column (Bitmap, Column (..), bitmapTestBit)
import DataFrame.Internal.Hash (fnvOffset, mixBytes, mixText, nullSalt)
import DataFrame.Internal.HashTable (htInsert, newHashTable)
import DataFrame.Internal.PackedText (
    PackedTextData,
    packedLength,
    packedSlice,
    sliceEqBytes,
 )

{- | Largest distinct-value count we will dictionary-encode. Above this the codes
no longer index a reasonable direct accumulator and the encode pass is pure
overhead, so the caller keeps the plain hash group-by.
-}
dictMaxCardinality :: Int
dictMaxCardinality = 1048576

{- | Dictionary-encode a text-like column to dense first-appearance @Int@ codes,
returning @Just (codes, cardinality)@ (a NULL row gets its own reserved code).
'Nothing' for non-text columns or cardinality above 'dictMaxCardinality'.
-}
dictEncodeColumn :: Column -> Maybe (VU.Vector Int, Int)
dictEncodeColumn = dictEncodeColumnUpTo dictMaxCardinality

{- | Dictionary-encode like 'dictEncodeColumn' but bail to 'Nothing' as soon as
the distinct count would exceed @maxCard@, letting a low-cardinality probe avoid
a full high-cardinality pass.
-}
dictEncodeColumnUpTo :: Int -> Column -> Maybe (VU.Vector Int, Int)
dictEncodeColumnUpTo maxCard (PackedText bm p) = encodePacked maxCard bm p
dictEncodeColumnUpTo maxCard (BoxedColumn bm (v :: V.Vector a)) =
    case testEquality (typeRep @a) (typeRep @T.Text) of
        Just Refl -> encodeBoxedText maxCard bm v
        Nothing -> Nothing
dictEncodeColumnUpTo _ _ = Nothing

{- | Encode a packed-text column: hash each row's raw UTF-8 bytes (the grouping
'mixBytes'), re-verify byte equality on collisions, assign dense codes in
first-appearance order. A null row hashes 'nullSalt'.
-}
encodePacked ::
    Int -> Maybe Bitmap -> PackedTextData -> Maybe (VU.Vector Int, Int)
encodePacked maxCard bm p =
    let !n = packedLength p
        valid i = case bm of
            Just b -> bitmapTestBit b i
            Nothing -> True
        hashAt i =
            if valid i
                then let (arr, o, l) = packedSlice p i in mixBytes fnvOffset arr o l
                else nullSalt
        eqAt a b =
            case (valid a, valid b) of
                (True, True) ->
                    let (arrA, oA, lA) = packedSlice p a
                        (arrB, oB, lB) = packedSlice p b
                     in sliceEqBytes arrA oA lA arrB oB lB
                (False, False) -> True
                _ -> False
     in buildCodes maxCard n hashAt eqAt

{- | Encode a boxed 'Data.Text.Text' column, mirroring 'encodePacked' but over
boxed values (used when a user-built Text column is grouped).
-}
encodeBoxedText ::
    Int -> Maybe Bitmap -> V.Vector T.Text -> Maybe (VU.Vector Int, Int)
encodeBoxedText maxCard bm v =
    let !n = V.length v
        valid i = case bm of
            Just b -> bitmapTestBit b i
            Nothing -> True
        hashAt i =
            if valid i then mixText fnvOffset (V.unsafeIndex v i) else nullSalt
        eqAt a b =
            case (valid a, valid b) of
                (True, True) -> V.unsafeIndex v a == V.unsafeIndex v b
                (False, False) -> True
                _ -> False
     in buildCodes maxCard n hashAt eqAt

{- | The shared code-assignment loop: bucket every row through an open-addressing
table on its precomputed hash, re-verify with @eqAt@ on a hit, assign dense
first-appearance codes. Bails to 'Nothing' once the distinct count exceeds @maxCard@.
-}
buildCodes ::
    Int -> Int -> (Int -> Int) -> (Int -> Int -> Bool) -> Maybe (VU.Vector Int, Int)
buildCodes maxCard n hashAt eqAt
    | n == 0 = Just (VU.empty, 0)
    | otherwise = runST $ do
        ht <- newHashTable (min n (maxCard + 1))
        codes <- VUM.new n
        let go !i !next
                | i >= n = pure (Just next)
                | next > maxCard = pure Nothing
                | otherwise = do
                    let !h = hashAt i
                    (code, isNew) <- htInsert ht eqAt next i h
                    VUM.unsafeWrite codes i code
                    go (i + 1) (if isNew then next + 1 else next)
        mres <- go 0 0
        case mres of
            Nothing -> pure Nothing
            Just card -> do
                frozen <- VU.unsafeFreeze codes
                pure (Just (frozen, card))