{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

module DataFrame.IO.Parquet.Utils (
    ColumnDescription (..),
    generateColumnDescriptions,
    getColumnNames,
    foldNonNullable,
    foldNonNullableUnboxed,
    foldNullable,
    foldNullableUnboxed,
    foldRepeated,
    foldRepeatedUnboxed,
) where

import Control.Monad.IO.Class (MonadIO (..))
import Data.Int (Int32)
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Vector as VB
import qualified Data.Vector.Mutable as VBM
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import Data.Word (Word8)
import DataFrame.IO.Parquet.Levels (
    stitchList,
    stitchList2,
    stitchList3,
 )
import DataFrame.IO.Parquet.Thrift (
    ConvertedType (..),
    FieldRepetitionType (..),
    LogicalType (..),
    SchemaElement (..),
    ThriftType,
    unField,
 )
import DataFrame.Internal.Column (
    Column (..),
    Columnable,
    buildBitmapFromValid,
    fromList,
 )

{- | A left-fold driver over a column's per-page triples
@(values, def-levels, rep-levels)@, as produced by
'DataFrame.IO.Parquet.Page.foldColumnPagesM'. Generic in the accumulator so a
single driver serves every fold strategy below.
-}
type PageFold m v a =
    forall acc.
    (acc -> (v a, VU.Vector Int, VU.Vector Int) -> m acc) -> acc -> m acc

data ColumnDescription = ColumnDescription
    { ColumnDescription -> Maybe ThriftType
colElementType :: !(Maybe ThriftType)
    , ColumnDescription -> Int32
maxDefinitionLevel :: !Int32
    , ColumnDescription -> Int32
maxRepetitionLevel :: !Int32
    , ColumnDescription -> Maybe LogicalType
colLogicalType :: !(Maybe LogicalType)
    , ColumnDescription -> Maybe ConvertedType
colConvertedType :: !(Maybe ConvertedType)
    , ColumnDescription -> Maybe Int32
typeLength :: !(Maybe Int32)
    }
    deriving (Int -> ColumnDescription -> ShowS
[ColumnDescription] -> ShowS
ColumnDescription -> String
(Int -> ColumnDescription -> ShowS)
-> (ColumnDescription -> String)
-> ([ColumnDescription] -> ShowS)
-> Show ColumnDescription
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ColumnDescription -> ShowS
showsPrec :: Int -> ColumnDescription -> ShowS
$cshow :: ColumnDescription -> String
show :: ColumnDescription -> String
$cshowList :: [ColumnDescription] -> ShowS
showList :: [ColumnDescription] -> ShowS
Show, ColumnDescription -> ColumnDescription -> Bool
(ColumnDescription -> ColumnDescription -> Bool)
-> (ColumnDescription -> ColumnDescription -> Bool)
-> Eq ColumnDescription
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ColumnDescription -> ColumnDescription -> Bool
== :: ColumnDescription -> ColumnDescription -> Bool
$c/= :: ColumnDescription -> ColumnDescription -> Bool
/= :: ColumnDescription -> ColumnDescription -> Bool
Eq)

levelContribution :: Maybe FieldRepetitionType -> (Int, Int)
levelContribution :: Maybe FieldRepetitionType -> (Int, Int)
levelContribution = \case
    Just (REPEATED Enumeration 2
_) -> (Int
1, Int
1)
    Just (OPTIONAL Enumeration 1
_) -> (Int
1, Int
0)
    Maybe FieldRepetitionType
_ -> (Int
0, Int
0) -- REQUIRED or absent

data SchemaTree = SchemaTree SchemaElement [SchemaTree]

buildTree :: [SchemaElement] -> (SchemaTree, [SchemaElement])
buildTree :: [SchemaElement] -> (SchemaTree, [SchemaElement])
buildTree [] = String -> (SchemaTree, [SchemaElement])
forall a. HasCallStack => String -> a
error String
"buildTree: schema ended unexpectedly"
buildTree (SchemaElement
se : [SchemaElement]
rest) =
    let n :: Int
n = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> Int32 -> Int
forall a b. (a -> b) -> a -> b
$ Int32 -> Maybe Int32 -> Int32
forall a. a -> Maybe a -> a
fromMaybe Int32
0 (Field 5 (Maybe Int32) -> Maybe Int32
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (SchemaElement -> Field 5 (Maybe Int32)
num_children SchemaElement
se)) :: Int
        ([SchemaTree]
children, [SchemaElement]
rest') = Int -> [SchemaElement] -> ([SchemaTree], [SchemaElement])
buildChildren Int
n [SchemaElement]
rest
     in (SchemaElement -> [SchemaTree] -> SchemaTree
SchemaTree SchemaElement
se [SchemaTree]
children, [SchemaElement]
rest')

-- | Build a forest of sibling trees from a flat depth-first element list.
buildForest :: [SchemaElement] -> ([SchemaTree], [SchemaElement])
buildForest :: [SchemaElement] -> ([SchemaTree], [SchemaElement])
buildForest [] = ([], [])
buildForest [SchemaElement]
xs =
    let (SchemaTree
tree, [SchemaElement]
rest') = [SchemaElement] -> (SchemaTree, [SchemaElement])
buildTree [SchemaElement]
xs
        ([SchemaTree]
siblings, [SchemaElement]
rest'') = [SchemaElement] -> ([SchemaTree], [SchemaElement])
buildForest [SchemaElement]
rest'
     in (SchemaTree
tree SchemaTree -> [SchemaTree] -> [SchemaTree]
forall a. a -> [a] -> [a]
: [SchemaTree]
siblings, [SchemaElement]
rest'')

-- | Build exactly @n@ child trees, each consuming only its own subtree.
buildChildren :: Int -> [SchemaElement] -> ([SchemaTree], [SchemaElement])
buildChildren :: Int -> [SchemaElement] -> ([SchemaTree], [SchemaElement])
buildChildren Int
0 [SchemaElement]
xs = ([], [SchemaElement]
xs)
buildChildren Int
n [SchemaElement]
xs =
    let (SchemaTree
child, [SchemaElement]
rest') = [SchemaElement] -> (SchemaTree, [SchemaElement])
buildTree [SchemaElement]
xs
        ([SchemaTree]
siblings, [SchemaElement]
rest'') = Int -> [SchemaElement] -> ([SchemaTree], [SchemaElement])
buildChildren (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [SchemaElement]
rest'
     in (SchemaTree
child SchemaTree -> [SchemaTree] -> [SchemaTree]
forall a. a -> [a] -> [a]
: [SchemaTree]
siblings, [SchemaElement]
rest'')

collectLeaves :: Int -> Int -> SchemaTree -> [ColumnDescription]
collectLeaves :: Int -> Int -> SchemaTree -> [ColumnDescription]
collectLeaves Int
defAcc Int
repAcc (SchemaTree SchemaElement
se [SchemaTree]
children) =
    let (Int
dInc, Int
rInc) = Maybe FieldRepetitionType -> (Int, Int)
levelContribution (Field 3 (Maybe FieldRepetitionType) -> Maybe FieldRepetitionType
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (SchemaElement -> Field 3 (Maybe FieldRepetitionType)
repetition_type SchemaElement
se))
        defLevel :: Int
defLevel = Int
defAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
dInc
        repLevel :: Int
repLevel = Int
repAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rInc
     in case [SchemaTree]
children of
            [] ->
                -- leaf: emit a description
                let pType :: Maybe ThriftType
pType = Field 1 (Maybe ThriftType) -> Maybe ThriftType
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (SchemaElement -> Field 1 (Maybe ThriftType)
schematype SchemaElement
se)
                 in [ Maybe ThriftType
-> Int32
-> Int32
-> Maybe LogicalType
-> Maybe ConvertedType
-> Maybe Int32
-> ColumnDescription
ColumnDescription
                        Maybe ThriftType
pType
                        (Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
defLevel)
                        (Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
repLevel)
                        (Field 10 (Maybe LogicalType) -> Maybe LogicalType
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (SchemaElement -> Field 10 (Maybe LogicalType)
logicalType SchemaElement
se))
                        (Field 6 (Maybe ConvertedType) -> Maybe ConvertedType
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (SchemaElement -> Field 6 (Maybe ConvertedType)
converted_type SchemaElement
se))
                        (Field 2 (Maybe Int32) -> Maybe Int32
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (SchemaElement -> Field 2 (Maybe Int32)
type_length SchemaElement
se))
                    ]
            [SchemaTree]
_ ->
                -- internal node: recurse into children
                (SchemaTree -> [ColumnDescription])
-> [SchemaTree] -> [ColumnDescription]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Int -> Int -> SchemaTree -> [ColumnDescription]
collectLeaves Int
defLevel Int
repLevel) [SchemaTree]
children

generateColumnDescriptions :: [SchemaElement] -> [ColumnDescription]
generateColumnDescriptions :: [SchemaElement] -> [ColumnDescription]
generateColumnDescriptions [] = []
generateColumnDescriptions (SchemaElement
_ : [SchemaElement]
rest) =
    -- drop schema root
    let ([SchemaTree]
forest, [SchemaElement]
_) = [SchemaElement] -> ([SchemaTree], [SchemaElement])
buildForest [SchemaElement]
rest
     in (SchemaTree -> [ColumnDescription])
-> [SchemaTree] -> [ColumnDescription]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Int -> Int -> SchemaTree -> [ColumnDescription]
collectLeaves Int
0 Int
0) [SchemaTree]
forest

getColumnNames :: [SchemaElement] -> [Text]
getColumnNames :: [SchemaElement] -> [Text]
getColumnNames [] = []
getColumnNames [SchemaElement]
schemaElements =
    let ([SchemaTree]
forest, [SchemaElement]
_) = [SchemaElement] -> ([SchemaTree], [SchemaElement])
buildForest [SchemaElement]
schemaElements
     in [SchemaTree] -> [Text] -> Bool -> [Text]
go [SchemaTree]
forest [] Bool
False
  where
    isRepeated :: SchemaElement -> Bool
isRepeated SchemaElement
se = case Field 3 (Maybe FieldRepetitionType) -> Maybe FieldRepetitionType
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (SchemaElement -> Field 3 (Maybe FieldRepetitionType)
repetition_type SchemaElement
se) of
        Just (REPEATED Enumeration 2
_) -> Bool
True
        Maybe FieldRepetitionType
_ -> Bool
False

    go :: [SchemaTree] -> [Text] -> Bool -> [Text]
go [] [Text]
_ Bool
_ = []
    go (SchemaTree SchemaElement
se [SchemaTree]
children : [SchemaTree]
rest) [Text]
path Bool
skipThis =
        case [SchemaTree]
children of
            -- Leaf node
            [] ->
                let newPath :: [Text]
newPath = if Bool
skipThis then [Text]
path else [Text]
path [Text] -> [Text] -> [Text]
forall a. [a] -> [a] -> [a]
++ [Field 4 Text -> Text
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (SchemaElement -> Field 4 Text
name SchemaElement
se)]
                    fullName :: Text
fullName = Text -> [Text] -> Text
T.intercalate Text
"." [Text]
newPath
                 in Text
fullName Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: [SchemaTree] -> [Text] -> Bool -> [Text]
go [SchemaTree]
rest [Text]
path Bool
skipThis
            -- REPEATED intermediate: skip this name; skip single child too
            [SchemaTree]
_
                | SchemaElement -> Bool
isRepeated SchemaElement
se ->
                    let skipChildren :: Bool
skipChildren = [SchemaTree] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SchemaTree]
children Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
                        childLeaves :: [Text]
childLeaves = [SchemaTree] -> [Text] -> Bool -> [Text]
go [SchemaTree]
children [Text]
path Bool
skipChildren
                     in [Text]
childLeaves [Text] -> [Text] -> [Text]
forall a. [a] -> [a] -> [a]
++ [SchemaTree] -> [Text] -> Bool -> [Text]
go [SchemaTree]
rest [Text]
path Bool
skipThis
            -- Name-skipped intermediate: recurse with skip cleared
            [SchemaTree]
_
                | Bool
skipThis ->
                    let childLeaves :: [Text]
childLeaves = [SchemaTree] -> [Text] -> Bool -> [Text]
go [SchemaTree]
children [Text]
path Bool
False
                     in [Text]
childLeaves [Text] -> [Text] -> [Text]
forall a. [a] -> [a] -> [a]
++ [SchemaTree] -> [Text] -> Bool -> [Text]
go [SchemaTree]
rest [Text]
path Bool
skipThis
            -- Normal intermediate: add name to path, recurse
            [SchemaTree]
_ ->
                let subPath :: [Text]
subPath = [Text]
path [Text] -> [Text] -> [Text]
forall a. [a] -> [a] -> [a]
++ [Field 4 Text -> Text
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (SchemaElement -> Field 4 Text
name SchemaElement
se)]
                    childLeaves :: [Text]
childLeaves = [SchemaTree] -> [Text] -> Bool -> [Text]
go [SchemaTree]
children [Text]
subPath Bool
False
                 in [Text]
childLeaves [Text] -> [Text] -> [Text]
forall a. [a] -> [a] -> [a]
++ [SchemaTree] -> [Text] -> Bool -> [Text]
go [SchemaTree]
rest [Text]
path Bool
skipThis

{- | Fold a column's value pages into a non-nullable 'Column'.

Pre-allocates a mutable vector of @totalRows@ and fills it page-by-page via a
single streaming left fold ('PageFold'), avoiding any intermediate list or
concatenation allocation. Only one page's values are live at a time.
-}
{-# INLINEABLE foldNonNullable #-}
foldNonNullable ::
    forall m a.
    (MonadIO m, Columnable a) =>
    Int ->
    PageFold m VB.Vector a ->
    m Column
foldNonNullable :: forall (m :: * -> *) a.
(MonadIO m, Columnable a) =>
Int -> PageFold m Vector a -> m Column
foldNonNullable Int
totalRows PageFold m Vector a
runPages = do
    MVector RealWorld a
mv <- IO (MVector RealWorld a) -> m (MVector RealWorld a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MVector RealWorld a) -> m (MVector RealWorld a))
-> IO (MVector RealWorld a) -> m (MVector RealWorld a)
forall a b. (a -> b) -> a -> b
$ Int -> IO (MVector (PrimState IO) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
VBM.unsafeNew Int
totalRows
    Int
_ <-
        (Int -> (Vector a, Vector Int, Vector Int) -> m Int)
-> Int -> m Int
PageFold m Vector a
runPages
            ( \Int
off (Vector a
chunk, Vector Int
_, Vector Int
_) -> IO Int -> m Int
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> m Int) -> IO Int -> m Int
forall a b. (a -> b) -> a -> b
$ do
                let n :: Int
n = Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
chunk
                MVector (PrimState IO) a -> Vector a -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Vector a -> m ()
VB.copy (Int -> Int -> MVector RealWorld a -> MVector RealWorld a
forall s a. Int -> Int -> MVector s a -> MVector s a
VBM.unsafeSlice Int
off Int
n MVector RealWorld a
mv) Vector a
chunk
                Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)
            )
            (Int
0 :: Int)
    Vector a
v <- IO (Vector a) -> m (Vector a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Vector a) -> m (Vector a)) -> IO (Vector a) -> m (Vector a)
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) a -> IO (Vector a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
VB.unsafeFreeze MVector RealWorld a
MVector (PrimState IO) a
mv
    Column -> m Column
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
forall a. Maybe a
Nothing Vector a
v)

