Encoding.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}

module DataFrame.IO.Parquet.Encoding (
    -- Kept from the original Encoding module (used by Levels)
    ceilLog2,
    bitWidthForMaxLevel,
    -- Vector-based RLE/bit-packed decoder (from new parser)
    decodeRLEBitPackedHybrid,
    rleHybridFold,
    extractBitsInto,
    fillRun,
    decodeDictIndices,
) where

import Control.Monad (when)
import Control.Monad.ST (ST, runST)
import Data.Bits
import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BSU
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import Data.Word
import DataFrame.IO.Parquet.Binary (readUVarInt)
import DataFrame.Internal.Binary (littleEndianWord32)

ceilLog2 :: Int -> Int
ceilLog2 x
    | x <= 1 = 0
    | otherwise = 1 + ceilLog2 ((x + 1) `div` 2)

bitWidthForMaxLevel :: Int -> Int
bitWidthForMaxLevel maxLevel = ceilLog2 (maxLevel + 1)

decodeRLEBitPackedHybrid ::
    -- | Bit width per value (0 = all zeros, use 'VU.replicate')
    Int ->
    -- | Exact number of values to decode
    Int ->
    BS.ByteString ->
    (VU.Vector Word32, BS.ByteString)
decodeRLEBitPackedHybrid bw need bs
    | bw == 0 = (VU.replicate need 0, bs)
    | otherwise = runST $ do
        mv <- VUM.new need
        rest <- go mv 0 bs
        dat <- VU.unsafeFreeze mv
        return (dat, rest)
  where
    !mask = if bw == 32 then maxBound else (1 `shiftL` bw) - 1 :: Word32
    go :: VUM.STVector s Word32 -> Int -> BS.ByteString -> ST s BS.ByteString
    go mv !filled !buf
        | filled >= need = return buf
        | BS.null buf = return buf
        | otherwise =
            let (hdr64, afterHdr) = readUVarInt buf
                isPacked = (hdr64 .&. 1) == 1
             in if isPacked
                    then do
                        let groups = fromIntegral (hdr64 `shiftR` 1) :: Int
                            totalVals = groups * 8
                            takeN = min (need - filled) totalVals
                            -- Consume the whole group even if we need less.
                            bytesN = (bw * totalVals + 7) `div` 8
                            (chunk, rest) = BS.splitAt bytesN afterHdr
                        extractBitsInto bw takeN chunk mv filled
                        go mv (filled + takeN) rest
                    else do
                        let runLen = fromIntegral (hdr64 `shiftR` 1) :: Int
                            nbytes = (bw + 7) `div` 8
                            val = littleEndianWord32 (BS.take 4 afterHdr) .&. mask
                            takeN = min (need - filled) runLen
                        -- Fill the run directly — no list, no reverse.
                        fillRun mv filled (filled + takeN) val
                        go mv (filled + takeN) (BS.drop nbytes afterHdr)
{-# INLINE decodeRLEBitPackedHybrid #-}

-- | Fill @mv[start..end-1]@ with @val@ using a bulk @memset@-style write.
fillRun :: VUM.STVector s Word32 -> Int -> Int -> Word32 -> ST s ()
fillRun mv i end = VUM.set (VUM.unsafeSlice i (end - i) mv)
{-# INLINE fillRun #-}

-- | Write bit-packed values into a mutable vector in a single pass.
extractBitsInto ::
    -- | Bit width
    Int ->
    -- | Number of values to extract
    Int ->
    BS.ByteString ->
    VUM.STVector s Word32 ->
    -- | Write offset into @mv@
    Int ->
    ST s ()
extractBitsInto bw count bs mv off = go 0 (0 :: Word64) 0 0
  where
    !mask = if bw == 32 then maxBound else (1 `unsafeShiftL` bw) - 1 :: Word64
    !len = BS.length bs
    go !byteIdx !acc !accBits !done
        | done >= count = return ()
        | accBits >= bw = do
            VUM.unsafeWrite mv (off + done) (fromIntegral (acc .&. mask))
            go byteIdx (acc `unsafeShiftR` bw) (accBits - bw) (done + 1)
        | byteIdx >= len = return ()
        | otherwise =
            let b = fromIntegral (BSU.unsafeIndex bs byteIdx) :: Word64
             in go (byteIdx + 1) (acc .|. (b `unsafeShiftL` accBits)) (accBits + 8) done
{-# INLINE extractBitsInto #-}

-- | Decode dictionary indices from a bit-width-prefixed DATA_PAGE stream.
decodeDictIndices :: Int -> BS.ByteString -> (VU.Vector Int, BS.ByteString)
decodeDictIndices need bs = case BS.uncons bs of
    Nothing -> error "decodeDictIndices: empty stream"
    Just (w0, rest0) ->
        let bw = fromIntegral w0 :: Int
            (raw, rest1) = decodeRLEBitPackedHybrid bw need rest0
         in (VU.map fromIntegral raw, rest1)
{-# INLINE decodeDictIndices #-}

-- | Stream RLE\/bit-packed hybrid values to a sink.
rleHybridFold ::
    Int -> Int -> BS.ByteString -> (Int -> Word32 -> ST s ()) -> ST s ()
rleHybridFold bw need bs sink
    | bw == 0 = zeros 0
    | otherwise = go 0 bs
  where
    !mask = if bw == 32 then maxBound else (1 `unsafeShiftL` bw) - 1 :: Word32
    zeros !i = when (i < need) (sink i 0 >> zeros (i + 1))
    go !filled !buf
        | filled >= need = pure ()
        | BS.null buf = pure ()
        | otherwise =
            let (hdr64, afterHdr) = readUVarInt buf
                isPacked = (hdr64 .&. 1) == 1
             in if isPacked
                    then do
                        let groups = fromIntegral (hdr64 `shiftR` 1) :: Int
                            totalVals = groups * 8
                            takeN = min (need - filled) totalVals
                            bytesN = (bw * totalVals + 7) `div` 8
                            (chunk, rest) = BS.splitAt bytesN afterHdr
                        extractBitsFold bw takeN chunk (\j w -> sink (filled + j) w)
                        go (filled + takeN) rest
                    else do
                        let runLen = fromIntegral (hdr64 `shiftR` 1) :: Int
                            nbytes = (bw + 7) `div` 8
                            val = littleEndianWord32 (BS.take 4 afterHdr) .&. mask
                            takeN = min (need - filled) runLen
                            run !j = when (j < takeN) (sink (filled + j) val >> run (j + 1))
                        run 0
                        go (filled + takeN) (BS.drop nbytes afterHdr)
{-# INLINE rleHybridFold #-}

-- | 'extractBitsInto' with a sink instead of a destination vector.
extractBitsFold ::
    Int -> Int -> BS.ByteString -> (Int -> Word32 -> ST s ()) -> ST s ()
extractBitsFold bw count bs sink = go 0 (0 :: Word64) 0 0
  where
    !mask = if bw == 32 then maxBound else (1 `unsafeShiftL` bw) - 1 :: Word64
    !len = BS.length bs
    go !byteIdx !acc !accBits !done
        | done >= count = return ()
        | accBits >= bw = do
            sink done (fromIntegral (acc .&. mask))
            go byteIdx (acc `unsafeShiftR` bw) (accBits - bw) (done + 1)
        | byteIdx >= len = return ()
        | otherwise =
            let b = fromIntegral (BSU.unsafeIndex bs byteIdx) :: Word64
             in go (byteIdx + 1) (acc .|. (b `unsafeShiftL` accBits)) (accBits + 8) done
{-# INLINE extractBitsFold #-}