{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE ScopedTypeVariables #-}

module DataFrame.IO.Parquet.Page (
    -- Types
    PageDecoder,
    UnboxedPageDecoder,
    -- Per-type decoders
    boolDecoder,
    int32Decoder,
    int64Decoder,
    int96Decoder,
    floatDecoder,
    doubleDecoder,
    byteArrayDecoder,
    fixedLenByteArrayDecoder,
    -- Page iteration
    foldColumnPagesM,
    foldColumnDataPagesM,
    RawPage,
    appendStringPageIO,
    appendNullableStringPageIO,
) where

import Control.Monad.IO.Class (MonadIO (liftIO))
import Control.Monad.ST (RealWorld, stToIO)
import Data.Bits (shiftL, shiftR, (.&.), (.|.))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BSU
import Data.Int (Int32, Int64)
import Data.Maybe (fromJust, fromMaybe)
import qualified Data.Text as T
import Data.Text.Encoding (decodeUtf8Lenient)
import Data.Time (UTCTime)
import qualified Data.Vector as VB
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Unboxed as VU
import Data.Word (Word8)
import DataFrame.IO.Parquet.Decompress (decompressData)
import DataFrame.IO.Parquet.Dictionary (
    DictVals (..),
    readDictVals,
 )
import DataFrame.IO.Parquet.Encoding (decodeDictIndices)
import DataFrame.IO.Parquet.Levels (readLevelsV1, readLevelsV2)
import DataFrame.IO.Parquet.Thrift (
    ColumnChunk (..),
    ColumnMetaData (..),
    CompressionCodec,
    DataPageHeader (..),
    DataPageHeaderV2 (..),
    DictionaryPageHeader (..),
    Encoding (..),
    PageHeader (..),
    PageType (..),
    ThriftType (..),
    unField,
 )
import DataFrame.IO.Parquet.Time (int96ToUTCTime)
import DataFrame.IO.Parquet.Utils (ColumnDescription (..))
import DataFrame.IO.Utils.RandomAccess (RandomAccess (..), Range (Range))
import DataFrame.Internal.Binary (
    littleEndianInt32,
    littleEndianWord32,
    littleEndianWord64,
 )
import DataFrame.Internal.ColumnBuilder (
    TextBuilder,
    appendNull,
    appendText,
    appendTextSliceFromPtr,
 )
import Foreign.Ptr (Ptr, castPtr, plusPtr)
import Foreign.Storable (peekByteOff)
import GHC.Float (castWord32ToFloat, castWord64ToDouble)
import Pinch (decodeWithLeftovers)
import qualified Pinch

-- ---------------------------------------------------------------------------
-- Types
-- ---------------------------------------------------------------------------

{- | A type-specific page decoder.
Given the optional dictionary, the page encoding, the number of present
values, and the decompressed value bytes, returns exactly @nPresent@ values.
-}
type PageDecoder a =
    Maybe DictVals -> Encoding -> Int -> BS.ByteString -> VB.Vector a

type UnboxedPageDecoder a =
    Maybe DictVals -> Encoding -> Int -> BS.ByteString -> VU.Vector a

-- ---------------------------------------------------------------------------
-- Per-type decoders
-- ---------------------------------------------------------------------------

boolDecoder :: UnboxedPageDecoder Bool
boolDecoder :: UnboxedPageDecoder Bool
boolDecoder Maybe DictVals
mDict Encoding
enc Int
nPresent ByteString
bs = case Encoding
enc of
    -- PLAIN bools are bit-packed (1 bit/value, LSB-first). Generate the
    -- unboxed vector directly by indexing the bit for each row, avoiding an
    -- intermediate @[Bool]@ list.
    PLAIN Enumeration 0
_ ->
        Int -> (Int -> Bool) -> Vector Bool
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
nPresent ((Int -> Bool) -> Vector Bool) -> (Int -> Bool) -> Vector Bool
forall a b. (a -> b) -> a -> b
$ \Int
i ->
            (ByteString -> Int -> Word8
BSU.unsafeIndex ByteString
bs (Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3) Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftR` (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
7)) Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
1 Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
1
    RLE_DICTIONARY Enumeration 8
_ -> Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> Bool) -> Vector Bool
forall a.
Unbox a =>
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
unboxedLookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Bool
getBool
    PLAIN_DICTIONARY Enumeration 2
_ -> Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> Bool) -> Vector Bool
forall a.
Unbox a =>
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
unboxedLookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Bool
getBool
    Encoding
_ -> [Char] -> Vector Bool
forall a. HasCallStack => [Char] -> a
error ([Char]
"boolDecoder: unsupported encoding " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Encoding -> [Char]
forall a. Show a => a -> [Char]
show Encoding
enc)
  where
    getBool :: DictVals -> Int -> Bool
getBool (DBool Vector Bool
ds) Int
i = Vector Bool
ds Vector Bool -> Int -> Bool
forall a. Vector a -> Int -> a
VB.! Int
i
    getBool DictVals
d Int
_ = [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error ([Char]
"boolDecoder: wrong dict type, got " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ DictVals -> [Char]
forall a. Show a => a -> [Char]
show DictVals
d)

int32Decoder :: UnboxedPageDecoder Int32
int32Decoder :: UnboxedPageDecoder Int32
int32Decoder Maybe DictVals
mDict Encoding
enc Int
nPresent ByteString
bs = case Encoding
enc of
    PLAIN Enumeration 0
_ -> Vector Int32 -> Vector Int32
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VU.convert (Int -> ByteString -> Vector Int32
readNInt32 Int
nPresent ByteString
bs)
    RLE_DICTIONARY Enumeration 8
_ -> Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> Int32) -> Vector Int32
forall a.
Unbox a =>
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
unboxedLookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Int32
getInt32
    PLAIN_DICTIONARY Enumeration 2
_ -> Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> Int32) -> Vector Int32
forall a.
Unbox a =>
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
unboxedLookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Int32
getInt32
    Encoding
_ -> [Char] -> Vector Int32
forall a. HasCallStack => [Char] -> a
error ([Char]
"int32Decoder: unsupported encoding " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Encoding -> [Char]
forall a. Show a => a -> [Char]
show Encoding
enc)
  where
    getInt32 :: DictVals -> Int -> Int32
getInt32 (DInt32 Vector Int32
ds) Int
i = Vector Int32
ds Vector Int32 -> Int -> Int32
forall a. Vector a -> Int -> a
VB.! Int
i
    getInt32 DictVals
d Int
_ = [Char] -> Int32
forall a. HasCallStack => [Char] -> a
error ([Char]
"int32Decoder: wrong dict type, got " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ DictVals -> [Char]
forall a. Show a => a -> [Char]
show DictVals
d)

int64Decoder :: UnboxedPageDecoder Int64
int64Decoder :: UnboxedPageDecoder Int64
int64Decoder Maybe DictVals
mDict Encoding
enc Int
nPresent ByteString
bs = case Encoding
enc of
    PLAIN Enumeration 0
_ -> Vector Int64 -> Vector Int64
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VU.convert (Int -> ByteString -> Vector Int64
readNInt64 Int
nPresent ByteString
bs)
    RLE_DICTIONARY Enumeration 8
_ -> Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> Int64) -> Vector Int64
forall a.
Unbox a =>
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
unboxedLookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Int64
getInt64
    PLAIN_DICTIONARY Enumeration 2
_ -> Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> Int64) -> Vector Int64
forall a.
Unbox a =>
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
unboxedLookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Int64
getInt64
    Encoding
_ -> [Char] -> Vector Int64
forall a. HasCallStack => [Char] -> a
error ([Char]
"int64Decoder: unsupported encoding " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Encoding -> [Char]
forall a. Show a => a -> [Char]
show Encoding
enc)
  where
    getInt64 :: DictVals -> Int -> Int64
getInt64 (DInt64 Vector Int64
ds) Int
i = Vector Int64
ds Vector Int64 -> Int -> Int64
forall a. Vector a -> Int -> a
VB.! Int
i
    getInt64 DictVals
d Int
_ = [Char] -> Int64
forall a. HasCallStack => [Char] -> a
error ([Char]
"int64Decoder: wrong dict type, got " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ DictVals -> [Char]
forall a. Show a => a -> [Char]
show DictVals
d)

int96Decoder :: PageDecoder UTCTime
int96Decoder :: PageDecoder UTCTime
int96Decoder Maybe DictVals
mDict Encoding
enc Int
nPresent ByteString
bs = case Encoding
enc of
    PLAIN Enumeration 0
_ -> [UTCTime] -> Vector UTCTime
forall a. [a] -> Vector a
VB.fromList (Int -> ByteString -> [UTCTime]
readNInt96 Int
nPresent ByteString
bs)
    RLE_DICTIONARY Enumeration 8
_ -> Maybe DictVals
-> Int
-> ByteString
-> (DictVals -> Int -> UTCTime)
-> Vector UTCTime
forall a.
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
lookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> UTCTime
getInt96
    PLAIN_DICTIONARY Enumeration 2
_ -> Maybe DictVals
-> Int
-> ByteString
-> (DictVals -> Int -> UTCTime)
-> Vector UTCTime
forall a.
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
lookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> UTCTime
getInt96
    Encoding
_ -> [Char] -> Vector UTCTime
forall a. HasCallStack => [Char] -> a
error ([Char]
"int96Decoder: unsupported encoding " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Encoding -> [Char]
forall a. Show a => a -> [Char]
show Encoding
enc)
  where
    getInt96 :: DictVals -> Int -> UTCTime
getInt96 (DInt96 Vector UTCTime
ds) Int
i = Vector UTCTime
ds Vector UTCTime -> Int -> UTCTime
forall a. Vector a -> Int -> a
VB.! Int
i
    getInt96 DictVals
d Int
_ = [Char] -> UTCTime
forall a. HasCallStack => [Char] -> a
error ([Char]
"int96Decoder: wrong dict type, got " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ DictVals -> [Char]
forall a. Show a => a -> [Char]
show DictVals
d)

floatDecoder :: UnboxedPageDecoder Float
floatDecoder :: UnboxedPageDecoder Float
floatDecoder Maybe DictVals
mDict Encoding
enc Int
nPresent ByteString
bs = case Encoding
enc of
    PLAIN Enumeration 0
_ -> Vector Float -> Vector Float
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VU.convert (Int -> ByteString -> Vector Float
readNFloat Int
nPresent ByteString
bs)
    RLE_DICTIONARY Enumeration 8
_ -> Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> Float) -> Vector Float
forall a.
Unbox a =>
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
unboxedLookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Float
getFloat
    PLAIN_DICTIONARY Enumeration 2
_ -> Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> Float) -> Vector Float
forall a.
Unbox a =>
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
unboxedLookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Float
getFloat
    Encoding
_ -> [Char] -> Vector Float
forall a. HasCallStack => [Char] -> a
error ([Char]
"floatDecoder: unsupported encoding " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Encoding -> [Char]
forall a. Show a => a -> [Char]
show Encoding
enc)
  where
    getFloat :: DictVals -> Int -> Float
getFloat (DFloat Vector Float
ds) Int
i = Vector Float
ds Vector Float -> Int -> Float
forall a. Vector a -> Int -> a
VB.! Int
i
    getFloat DictVals
d Int
_ = [Char] -> Float
forall a. HasCallStack => [Char] -> a
error ([Char]
"floatDecoder: wrong dict type, got " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ DictVals -> [Char]
forall a. Show a => a -> [Char]
show DictVals
d)

doubleDecoder :: UnboxedPageDecoder Double
doubleDecoder :: UnboxedPageDecoder Double
doubleDecoder Maybe DictVals
mDict Encoding
enc Int
nPresent ByteString
bs = case Encoding
enc of
    PLAIN Enumeration 0
_ -> Vector Double -> Vector Double
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VU.convert (Int -> ByteString -> Vector Double
readNDouble Int
nPresent ByteString
bs)
    RLE_DICTIONARY Enumeration 8
_ -> Maybe DictVals
-> Int
-> ByteString
-> (DictVals -> Int -> Double)
-> Vector Double
forall a.
Unbox a =>
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
unboxedLookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Double
getDouble
    PLAIN_DICTIONARY Enumeration 2
_ -> Maybe DictVals
-> Int
-> ByteString
-> (DictVals -> Int -> Double)
-> Vector Double
forall a.
Unbox a =>
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
unboxedLookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Double
getDouble
    Encoding
_ -> [Char] -> Vector Double
forall a. HasCallStack => [Char] -> a
error ([Char]
"doubleDecoder: unsupported encoding " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Encoding -> [Char]
forall a. Show a => a -> [Char]
show Encoding
enc)
  where
    getDouble :: DictVals -> Int -> Double
getDouble (DDouble Vector Double
ds) Int
i = Vector Double
ds Vector Double -> Int -> Double
forall a. Vector a -> Int -> a
VB.! Int
i
    getDouble DictVals
d Int
_ = [Char] -> Double
forall a. HasCallStack => [Char] -> a
error ([Char]
"doubleDecoder: wrong dict type, got " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ DictVals -> [Char]
forall a. Show a => a -> [Char]
show DictVals
d)

byteArrayDecoder :: PageDecoder T.Text
byteArrayDecoder :: PageDecoder Text
byteArrayDecoder Maybe DictVals
mDict Encoding
enc Int
nPresent ByteString
bs = case Encoding
enc of
    PLAIN Enumeration 0
_ -> [Text] -> Vector Text
forall a. [a] -> Vector a
VB.fromList (Int -> ByteString -> [Text]
readNTexts Int
nPresent ByteString
bs)
    RLE_DICTIONARY Enumeration 8
_ -> Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> Text) -> Vector Text
forall a.
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
lookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Text
getText
    PLAIN_DICTIONARY Enumeration 2
_ -> Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> Text) -> Vector Text
forall a.
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
lookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Text
getText
    Encoding
_ -> [Char] -> Vector Text
forall a. HasCallStack => [Char] -> a
error ([Char]
"byteArrayDecoder: unsupported encoding " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Encoding -> [Char]
forall a. Show a => a -> [Char]
show Encoding
enc)
  where
    getText :: DictVals -> Int -> Text
getText (DText Vector Text
ds) Int
i = Vector Text
ds Vector Text -> Int -> Text
forall a. Vector a -> Int -> a
VB.! Int
i
    getText DictVals
d Int
_ = [Char] -> Text
forall a. HasCallStack => [Char] -> a
error ([Char]
"byteArrayDecoder: wrong dict type, got " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ DictVals -> [Char]
forall a. Show a => a -> [Char]
show DictVals
d)

fixedLenByteArrayDecoder :: Int -> PageDecoder T.Text
fixedLenByteArrayDecoder :: Int -> PageDecoder Text
fixedLenByteArrayDecoder Int
len Maybe DictVals
mDict Encoding
enc Int
nPresent ByteString
bs = case Encoding
enc of
    PLAIN Enumeration 0
_ -> [Text] -> Vector Text
forall a. [a] -> Vector a
VB.fromList (Int -> Int -> ByteString -> [Text]
readNFixedTexts Int
len Int
nPresent ByteString
bs)
    RLE_DICTIONARY Enumeration 8
_ -> Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> Text) -> Vector Text
forall a.
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
lookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Text
getText
    PLAIN_DICTIONARY Enumeration 2
_ -> Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> Text) -> Vector Text
forall a.
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
lookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> Text
getText
    Encoding
_ -> [Char] -> Vector Text
forall a. HasCallStack => [Char] -> a
error ([Char]
"fixedLenByteArrayDecoder: unsupported encoding " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Encoding -> [Char]
forall a. Show a => a -> [Char]
show Encoding
enc)
  where
    getText :: DictVals -> Int -> Text
getText (DText Vector Text
ds) Int
i = Vector Text
ds Vector Text -> Int -> Text
forall a. Vector a -> Int -> a
VB.! Int
i
    getText DictVals
d Int
_ = [Char] -> Text
forall a. HasCallStack => [Char] -> a
error ([Char]
"fixedLenByteArrayDecoder: wrong dict type, got " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ DictVals -> [Char]
forall a. Show a => a -> [Char]
show DictVals
d)

{- | Shared dictionary-path helper: decode @nPresent@ RLE/bit-packed indices
and look each one up in the dictionary.
-}
lookupDict ::
    Maybe DictVals ->
    Int ->
    BS.ByteString ->
    (DictVals -> Int -> a) ->
    VB.Vector a
lookupDict :: forall a.
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
lookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> a
f = case Maybe DictVals
mDict of
    Maybe DictVals
Nothing -> [Char] -> Vector a
forall a. HasCallStack => [Char] -> a
error [Char]
"Dictionary-encoded page but no dictionary page seen"
    Just DictVals
dict ->
        let (Vector Int
idxs, ByteString
_) = Int -> ByteString -> (Vector Int, ByteString)
decodeDictIndices Int
nPresent ByteString
bs
         in Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
VB.generate Int
nPresent (DictVals -> Int -> a
f DictVals
dict (Int -> a) -> (Int -> Int) -> Int -> a
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
idxs)

unboxedLookupDict ::
    (VU.Unbox a) =>
    Maybe DictVals ->
    Int ->
    BS.ByteString ->
    (DictVals -> Int -> a) ->
    VU.Vector a
unboxedLookupDict :: forall a.
Unbox a =>
Maybe DictVals
-> Int -> ByteString -> (DictVals -> Int -> a) -> Vector a
unboxedLookupDict Maybe DictVals
mDict Int
nPresent ByteString
bs DictVals -> Int -> a
f = case Maybe DictVals
mDict of
    Maybe DictVals
Nothing -> [Char] -> Vector a
forall a. HasCallStack => [Char] -> a
error [Char]
"Dictionary-encoded page but no dictionary page seen"
    Just DictVals
dict ->
        let (Vector Int
idxs, ByteString
_) = Int -> ByteString -> (Vector Int, ByteString)
decodeDictIndices Int
nPresent ByteString
bs
         in Int -> (Int -> a) -> Vector a
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
nPresent (DictVals -> Int -> a
f DictVals
dict (Int -> a) -> (Int -> Int) -> Int -> a
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
idxs)

-- ---------------------------------------------------------------------------
-- Core page-iteration loop
-- ---------------------------------------------------------------------------

-- | Read the raw (compressed) byte range for a column chunk.
readChunkBytes ::
    (RandomAccess m) =>
    ColumnChunk ->
    m (CompressionCodec, ThriftType, BS.ByteString)
readChunkBytes :: forall (m :: * -> *).
RandomAccess m =>
ColumnChunk -> m (CompressionCodec, ThriftType, ByteString)
readChunkBytes ColumnChunk
columnChunk = do
    let meta :: ColumnMetaData
meta = Maybe ColumnMetaData -> ColumnMetaData
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe ColumnMetaData -> ColumnMetaData)
-> (Field 3 (Maybe ColumnMetaData) -> Maybe ColumnMetaData)
-> Field 3 (Maybe ColumnMetaData)
-> ColumnMetaData
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Field 3 (Maybe ColumnMetaData) -> Maybe ColumnMetaData
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (Field 3 (Maybe ColumnMetaData) -> ColumnMetaData)
-> Field 3 (Maybe ColumnMetaData) -> ColumnMetaData
forall a b. (a -> b) -> a -> b
$ ColumnChunk
columnChunk.cc_meta_data
        codec :: CompressionCodec
codec = Field 4 CompressionCodec -> CompressionCodec
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField ColumnMetaData
meta.cmd_codec
        pType :: ThriftType
pType = Field 1 ThriftType -> ThriftType
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField ColumnMetaData
meta.cmd_type
        dataOffset :: Integer
dataOffset = Int64 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Integer)
-> (Field 9 Int64 -> Int64) -> Field 9 Int64 -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Field 9 Int64 -> Int64
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (Field 9 Int64 -> Integer) -> Field 9 Int64 -> Integer
forall a b. (a -> b) -> a -> b
$ ColumnMetaData
meta.cmd_data_page_offset
        dictOffset :: Maybe Integer
dictOffset = Int64 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Integer) -> Maybe Int64 -> Maybe Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Field 11 (Maybe Int64) -> Maybe Int64
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField ColumnMetaData
meta.cmd_dictionary_page_offset
        offset :: Integer
offset = Integer -> Maybe Integer -> Integer
forall a. a -> Maybe a -> a
fromMaybe Integer
dataOffset Maybe Integer
dictOffset
        compLen :: Int
compLen = Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int) -> (Field 7 Int64 -> Int64) -> Field 7 Int64 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Field 7 Int64 -> Int64
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (Field 7 Int64 -> Int) -> Field 7 Int64 -> Int
forall a b. (a -> b) -> a -> b
$ ColumnMetaData
meta.cmd_total_compressed_size
    ByteString
rawBytes <- Range -> m ByteString
forall (m :: * -> *). RandomAccess m => Range -> m ByteString
readBytes (Integer -> Int -> Range
Range Integer
offset Int
compLen)
    (CompressionCodec, ThriftType, ByteString)
-> m (CompressionCodec, ThriftType, ByteString)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (CompressionCodec
codec, ThriftType
pType, ByteString
rawBytes)

{- | A decoded DATA page handed to a fold step: the running dictionary, the
page encoding, the present-value count, the (decompressed) value bytes, and the
definition/repetition level vectors.
-}
type RawPage =
    (Maybe DictVals, Encoding, Int, BS.ByteString, VU.Vector Int, VU.Vector Int)

{- | Left-fold a monadic step over every DATA page (V1 or V2) of every column
chunk, in order, threading the running dictionary internally and handing each
page to @step@ as a 'RawPage' (no value decoding — the step decides how to
materialize, e.g. into a typed vector or straight into a text buffer).

Dictionary pages update the running dictionary; @INDEX_PAGE@s are skipped. Only
one page's bytes are live at a time.

-- TODO: when a page index is available, use it here to compute which page
-- byte ranges to request from the RandomAccess layer instead of reading the
-- entire column chunk in one contiguous read.

-- TODO: accept an optional row-range and use the column/offset page index
-- (when present in file metadata) to skip pages whose row range does not
-- overlap the requested range, avoiding decompression of irrelevant pages
-- entirely.
-}
{-# INLINEABLE foldColumnDataPagesM #-}
foldColumnDataPagesM ::
    forall m acc.
    (RandomAccess m, MonadIO m) =>
    ColumnDescription ->
    [ColumnChunk] ->
    (acc -> RawPage -> m acc) ->
    acc ->
    m acc
foldColumnDataPagesM :: forall (m :: * -> *) acc.
(RandomAccess m, MonadIO m) =>
ColumnDescription
-> [ColumnChunk] -> (acc -> RawPage -> m acc) -> acc -> m acc
foldColumnDataPagesM ColumnDescription
description [ColumnChunk]
chunks acc -> RawPage -> m acc
step = [ColumnChunk] -> acc -> m acc
goChunks [ColumnChunk]
chunks
  where
    maxDef :: Int
maxDef = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ColumnDescription
description.maxDefinitionLevel :: Int
    maxRep :: Int
maxRep = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ColumnDescription
description.maxRepetitionLevel :: Int

    goChunks :: [ColumnChunk] -> acc -> m acc
goChunks [] !acc
acc = acc -> m acc
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return acc
acc
    goChunks (ColumnChunk
cc : [ColumnChunk]
ccs) !acc
acc = do
        (CompressionCodec
codec, ThriftType
pType, ByteString
rawBytes) <- ColumnChunk -> m (CompressionCodec, ThriftType, ByteString)
forall (m :: * -> *).
RandomAccess m =>
ColumnChunk -> m (CompressionCodec, ThriftType, ByteString)
readChunkBytes ColumnChunk
cc
        acc
acc' <- Maybe DictVals
-> CompressionCodec -> ThriftType -> ByteString -> acc -> m acc
goPages Maybe DictVals
forall a. Maybe a
Nothing CompressionCodec
codec ThriftType
pType ByteString
rawBytes acc
acc
        [ColumnChunk] -> acc -> m acc
goChunks [ColumnChunk]
ccs acc
acc'

    goPages :: Maybe DictVals
-> CompressionCodec -> ThriftType -> ByteString -> acc -> m acc
goPages Maybe DictVals
dict CompressionCodec
codec ThriftType
pType ByteString
bs !acc
acc
        | ByteString -> Bool
BS.null ByteString
bs = acc -> m acc
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return acc
acc
        | Bool
otherwise = case ByteString -> Either [Char] (ByteString, PageHeader)
parsePageHeader ByteString
bs of
            Left [Char]
e -> [Char] -> m acc
forall a. HasCallStack => [Char] -> a
error ([Char]
"foldColumnDataPagesM: failed to parse page header: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
e)
            Right (ByteString
rest, PageHeader
hdr) -> do
                let compSz :: Int
compSz = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> (Field 3 Int32 -> Int32) -> Field 3 Int32 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Field 3 Int32 -> Int32
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (Field 3 Int32 -> Int) -> Field 3 Int32 -> Int
forall a b. (a -> b) -> a -> b
$ PageHeader
hdr.ph_compressed_page_size
                    uncmpSz :: Int
uncmpSz = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> (Field 2 Int32 -> Int32) -> Field 2 Int32 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Field 2 Int32 -> Int32
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (Field 2 Int32 -> Int) -> Field 2 Int32 -> Int
forall a b. (a -> b) -> a -> b
$ PageHeader
hdr.ph_uncompressed_page_size
                    (ByteString
pageData, ByteString
rest') = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
compSz ByteString
rest
                case Field 1 PageType -> PageType
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField PageHeader
hdr.ph_type of
                    DICTIONARY_PAGE Enumeration 2
_ -> do
                        let dictHdr :: DictionaryPageHeader
dictHdr =
                                DictionaryPageHeader
-> Maybe DictionaryPageHeader -> DictionaryPageHeader
forall a. a -> Maybe a -> a
fromMaybe
                                    ([Char] -> DictionaryPageHeader
forall a. HasCallStack => [Char] -> a
error [Char]
"DICTIONARY_PAGE: missing dictionary page header")
                                    (Field 7 (Maybe DictionaryPageHeader) -> Maybe DictionaryPageHeader
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField PageHeader
hdr.ph_dictionary_page_header)
                            numVals :: Int32
numVals = Field 1 Int32 -> Int32
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField DictionaryPageHeader
dictHdr.diph_num_values
                        ByteString
decompressed <- IO ByteString -> m ByteString
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> m ByteString) -> IO ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ Int -> CompressionCodec -> ByteString -> IO ByteString
decompressData Int
uncmpSz CompressionCodec
codec ByteString
pageData
                        let d :: DictVals
d = ThriftType -> ByteString -> Int32 -> Maybe Int32 -> DictVals
readDictVals ThriftType
pType ByteString
decompressed Int32
numVals ColumnDescription
description.typeLength
                        Maybe DictVals
-> CompressionCodec -> ThriftType -> ByteString -> acc -> m acc
goPages (DictVals -> Maybe DictVals
forall a. a -> Maybe a
Just DictVals
d) CompressionCodec
codec ThriftType
pType ByteString
rest' acc
acc
                    DATA_PAGE Enumeration 0
_ -> do
                        let dph :: DataPageHeader
dph =
                                DataPageHeader -> Maybe DataPageHeader -> DataPageHeader
forall a. a -> Maybe a -> a
fromMaybe
                                    ([Char] -> DataPageHeader
forall a. HasCallStack => [Char] -> a
error [Char]
"DATA_PAGE: missing data page header")
                                    (Field 5 (Maybe DataPageHeader) -> Maybe DataPageHeader
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField PageHeader
hdr.ph_data_page_header)
                            n :: Int
n = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> (Field 1 Int32 -> Int32) -> Field 1 Int32 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Field 1 Int32 -> Int32
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (Field 1 Int32 -> Int) -> Field 1 Int32 -> Int
forall a b. (a -> b) -> a -> b
$ DataPageHeader
dph.dph_num_values
                            enc :: Encoding
enc = Field 2 Encoding -> Encoding
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField DataPageHeader
dph.dph_encoding
                        ByteString
decompressed <- IO ByteString -> m ByteString
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> m ByteString) -> IO ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ Int -> CompressionCodec -> ByteString -> IO ByteString
decompressData Int
uncmpSz CompressionCodec
codec ByteString
pageData
                        let (Vector Int
defLvls, Vector Int
repLvls, Int
nPresent, ByteString
valBytes) =
                                Int
-> Int
-> Int
-> ByteString
-> (Vector Int, Vector Int, Int, ByteString)
readLevelsV1 Int
n Int
maxDef Int
maxRep ByteString
decompressed
                        acc
acc' <- acc -> RawPage -> m acc
step acc
acc (Maybe DictVals
dict, Encoding
enc, Int
nPresent, ByteString
valBytes, Vector Int
defLvls, Vector Int
repLvls)
                        Maybe DictVals
-> CompressionCodec -> ThriftType -> ByteString -> acc -> m acc
goPages Maybe DictVals
dict CompressionCodec
codec ThriftType
pType ByteString
rest' acc
acc'
                    DATA_PAGE_V2 Enumeration 3
_ -> do
                        let dph2 :: DataPageHeaderV2
dph2 =
                                DataPageHeaderV2 -> Maybe DataPageHeaderV2 -> DataPageHeaderV2
forall a. a -> Maybe a -> a
fromMaybe
                                    ([Char] -> DataPageHeaderV2
forall a. HasCallStack => [Char] -> a
error [Char]
"DATA_PAGE_V2: missing data page header v2")
                                    (Field 8 (Maybe DataPageHeaderV2) -> Maybe DataPageHeaderV2
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField PageHeader
hdr.ph_data_page_header_v2)
                            n :: Int
n = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> (Field 1 Int32 -> Int32) -> Field 1 Int32 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Field 1 Int32 -> Int32
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (Field 1 Int32 -> Int) -> Field 1 Int32 -> Int
forall a b. (a -> b) -> a -> b
$ DataPageHeaderV2
dph2.dph2_num_values
                            enc :: Encoding
enc = Field 4 Encoding -> Encoding
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField DataPageHeaderV2
dph2.dph2_encoding
                            defLen :: Int32
defLen = Field 5 Int32 -> Int32
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField DataPageHeaderV2
dph2.dph2_definition_levels_byte_length
                            repLen :: Int32
repLen = Field 6 Int32 -> Int32
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField DataPageHeaderV2
dph2.dph2_repetition_levels_byte_length
                            -- V2: levels are never compressed; only the value
                            -- payload is (optionally) compressed.
                            isCompressed :: Bool
isCompressed = Bool -> Maybe Bool -> Bool
forall a. a -> Maybe a -> a
fromMaybe Bool
True (Field 7 (Maybe Bool) -> Maybe Bool
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField DataPageHeaderV2
dph2.dph2_is_compressed)
                            (Vector Int
defLvls, Vector Int
repLvls, Int
nPresent, ByteString
compValBytes) =
                                Int
-> Int
-> Int
-> Int32
-> Int32
-> ByteString
-> (Vector Int, Vector Int, Int, ByteString)
readLevelsV2 Int
n Int
maxDef Int
maxRep Int32
repLen Int32
defLen ByteString
pageData
                        ByteString
valBytes <-
                            if Bool
isCompressed
                                then IO ByteString -> m ByteString
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> m ByteString) -> IO ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ Int -> CompressionCodec -> ByteString -> IO ByteString
decompressData Int
uncmpSz CompressionCodec
codec ByteString
compValBytes
                                else ByteString -> m ByteString
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
compValBytes
                        acc
acc' <- acc -> RawPage -> m acc
step acc
acc (Maybe DictVals
dict, Encoding
enc, Int
nPresent, ByteString
valBytes, Vector Int
defLvls, Vector Int
repLvls)
                        Maybe DictVals
-> CompressionCodec -> ThriftType -> ByteString -> acc -> m acc
goPages Maybe DictVals
dict CompressionCodec
codec ThriftType
pType ByteString
rest' acc
acc'
                    INDEX_PAGE Enumeration 1
_ -> Maybe DictVals
-> CompressionCodec -> ThriftType -> ByteString -> acc -> m acc
goPages Maybe DictVals
dict CompressionCodec
codec ThriftType
pType ByteString
rest' acc
acc

{- | Left-fold over per-page value triples, decoding each page with @decoder@.
A thin wrapper over 'foldColumnDataPagesM' for the typed (numeric/boxed) column
paths.
-}
{-# INLINEABLE foldColumnPagesM #-}
foldColumnPagesM ::
    forall m v a acc.
    (RandomAccess m, MonadIO m, VG.Vector v a) =>
    ColumnDescription ->
    (Maybe DictVals -> Encoding -> Int -> BS.ByteString -> v a) ->
    [ColumnChunk] ->
    (acc -> (v a, VU.Vector Int, VU.Vector Int) -> m acc) ->
    acc ->
    m acc
foldColumnPagesM :: forall (m :: * -> *) (v :: * -> *) a acc.
(RandomAccess m, MonadIO m, Vector v a) =>
ColumnDescription
-> (Maybe DictVals -> Encoding -> Int -> ByteString -> v a)
-> [ColumnChunk]
-> (acc -> (v a, Vector Int, Vector Int) -> m acc)
-> acc
-> m acc
foldColumnPagesM ColumnDescription
description Maybe DictVals -> Encoding -> Int -> ByteString -> v a
decoder [ColumnChunk]
chunks acc -> (v a, Vector Int, Vector Int) -> m acc
step =
    ColumnDescription
-> [ColumnChunk] -> (acc -> RawPage -> m acc) -> acc -> m acc
forall (m :: * -> *) acc.
(RandomAccess m, MonadIO m) =>
ColumnDescription
-> [ColumnChunk] -> (acc -> RawPage -> m acc) -> acc -> m acc
foldColumnDataPagesM ColumnDescription
description [ColumnChunk]
chunks ((acc -> RawPage -> m acc) -> acc -> m acc)
-> (acc -> RawPage -> m acc) -> acc -> m acc
forall a b. (a -> b) -> a -> b
$
        \acc
acc (Maybe DictVals
dict, Encoding
enc, Int
nPresent, ByteString
valBytes, Vector Int
defLvls, Vector Int
repLvls) ->
            acc -> (v a, Vector Int, Vector Int) -> m acc
step acc
acc (Maybe DictVals -> Encoding -> Int -> ByteString -> v a
decoder Maybe DictVals
dict Encoding
enc Int
nPresent ByteString
valBytes, Vector Int
defLvls, Vector Int
repLvls)

{- | Append a non-nullable BYTE_ARRAY page's strings straight into a
'TextBuilder', no intermediate boxed 'Text' vector.

PLAIN pages copy each value's UTF-8 bytes directly from the page buffer pointer
('appendTextSliceFromPtr'), skipping the per-value 'decodeUtf8Lenient' +
'ByteString' slicing + 'Text' allocation that @readNTexts@ would do.
Dictionary pages append the shared dictionary 'Text's by reference.
-}
appendStringPageIO ::
    TextBuilder RealWorld ->
    Maybe DictVals ->
    Encoding ->
    Int ->
    BS.ByteString ->
    IO ()
appendStringPageIO :: TextBuilder RealWorld
-> Maybe DictVals -> Encoding -> Int -> ByteString -> IO ()
appendStringPageIO TextBuilder RealWorld
builder Maybe DictVals
mDict Encoding
enc Int
nPresent ByteString
bs = case Encoding
enc of
    PLAIN Enumeration 0
_ -> ByteString -> (CStringLen -> IO ()) -> IO ()
forall a. ByteString -> (CStringLen -> IO a) -> IO a
BSU.unsafeUseAsCStringLen ByteString
bs ((CStringLen -> IO ()) -> IO ()) -> (CStringLen -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
cptr, Int
_len) -> do
        let p :: Ptr Word8
p = Ptr CChar -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr CChar
cptr :: Ptr Word8
            go :: Int -> Int -> IO ()
go !Int
_ !Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nPresent = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            go !Int
off !Int
i = do
                Int
len <- Ptr Word8 -> Int -> IO Int
readLenLE Ptr Word8
p Int
off
                ST RealWorld () -> IO ()
forall a. ST RealWorld a -> IO a
stToIO (TextBuilder RealWorld -> Ptr Word8 -> Int -> ST RealWorld ()
forall s. TextBuilder s -> Ptr Word8 -> Int -> ST s ()
appendTextSliceFromPtr TextBuilder RealWorld
builder (Ptr Word8
p Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4)) Int
len)
                Int -> Int -> IO ()
go (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        Int -> Int -> IO ()
go Int
0 Int
0
    RLE_DICTIONARY Enumeration 8
_ -> IO ()
dictAppend
    PLAIN_DICTIONARY Enumeration 2
_ -> IO ()
dictAppend
    Encoding
_ -> [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char]
"appendStringPageIO: unsupported encoding " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Encoding -> [Char]
forall a. Show a => a -> [Char]
show Encoding
enc)
  where
    dictAppend :: IO ()
dictAppend = case Maybe DictVals
mDict of
        Just (DText Vector Text
ds) ->
            let (Vector Int
idxs, ByteString
_) = Int -> ByteString -> (Vector Int, ByteString)
decodeDictIndices Int
nPresent ByteString
bs
             in ST RealWorld () -> IO ()
forall a. ST RealWorld a -> IO a
stToIO ((Int -> ST RealWorld ()) -> Vector Int -> ST RealWorld ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
(a -> m b) -> Vector a -> m ()
VU.mapM_ (\Int
i -> TextBuilder RealWorld -> Text -> ST RealWorld ()
forall s. TextBuilder s -> Text -> ST s ()
appendText TextBuilder RealWorld
builder (Vector Text
ds Vector Text -> Int -> Text
forall a. Vector a -> Int -> a
VB.! Int
i)) Vector Int
idxs)
        Just DictVals
d -> [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char]
"appendStringPageIO: wrong dict type, got " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ DictVals -> [Char]
forall a. Show a => a -> [Char]
show DictVals
d)
        Maybe DictVals
Nothing -> [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"appendStringPageIO: dictionary-encoded page but no dictionary seen"

{- | Read a little-endian 4-byte length prefix at byte @off@ from a raw page
pointer.
-}
readLenLE :: Ptr Word8 -> Int -> IO Int
readLenLE :: Ptr Word8 -> Int -> IO Int
readLenLE Ptr Word8
p Int
off = do
    Word8
b0 <- Ptr Word8 -> Int -> IO Word8
forall b. Ptr b -> Int -> IO Word8
forall a b. Storable a => Ptr b -> Int -> IO a
peekByteOff Ptr Word8
p Int
off :: IO Word8
    Word8
b1 <- Ptr Word8 -> Int -> IO Word8
forall b. Ptr b -> Int -> IO Word8
forall a b. Storable a => Ptr b -> Int -> IO a
peekByteOff Ptr Word8
p (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) :: IO Word8
    Word8
b2 <- Ptr Word8 -> Int -> IO Word8
forall b. Ptr b -> Int -> IO Word8
forall a b. Storable a => Ptr b -> Int -> IO a
peekByteOff Ptr Word8
p (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) :: IO Word8
    Word8
b3 <- Ptr Word8 -> Int -> IO Word8
forall b. Ptr b -> Int -> IO Word8
forall a b. Storable a => Ptr b -> Int -> IO a
peekByteOff Ptr Word8
p (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3) :: IO Word8
    Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$
        Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b0
            Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. (Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
8)
            Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. (Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b2 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
16)
            Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. (Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b3 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
24)
{-# INLINE readLenLE #-}

{- | Append a NULLABLE BYTE_ARRAY page into a 'TextBuilder', interleaving nulls
by walking the definition levels: a present row (@def == maxDef@) consumes the
next stored value (dictionary entry by reference, or the next length-prefixed
PLAIN slice by memcpy); a null row ('appendNull') writes an offset and clears a
validity bit. Builds 'PackedText' + validity bitmap with no per-row 'Text' or
@[Maybe a]@ allocation. Only present values are stored in the page payload.
-}
appendNullableStringPageIO ::
    TextBuilder RealWorld ->
    Int ->
    Maybe DictVals ->
    Encoding ->
    Int ->
    BS.ByteString ->
    VU.Vector Int ->
    IO ()
appendNullableStringPageIO :: TextBuilder RealWorld
-> Int
-> Maybe DictVals
-> Encoding
-> Int
-> ByteString
-> Vector Int
-> IO ()
appendNullableStringPageIO TextBuilder RealWorld
builder Int
maxDef Maybe DictVals
mDict Encoding
enc Int
nPresent ByteString
bs Vector Int
defs = case Encoding
enc of
    PLAIN Enumeration 0
_ -> ByteString -> (CStringLen -> IO ()) -> IO ()
forall a. ByteString -> (CStringLen -> IO a) -> IO a
BSU.unsafeUseAsCStringLen ByteString
bs ((CStringLen -> IO ()) -> IO ()) -> (CStringLen -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
cptr, Int
_len) -> do
        let p :: Ptr Word8
p = Ptr CChar -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr CChar
cptr :: Ptr Word8
            go :: Int -> Int -> IO ()
go !Int
_ !Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nDefs = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            go !Int
off !Int
i
                | Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
defs Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
maxDef = do
                    Int
len <- Ptr Word8 -> Int -> IO Int
readLenLE Ptr Word8
p Int
off
                    ST RealWorld () -> IO ()
forall a. ST RealWorld a -> IO a
stToIO (TextBuilder RealWorld -> Ptr Word8 -> Int -> ST RealWorld ()
forall s. TextBuilder s -> Ptr Word8 -> Int -> ST s ()
appendTextSliceFromPtr TextBuilder RealWorld
builder (Ptr Word8
p Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4)) Int
len)
                    Int -> Int -> IO ()
go (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                | Bool
otherwise = ST RealWorld () -> IO ()
forall a. ST RealWorld a -> IO a
stToIO (TextBuilder RealWorld -> ST RealWorld ()
forall s. TextBuilder s -> ST s ()
forall (b :: * -> *) s. ColumnBuilder b => b s -> ST s ()
appendNull TextBuilder RealWorld
builder) IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> IO ()
go Int
off (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        Int -> Int -> IO ()
go Int
0 Int
0
    RLE_DICTIONARY Enumeration 8
_ -> IO ()
dictAppend
    PLAIN_DICTIONARY Enumeration 2
_ -> IO ()
dictAppend
    Encoding
_ -> [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char]
"appendNullableStringPageIO: unsupported encoding " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Encoding -> [Char]
forall a. Show a => a -> [Char]
show Encoding
enc)
  where
    nDefs :: Int
nDefs = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
defs
    dictAppend :: IO ()
dictAppend = case Maybe DictVals
mDict of
        Just (DText Vector Text
ds) -> do
            let (Vector Int
idxs, ByteString
_) = Int -> ByteString -> (Vector Int, ByteString)
decodeDictIndices Int
nPresent ByteString
bs
                go :: Int -> Int -> IO ()
go !Int
_ !Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nDefs = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                go !Int
j !Int
i
                    | Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
defs Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
maxDef =
                        ST RealWorld () -> IO ()
forall a. ST RealWorld a -> IO a
stToIO (TextBuilder RealWorld -> Text -> ST RealWorld ()
forall s. TextBuilder s -> Text -> ST s ()
appendText TextBuilder RealWorld
builder (Vector Text
ds Vector Text -> Int -> Text
forall a. Vector a -> Int -> a
VB.! Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
idxs Int
j))
                            IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> IO ()
go (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    | Bool
otherwise = ST RealWorld () -> IO ()
forall a. ST RealWorld a -> IO a
stToIO (TextBuilder RealWorld -> ST RealWorld ()
forall s. TextBuilder s -> ST s ()
forall (b :: * -> *) s. ColumnBuilder b => b s -> ST s ()
appendNull TextBuilder RealWorld
builder) IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> IO ()
go Int
j (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            Int -> Int -> IO ()
go Int
0 Int
0
        Just DictVals
d -> [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char]
"appendNullableStringPageIO: wrong dict type, got " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ DictVals -> [Char]
forall a. Show a => a -> [Char]
show DictVals
d)
        Maybe DictVals
Nothing ->
            [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error
                [Char]
"appendNullableStringPageIO: dictionary-encoded page but no dictionary seen"

-- ---------------------------------------------------------------------------
-- Page header parsing
-- ---------------------------------------------------------------------------

parsePageHeader :: BS.ByteString -> Either String (BS.ByteString, PageHeader)
parsePageHeader :: ByteString -> Either [Char] (ByteString, PageHeader)
parsePageHeader = Protocol -> ByteString -> Either [Char] (ByteString, PageHeader)
forall a.
Pinchable a =>
Protocol -> ByteString -> Either [Char] (ByteString, a)
decodeWithLeftovers Protocol
Pinch.compactProtocol

-- ---------------------------------------------------------------------------
-- Batch value readers
-- ---------------------------------------------------------------------------

readNInt32 :: Int -> BS.ByteString -> VU.Vector Int32
readNInt32 :: Int -> ByteString -> Vector Int32
readNInt32 Int
n ByteString
bs = Int -> (Int -> Int32) -> Vector Int32
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n ((Int -> Int32) -> Vector Int32) -> (Int -> Int32) -> Vector Int32
forall a b. (a -> b) -> a -> b
$ \Int
i -> ByteString -> Int32
littleEndianInt32 (Int -> ByteString -> ByteString
BS.drop (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i) ByteString
bs)

readNInt64 :: Int -> BS.ByteString -> VU.Vector Int64
readNInt64 :: Int -> ByteString -> Vector Int64
readNInt64 Int
n ByteString
bs = Int -> (Int -> Int64) -> Vector Int64
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n ((Int -> Int64) -> Vector Int64) -> (Int -> Int64) -> Vector Int64
forall a b. (a -> b) -> a -> b
$ \Int
i ->
    Word64 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Word64
littleEndianWord64 (Int -> ByteString -> ByteString
BS.drop (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i) ByteString
bs))

readNInt96 :: Int -> BS.ByteString -> [UTCTime]
readNInt96 :: Int -> ByteString -> [UTCTime]
readNInt96 Int
0 ByteString
_ = []
readNInt96 Int
n ByteString
bs = ByteString -> UTCTime
int96ToUTCTime (Int -> ByteString -> ByteString
BS.take Int
12 ByteString
bs) UTCTime -> [UTCTime] -> [UTCTime]
forall a. a -> [a] -> [a]
: Int -> ByteString -> [UTCTime]
readNInt96 (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int -> ByteString -> ByteString
BS.drop Int
12 ByteString
bs)

readNFloat :: Int -> BS.ByteString -> VU.Vector Float
readNFloat :: Int -> ByteString -> Vector Float
readNFloat Int
n ByteString
bs = Int -> (Int -> Float) -> Vector Float
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n ((Int -> Float) -> Vector Float) -> (Int -> Float) -> Vector Float
forall a b. (a -> b) -> a -> b
$ \Int
i ->
    Word32 -> Float
castWord32ToFloat (ByteString -> Word32
littleEndianWord32 (Int -> ByteString -> ByteString
BS.drop (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i) ByteString
bs))

readNDouble :: Int -> BS.ByteString -> VU.Vector Double
readNDouble :: Int -> ByteString -> Vector Double
readNDouble Int
n ByteString
bs = Int -> (Int -> Double) -> Vector Double
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n ((Int -> Double) -> Vector Double)
-> (Int -> Double) -> Vector Double
forall a b. (a -> b) -> a -> b
$ \Int
i ->
    Word64 -> Double
castWord64ToDouble (ByteString -> Word64
littleEndianWord64 (Int -> ByteString -> ByteString
BS.drop (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i) ByteString
bs))

readNTexts :: Int -> BS.ByteString -> [T.Text]
readNTexts :: Int -> ByteString -> [Text]
readNTexts Int
0 ByteString
_ = []
readNTexts Int
n ByteString
bs =
    let len :: Int
len = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> (ByteString -> Int32) -> ByteString -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Int32
littleEndianInt32 (ByteString -> Int32)
-> (ByteString -> ByteString) -> ByteString -> Int32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ByteString -> ByteString
BS.take Int
4 (ByteString -> Int) -> ByteString -> Int
forall a b. (a -> b) -> a -> b
$ ByteString
bs
        text :: Text
text = ByteString -> Text
decodeUtf8Lenient (ByteString -> Text)
-> (ByteString -> ByteString) -> ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ByteString -> ByteString
BS.take Int
len (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ByteString -> ByteString
BS.drop Int
4 (ByteString -> Text) -> ByteString -> Text
forall a b. (a -> b) -> a -> b
$ ByteString
bs
     in Text
text Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Int -> ByteString -> [Text]
readNTexts (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int -> ByteString -> ByteString
BS.drop (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) ByteString
bs)

readNFixedTexts :: Int -> Int -> BS.ByteString -> [T.Text]
readNFixedTexts :: Int -> Int -> ByteString -> [Text]
readNFixedTexts Int
_ Int
0 ByteString
_ = []
readNFixedTexts Int
len Int
n ByteString
bs =
    ByteString -> Text
decodeUtf8Lenient (Int -> ByteString -> ByteString
BS.take Int
len ByteString
bs)
        Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Int -> Int -> ByteString -> [Text]
readNFixedTexts Int
len (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int -> ByteString -> ByteString
BS.drop Int
len ByteString
bs)