{-# INLINEABLE foldNonNullableUnboxed #-}
foldNonNullableUnboxed ::
    forall m a.
    (MonadIO m, Columnable a, VU.Unbox a) =>
    Int ->
    PageFold m VU.Vector a ->
    m Column
foldNonNullableUnboxed :: forall (m :: * -> *) a.
(MonadIO m, Columnable a, Unbox a) =>
Int -> PageFold m Vector a -> m Column
foldNonNullableUnboxed Int
totalRows PageFold m Vector a
runPages = do
    MVector RealWorld a
mv <- IO (MVector RealWorld a) -> m (MVector RealWorld a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MVector RealWorld a) -> m (MVector RealWorld a))
-> IO (MVector RealWorld a) -> m (MVector RealWorld a)
forall a b. (a -> b) -> a -> b
$ Int -> IO (MVector (PrimState IO) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
totalRows
    Int
_ <-
        (Int -> (Vector a, Vector Int, Vector Int) -> m Int)
-> Int -> m Int
PageFold m Vector a
runPages
            ( \Int
off (Vector a
chunk, Vector Int
_, Vector Int
_) -> IO Int -> m Int
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> m Int) -> IO Int -> m Int
forall a b. (a -> b) -> a -> b
$ do
                let n :: Int
n = Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
chunk
                    go :: Int -> IO ()
go Int
i
                        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                        | Bool
otherwise = do
                            MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite
                                MVector RealWorld a
MVector (PrimState IO) a
mv
                                (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i)
                                (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
chunk Int
i)
                            Int -> IO ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                Int -> IO ()
go Int
0
                Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)
            )
            (Int
0 :: Int)
    Vector a
dat <- IO (Vector a) -> m (Vector a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Vector a) -> m (Vector a)) -> IO (Vector a) -> m (Vector a)
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) a -> IO (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld a
MVector (PrimState IO) a
mv
    Column -> m Column
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
forall a. Maybe a
Nothing Vector a
dat)

{- | Fold a column's (values, def-levels) pages into a nullable 'Column'.

Pre-allocates the output buffer and a valid-mask vector of @totalRows@, then
scatters values inline during a single streaming left fold ('PageFold').

A 'hasNull' flag is accumulated during the scatter so the
'buildBitmapFromValid' call is skipped entirely when all values are present.
-}
{-# INLINEABLE foldNullable #-}
foldNullable ::
    forall m a.
    (MonadIO m, Columnable a) =>
    Int ->
    Int ->
    PageFold m VB.Vector a ->
    m Column
foldNullable :: forall (m :: * -> *) a.
(MonadIO m, Columnable a) =>
Int -> Int -> PageFold m Vector a -> m Column
foldNullable Int
maxDef Int
totalRows PageFold m Vector a
runPages = do
    -- null slots hold an error thunk, guarded by bitmap.
    --
    -- IMPORTANT: 'VBM.unsafeWrite' for boxed vectors stores a *pointer* to
    -- the value without evaluating it, so unsupported-encoding error thunks
    -- would be silently swallowed into the column data and only fire lazily
    -- when user code reads a cell. The '!v' bang pattern forces each value
    -- to WHNF before the write, surfacing decoder errors immediately.
    MVector RealWorld a
mvDat <-
        IO (MVector RealWorld a) -> m (MVector RealWorld a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MVector RealWorld a) -> m (MVector RealWorld a))
-> IO (MVector RealWorld a) -> m (MVector RealWorld a)
forall a b. (a -> b) -> a -> b
$ Int -> a -> IO (MVector (PrimState IO) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MVector (PrimState m) a)
VBM.replicate Int
totalRows (String -> a
forall a. HasCallStack => String -> a
error String
"parquet: null slot accessed")
    IOVector Word8
mvValid <- IO (IOVector Word8) -> m (IOVector Word8)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Int -> IO (MVector (PrimState IO) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
totalRows :: IO (VUM.IOVector Word8))
    (Int
_, Bool
hasNull) <-
        ((Int, Bool)
 -> (Vector a, Vector Int, Vector Int) -> m (Int, Bool))
-> (Int, Bool) -> m (Int, Bool)
PageFold m Vector a
runPages
            ( \(Int
rowOff, Bool
anyNull) (Vector a
vals, Vector Int
defs, Vector Int
_) -> IO (Int, Bool) -> m (Int, Bool)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Int, Bool) -> m (Int, Bool))
-> IO (Int, Bool) -> m (Int, Bool)
forall a b. (a -> b) -> a -> b
$ do
                let nDefs :: Int
nDefs = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
defs
                    go :: Int -> Int -> Bool -> IO Bool
go Int
i Int
j Bool
acc
                        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nDefs = Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
acc
                        | 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
                            let !v :: a
v = Vector a -> Int -> a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector a
vals Int
j
                            MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
VBM.unsafeWrite MVector RealWorld a
MVector (PrimState IO) a
mvDat (Int
rowOff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) a
v
                            MVector (PrimState IO) Word8 -> Int -> Word8 -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite IOVector Word8
MVector (PrimState IO) Word8
mvValid (Int
rowOff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) Word8
1
                            Int -> Int -> Bool -> IO Bool
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bool
acc
                        | Bool
otherwise = Int -> Int -> Bool -> IO Bool
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
j Bool
True
                Bool
newNull <- Int -> Int -> Bool -> IO Bool
go Int
0 Int
0 Bool
False
                (Int, Bool) -> IO (Int, Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
rowOff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nDefs, Bool
anyNull Bool -> Bool -> Bool
|| Bool
newNull)
            )
            (Int
0 :: Int, Bool
False)
    Vector a
dat <- IO (Vector a) -> m (Vector a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Vector a) -> m (Vector a)) -> IO (Vector a) -> m (Vector a)
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) a -> IO (Vector a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
VB.unsafeFreeze MVector RealWorld a
MVector (PrimState IO) a
mvDat
    Maybe Bitmap
maybeBm <-
        if Bool
hasNull
            then do
                Bitmap
validV <- IO Bitmap -> m Bitmap
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Bitmap -> m Bitmap) -> IO Bitmap -> m Bitmap
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) Word8 -> IO Bitmap
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze IOVector Word8
MVector (PrimState IO) Word8
mvValid
                Maybe Bitmap -> m (Maybe Bitmap)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Bitmap -> Bitmap
buildBitmapFromValid Bitmap
validV))
            else Maybe Bitmap -> m (Maybe Bitmap)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Bitmap
forall a. Maybe a
Nothing
    Column -> m Column
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
maybeBm Vector a
dat)

{-# INLINEABLE foldNullableUnboxed #-}
foldNullableUnboxed ::
    forall m a.
    (MonadIO m, Columnable a, VU.Unbox a) =>
    Int ->
    Int ->
    PageFold m VU.Vector a ->
    m Column
foldNullableUnboxed :: forall (m :: * -> *) a.
(MonadIO m, Columnable a, Unbox a) =>
Int -> Int -> PageFold m Vector a -> m Column
foldNullableUnboxed Int
maxDef Int
totalRows PageFold m Vector a
runPages = do
    -- zero-init means null slots silently hold 0, guarded by bitmap.
    MVector RealWorld a
mvDat <- IO (MVector RealWorld a) -> m (MVector RealWorld a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MVector RealWorld a) -> m (MVector RealWorld a))
-> IO (MVector RealWorld a) -> m (MVector RealWorld a)
forall a b. (a -> b) -> a -> b
$ Int -> IO (MVector (PrimState IO) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
totalRows
    IOVector Word8
mvValid <- IO (IOVector Word8) -> m (IOVector Word8)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Int -> IO (MVector (PrimState IO) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
totalRows :: IO (VUM.IOVector Word8))
    (Int
_, Bool
hasNull) <-
        ((Int, Bool)
 -> (Vector a, Vector Int, Vector Int) -> m (Int, Bool))
-> (Int, Bool) -> m (Int, Bool)
PageFold m Vector a
runPages
            ( \(Int
rowOff, Bool
anyNull) (Vector a
vals, Vector Int
defs, Vector Int
_) -> IO (Int, Bool) -> m (Int, Bool)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Int, Bool) -> m (Int, Bool))
-> IO (Int, Bool) -> m (Int, Bool)
forall a b. (a -> b) -> a -> b
$ do
                let !nDefs :: Int
nDefs = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
defs
                    go :: Int -> Int -> Bool -> IO Bool
go !Int
i !Int
j !Bool
acc
                        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nDefs = Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
acc
                        | 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
                            MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld a
MVector (PrimState IO) a
mvDat (Int
rowOff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
vals Int
j)
                            MVector (PrimState IO) Word8 -> Int -> Word8 -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite IOVector Word8
MVector (PrimState IO) Word8
mvValid (Int
rowOff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) Word8
1
                            Int -> Int -> Bool -> IO Bool
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bool
acc
                        | Bool
otherwise = Int -> Int -> Bool -> IO Bool
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
j Bool
True
                !Bool
newNull <- Int -> Int -> Bool -> IO Bool
go Int
0 Int
0 Bool
False
                (Int, Bool) -> IO (Int, Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
rowOff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nDefs, Bool
anyNull Bool -> Bool -> Bool
|| Bool
newNull)
            )
            (Int
0 :: Int, Bool
False)
    Vector a
dat <- IO (Vector a) -> m (Vector a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Vector a) -> m (Vector a)) -> IO (Vector a) -> m (Vector a)
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) a -> IO (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld a
MVector (PrimState IO) a
mvDat
    Maybe Bitmap
maybeBm <-
        if Bool
hasNull
            then do
                Bitmap
validV <- IO Bitmap -> m Bitmap
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Bitmap -> m Bitmap) -> IO Bitmap -> m Bitmap
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) Word8 -> IO Bitmap
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze IOVector Word8
MVector (PrimState IO) Word8
mvValid
                Maybe Bitmap -> m (Maybe Bitmap)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Bitmap -> Bitmap
buildBitmapFromValid Bitmap
validV))
            else Maybe Bitmap -> m (Maybe Bitmap)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Bitmap
forall a. Maybe a
Nothing
    Column -> m Column
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
maybeBm Vector a
dat)

{- | Fold a stream of (values, def-levels, rep-levels) triples into a
repeated (list) 'Column' using Dremel-style level stitching.

The stitching function is selected by @maxRep@:

  * @maxRep == 1@  →  'stitchList'   → @[Maybe [Maybe a]]@
  * @maxRep == 2@  →  'stitchList2'  → @[Maybe [Maybe [Maybe a]]]@
  * @maxRep >= 3@  →  'stitchList3'  → @[Maybe [Maybe [Maybe [Maybe a]]]]@

Threshold formula: @defT_r = maxDef - 2 * (maxRep - r)@.
-}
{-# INLINEABLE foldRepeated #-}
foldRepeated ::
    forall m a.
    ( MonadIO m
    , Columnable a
    , Columnable (Maybe [Maybe a])
    , Columnable (Maybe [Maybe [Maybe a]])
    , Columnable (Maybe [Maybe [Maybe [Maybe a]]])
    ) =>
    Int ->
    Int ->
    PageFold m VB.Vector a ->
    m Column
foldRepeated :: forall (m :: * -> *) a.
(MonadIO m, Columnable a, Columnable (Maybe [Maybe a]),
 Columnable (Maybe [Maybe [Maybe a]]),
 Columnable (Maybe [Maybe [Maybe [Maybe a]]])) =>
Int -> Int -> PageFold m Vector a -> m Column
foldRepeated Int
maxRep Int
maxDef PageFold m Vector a
runPages = do
    [(Vector a, Vector Int, Vector Int)]
chunks <- PageFold m Vector a -> m [(Vector a, Vector Int, Vector Int)]
forall (m :: * -> *) (v :: * -> *) a.
Monad m =>
PageFold m v a -> m [(v a, Vector Int, Vector Int)]
collectPages (acc -> (Vector a, Vector Int, Vector Int) -> m acc)
-> acc -> m acc
PageFold m Vector a
runPages
    let allVals :: Vector a
allVals = [Vector a] -> Vector a
forall a. [Vector a] -> Vector a
VB.concat [Vector a
vs | (Vector a
vs, Vector Int
_, Vector Int
_) <- [(Vector a, Vector Int, Vector Int)]
chunks]
        allDefs :: Vector Int
allDefs = [Vector Int] -> Vector Int
forall a. Unbox a => [Vector a] -> Vector a
VU.concat [Vector Int
ds | (Vector a
_, Vector Int
ds, Vector Int
_) <- [(Vector a, Vector Int, Vector Int)]
chunks]
        allReps :: Vector Int
allReps = [Vector Int] -> Vector Int
forall a. Unbox a => [Vector a] -> Vector a
VU.concat [Vector Int
rs | (Vector a
_, Vector Int
_, Vector Int
rs) <- [(Vector a, Vector Int, Vector Int)]
chunks]
    Column -> m Column
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Column -> m Column) -> Column -> m Column
forall a b. (a -> b) -> a -> b
$ case Int
maxRep of
        Int
2 -> [Maybe [Maybe [Maybe a]]] -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList (Int
-> Int
-> Vector Int
-> Vector Int
-> Vector a
-> [Maybe [Maybe [Maybe a]]]
forall a.
Int
-> Int
-> Vector Int
-> Vector Int
-> Vector a
-> [Maybe [Maybe [Maybe a]]]
stitchList2 (Int
maxDef Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) Int
maxDef Vector Int
allReps Vector Int
allDefs Vector a
allVals)
        Int
3 ->
            [Maybe [Maybe [Maybe [Maybe a]]]] -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList (Int
-> Int
-> Int
-> Vector Int
-> Vector Int
-> Vector a
-> [Maybe [Maybe [Maybe [Maybe a]]]]
forall a.
Int
-> Int
-> Int
-> Vector Int
-> Vector Int
-> Vector a
-> [Maybe [Maybe [Maybe [Maybe a]]]]
stitchList3 (Int
maxDef Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
4) (Int
maxDef Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) Int
maxDef Vector Int
allReps Vector Int
allDefs Vector a
allVals)
        Int
_ -> [Maybe [Maybe a]] -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList (Int -> Vector Int -> Vector Int -> Vector a -> [Maybe [Maybe a]]
forall a.
Int -> Vector Int -> Vector Int -> Vector a -> [Maybe [Maybe a]]
stitchList Int
maxDef Vector Int
allReps Vector Int
allDefs Vector a
allVals)

{-# INLINEABLE foldRepeatedUnboxed #-}
foldRepeatedUnboxed ::
    forall m a.
    ( MonadIO m
    , Columnable a
    , VU.Unbox a
    , Columnable (Maybe [Maybe a])
    , Columnable (Maybe [Maybe [Maybe a]])
    , Columnable (Maybe [Maybe [Maybe [Maybe a]]])
    ) =>
    Int ->
    Int ->
    PageFold m VU.Vector a ->
    m Column
foldRepeatedUnboxed :: forall (m :: * -> *) a.
(MonadIO m, Columnable a, Unbox a, Columnable (Maybe [Maybe a]),
 Columnable (Maybe [Maybe [Maybe a]]),
 Columnable (Maybe [Maybe [Maybe [Maybe a]]])) =>
Int -> Int -> PageFold m Vector a -> m Column
foldRepeatedUnboxed Int
maxRep Int
maxDef PageFold m Vector a
runPages = do
    [(Vector a, Vector Int, Vector Int)]
chunks <- PageFold m Vector a -> m [(Vector a, Vector Int, Vector Int)]
forall (m :: * -> *) (v :: * -> *) a.
Monad m =>
PageFold m v a -> m [(v a, Vector Int, Vector Int)]
collectPages (acc -> (Vector a, Vector Int, Vector Int) -> m acc)
-> acc -> m acc
PageFold m Vector a
runPages
    let allVals :: Vector a
allVals = Vector a -> Vector a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VB.convert (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ [Vector a] -> Vector a
forall a. Unbox a => [Vector a] -> Vector a
VU.concat [Vector a
vs | (Vector a
vs, Vector Int
_, Vector Int
_) <- [(Vector a, Vector Int, Vector Int)]
chunks]
        allDefs :: Vector Int
allDefs = [Vector Int] -> Vector Int
forall a. Unbox a => [Vector a] -> Vector a
VU.concat [Vector Int
ds | (Vector a
_, Vector Int
ds, Vector Int
_) <- [(Vector a, Vector Int, Vector Int)]
chunks]
        allReps :: Vector Int
allReps = [Vector Int] -> Vector Int
forall a. Unbox a => [Vector a] -> Vector a
VU.concat [Vector Int
rs | (Vector a
_, Vector Int
_, Vector Int
rs) <- [(Vector a, Vector Int, Vector Int)]
chunks]
    Column -> m Column
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Column -> m Column) -> Column -> m Column
forall a b. (a -> b) -> a -> b
$ case Int
maxRep of
        Int
2 -> [Maybe [Maybe [Maybe a]]] -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList (Int
-> Int
-> Vector Int
-> Vector Int
-> Vector a
-> [Maybe [Maybe [Maybe a]]]
forall a.
Int
-> Int
-> Vector Int
-> Vector Int
-> Vector a
-> [Maybe [Maybe [Maybe a]]]
stitchList2 (Int
maxDef Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) Int
maxDef Vector Int
allReps Vector Int
allDefs Vector a
allVals)
        Int
3 ->
            [Maybe [Maybe [Maybe [Maybe a]]]] -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList (Int
-> Int
-> Int
-> Vector Int
-> Vector Int
-> Vector a
-> [Maybe [Maybe [Maybe [Maybe a]]]]
forall a.
Int
-> Int
-> Int
-> Vector Int
-> Vector Int
-> Vector a
-> [Maybe [Maybe [Maybe [Maybe a]]]]
stitchList3 (Int
maxDef Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
4) (Int
maxDef Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) Int
maxDef Vector Int
allReps Vector Int
allDefs Vector a
allVals)
        Int
_ -> [Maybe [Maybe a]] -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList (Int -> Vector Int -> Vector Int -> Vector a -> [Maybe [Maybe a]]
forall a.
Int -> Vector Int -> Vector Int -> Vector a -> [Maybe [Maybe a]]
stitchList Int
maxDef Vector Int
allReps Vector Int
allDefs Vector a
allVals)

{- | Collect all of a column's page triples into a list, in page order. Used by
the repeated/list folds, where Dremel level-stitching needs the full
concatenated rep/def/value arrays at once (so streaming gives no benefit here).
-}
collectPages ::
    (Monad m) =>
    PageFold m v a ->
    m [(v a, VU.Vector Int, VU.Vector Int)]
collectPages :: forall (m :: * -> *) (v :: * -> *) a.
Monad m =>
PageFold m v a -> m [(v a, Vector Int, Vector Int)]
collectPages PageFold m v a
runPages = [(v a, Vector Int, Vector Int)] -> [(v a, Vector Int, Vector Int)]
forall a. [a] -> [a]
reverse ([(v a, Vector Int, Vector Int)]
 -> [(v a, Vector Int, Vector Int)])
-> m [(v a, Vector Int, Vector Int)]
-> m [(v a, Vector Int, Vector Int)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([(v a, Vector Int, Vector Int)]
 -> (v a, Vector Int, Vector Int)
 -> m [(v a, Vector Int, Vector Int)])
-> [(v a, Vector Int, Vector Int)]
-> m [(v a, Vector Int, Vector Int)]
PageFold m v a
runPages (\[(v a, Vector Int, Vector Int)]
acc (v a, Vector Int, Vector Int)
triple -> [(v a, Vector Int, Vector Int)]
-> m [(v a, Vector Int, Vector Int)]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((v a, Vector Int, Vector Int)
triple (v a, Vector Int, Vector Int)
-> [(v a, Vector Int, Vector Int)]
-> [(v a, Vector Int, Vector Int)]
forall a. a -> [a] -> [a]
: [(v a, Vector Int, Vector Int)]
acc)) []