packages feed

wireform-arrow-0.1.0.0: src/Arrow/Column/Internal.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}

{- | Raw representation of Arrow columns.

The data constructors here establish nothing: every invariant listed
on 'ColumnArray' must hold for each value built with them, because
the accessors in "Arrow.Column" read buffers without re-checking
(after an O(1) shape check and an index check). Only the IPC reader,
the writer, "Arrow.Column.Builder", the smart constructors in
"Arrow.Column" and malformed-input tests should import this module.
Everyone else uses the matching-only pattern synonyms, the @mk*@
constructors and the builders that "Arrow.Column" exports.

The low-level buffer types ('Bitmap', 'Validity' and the element
types) are re-exported with their constructors.
-}
module Arrow.Column.Internal (
  -- * Columns
  ColumnArray (..),
  columnBuffers,

  -- * Fixed-width element tags
  PrimType (..),
  withPrim,
  primWidth,
  swapPattern,
  primTypeName,
  samePrimTag,
  samePrimType,
  SomePrimType (..),
  primTypeFor,
  IntegralPrim (..),
  integralPrim,

  -- * Structural validators (shared by the reader and the smart constructors)
  Offset (..),
  validateOffsets,
  validateUtf8,
  validateKeys,
  validateRunEnds,
  validateListView,
  validateDenseUnion,
  validateSparseUnionTypes,
  validateViews,
  withValidityPtr,

  -- * Offsets
  ChildRange (..),
  rebaseToZero,

  -- * Buffers
  module Arrow.Column.Buffer,
) where

import Arrow.Column.Buffer
import Arrow.Types (ArrowType (..), DateUnit (..), IntervalUnit (..), Precision (..), TimeUnit (..))
import Arrow.Vector.Internal (FixedWidth)
import Columnar.SIMD qualified as K
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Unsafe qualified as BSU
import Data.Int (Int16, Int32, Int64, Int8)
import Data.Text (Text)
import Data.Type.Equality ((:~:) (..))
import Data.Vector qualified as V
import Data.Vector.Storable qualified as VS
import Data.Word (Word16, Word32, Word64, Word8)
import Foreign.Marshal.Array (allocaArray, pokeArray)
import Foreign.Ptr (Ptr, castPtr, nullPtr)
import Foreign.Storable (Storable (..))
import System.IO.Unsafe (unsafeDupablePerformIO)


{- | Element tag of a fixed-width column. Several tags share a Haskell
element type (dates, times, timestamps and durations are 'Int32' or
'Int64'), so the tag, not the type, says what the column is.
Decimal tags carry precision and scale.
-}
data PrimType a where
  PInt8 :: PrimType Int8
  PInt16 :: PrimType Int16
  PInt32 :: PrimType Int32
  PInt64 :: PrimType Int64
  PUInt8 :: PrimType Word8
  PUInt16 :: PrimType Word16
  PUInt32 :: PrimType Word32
  PUInt64 :: PrimType Word64
  PFloat16 :: PrimType Float16
  PFloat :: PrimType Float
  PDouble :: PrimType Double
  PDate32 :: PrimType Int32
  PDate64 :: PrimType Int64
  PTime32 :: PrimType Int32
  PTime64 :: PrimType Int64
  PTimestamp :: PrimType Int64
  PDuration :: PrimType Int64
  PIntervalYearMonth :: PrimType Int32
  PIntervalDayTime :: PrimType IntervalDayTime
  PIntervalMonthDayNano :: PrimType IntervalMonthDayNano
  PDecimal32 :: {-# UNPACK #-} !Int -> {-# UNPACK #-} !Int -> PrimType Int32
  PDecimal64 :: {-# UNPACK #-} !Int -> {-# UNPACK #-} !Int -> PrimType Int64
  PDecimal128 :: {-# UNPACK #-} !Int -> {-# UNPACK #-} !Int -> PrimType Decimal128
  PDecimal256 :: {-# UNPACK #-} !Int -> {-# UNPACK #-} !Int -> PrimType Decimal256


deriving stock instance Show (PrimType a)


deriving stock instance Eq (PrimType a)


{- | Bring the element type's instances into scope from the tag
('FixedWidth' includes 'Storable').
-}
withPrim :: PrimType a -> ((FixedWidth a, Eq a, Show a) => r) -> r
withPrim t k = case t of
  PInt8 -> k
  PInt16 -> k
  PInt32 -> k
  PInt64 -> k
  PUInt8 -> k
  PUInt16 -> k
  PUInt32 -> k
  PUInt64 -> k
  PFloat16 -> k
  PFloat -> k
  PDouble -> k
  PDate32 -> k
  PDate64 -> k
  PTime32 -> k
  PTime64 -> k
  PTimestamp -> k
  PDuration -> k
  PIntervalYearMonth -> k
  PIntervalDayTime -> k
  PIntervalMonthDayNano -> k
  PDecimal32 _ _ -> k
  PDecimal64 _ _ -> k
  PDecimal128 _ _ -> k
  PDecimal256 _ _ -> k
{-# INLINE withPrim #-}


-- | Element width in bytes.
primWidth :: PrimType a -> Int
primWidth = \case
  PInt8 -> 1
  PInt16 -> 2
  PInt32 -> 4
  PInt64 -> 8
  PUInt8 -> 1
  PUInt16 -> 2
  PUInt32 -> 4
  PUInt64 -> 8
  PFloat16 -> 2
  PFloat -> 4
  PDouble -> 8
  PDate32 -> 4
  PDate64 -> 8
  PTime32 -> 4
  PTime64 -> 8
  PTimestamp -> 8
  PDuration -> 8
  PIntervalYearMonth -> 4
  PIntervalDayTime -> 8
  PIntervalMonthDayNano -> 16
  PDecimal32 _ _ -> 4
  PDecimal64 _ _ -> 8
  PDecimal128 _ _ -> 16
  PDecimal256 _ _ -> 32
{-# INLINE primWidth #-}


-- | The tag's name without the @P@ (the column constructor is @Col@ ++ this).
primTypeName :: PrimType a -> String
primTypeName = \case
  PInt8 -> "Int8"
  PInt16 -> "Int16"
  PInt32 -> "Int32"
  PInt64 -> "Int64"
  PUInt8 -> "UInt8"
  PUInt16 -> "UInt16"
  PUInt32 -> "UInt32"
  PUInt64 -> "UInt64"
  PFloat16 -> "Float16"
  PFloat -> "Float"
  PDouble -> "Double"
  PDate32 -> "Date32"
  PDate64 -> "Date64"
  PTime32 -> "Time32"
  PTime64 -> "Time64"
  PTimestamp -> "Timestamp"
  PDuration -> "Duration"
  PIntervalYearMonth -> "IntervalYearMonth"
  PIntervalDayTime -> "IntervalDayTime"
  PIntervalMonthDayNano -> "IntervalMonthDayNano"
  PDecimal32 _ _ -> "Decimal32"
  PDecimal64 _ _ -> "Decimal64"
  PDecimal128 _ _ -> "Decimal128"
  PDecimal256 _ _ -> "Decimal256"


{- | Same tag (decimal precision and scale ignored), with the element
type equality as evidence.
-}
samePrimTag :: PrimType a -> PrimType b -> Maybe (a :~: b)
samePrimTag a b = case (a, b) of
  (PInt8, PInt8) -> Just Refl
  (PInt16, PInt16) -> Just Refl
  (PInt32, PInt32) -> Just Refl
  (PInt64, PInt64) -> Just Refl
  (PUInt8, PUInt8) -> Just Refl
  (PUInt16, PUInt16) -> Just Refl
  (PUInt32, PUInt32) -> Just Refl
  (PUInt64, PUInt64) -> Just Refl
  (PFloat16, PFloat16) -> Just Refl
  (PFloat, PFloat) -> Just Refl
  (PDouble, PDouble) -> Just Refl
  (PDate32, PDate32) -> Just Refl
  (PDate64, PDate64) -> Just Refl
  (PTime32, PTime32) -> Just Refl
  (PTime64, PTime64) -> Just Refl
  (PTimestamp, PTimestamp) -> Just Refl
  (PDuration, PDuration) -> Just Refl
  (PIntervalYearMonth, PIntervalYearMonth) -> Just Refl
  (PIntervalDayTime, PIntervalDayTime) -> Just Refl
  (PIntervalMonthDayNano, PIntervalMonthDayNano) -> Just Refl
  (PDecimal32 _ _, PDecimal32 _ _) -> Just Refl
  (PDecimal64 _ _, PDecimal64 _ _) -> Just Refl
  (PDecimal128 _ _, PDecimal128 _ _) -> Just Refl
  (PDecimal256 _ _, PDecimal256 _ _) -> Just Refl
  _ -> Nothing


-- | Same tag including decimal precision and scale.
samePrimType :: PrimType a -> PrimType b -> Bool
samePrimType a b = case (a, b) of
  (PDecimal32 p s, PDecimal32 p' s') -> p == p' && s == s'
  (PDecimal64 p s, PDecimal64 p' s') -> p == p' && s == s'
  (PDecimal128 p s, PDecimal128 p' s') -> p == p' && s == s'
  (PDecimal256 p s, PDecimal256 p' s') -> p == p' && s == s'
  _ -> case samePrimTag a b of
    Just Refl -> True
    Nothing -> False


data SomePrimType = forall a. SomePrimType !(PrimType a)


instance Show SomePrimType where
  showsPrec d (SomePrimType t) = showParen (d > 10) $ showString "SomePrimType " . showsPrec 11 t


{- | The fixed-width tag of an Arrow type, if it has one. Integer
widths other than 8, 16, 32 and 64 have none.
-}
primTypeFor :: ArrowType -> Maybe SomePrimType
primTypeFor = \case
  AInt 8 True -> some PInt8
  AInt 16 True -> some PInt16
  AInt 32 True -> some PInt32
  AInt 64 True -> some PInt64
  AInt 8 False -> some PUInt8
  AInt 16 False -> some PUInt16
  AInt 32 False -> some PUInt32
  AInt 64 False -> some PUInt64
  AFloatingPoint Half -> some PFloat16
  AFloatingPoint Single -> some PFloat
  AFloatingPoint DoublePrecision -> some PDouble
  ADate DateDay -> some PDate32
  ADate DateMillisecond -> some PDate64
  ATime u _
    | u == Second || u == Millisecond -> some PTime32
    | otherwise -> some PTime64
  ATimestamp _ _ -> some PTimestamp
  ADuration _ -> some PDuration
  AInterval YearMonth -> some PIntervalYearMonth
  AInterval DayTime -> some PIntervalDayTime
  AInterval MonthDayNano -> some PIntervalMonthDayNano
  ADecimal32 p s -> some (PDecimal32 p s)
  ADecimal64 p s -> some (PDecimal64 p s)
  ADecimal p s -> some (PDecimal128 p s)
  ADecimal256 p s -> some (PDecimal256 p s)
  _ -> Nothing
  where
    some :: PrimType a -> Maybe SomePrimType
    some = Just . SomePrimType


-- | Evidence that a tag is a plain integer (the dictionary key and run-end types).
data IntegralPrim a where
  IntegralPrim :: (Integral a, Bounded a, Storable a, Show a) => IntegralPrim a


-- | 'Just' for the eight integer tags only (not dates, times or intervals).
integralPrim :: PrimType a -> Maybe (IntegralPrim a)
integralPrim = \case
  PInt8 -> Just IntegralPrim
  PInt16 -> Just IntegralPrim
  PInt32 -> Just IntegralPrim
  PInt64 -> Just IntegralPrim
  PUInt8 -> Just IntegralPrim
  PUInt16 -> Just IntegralPrim
  PUInt32 -> Just IntegralPrim
  PUInt64 -> Just IntegralPrim
  _ -> Nothing
{-# INLINE integralPrim #-}


{- | One Arrow array. Buffers are Arrow-native and usually alias the
IPC input (see "Arrow.Column" for retention and 'Arrow.Column.copyColumn').

Invariants (not checked by these constructors):

* Validity, when present, has exactly as many bits as the array has
  rows and a null count > 0 that matches its bits ('Validity').
* 'ColPrim': rows = vector length.
* 'ColBool': rows = bitmap length.
* Var-length ('ColUtf8', 'ColBinary', 'ColLargeUtf8',
  'ColLargeBinary'): offsets length = rows + 1 (at least 1),
  @offs[0] >= 0@, monotonic, @offs[rows] <= length data@. Utf8 data in
  @[offs[0], offs[rows])@ is valid UTF-8 and every offset falls on a
  character boundary. Offsets need not start at 0.
* 'ColFixedSizeBinary' width rows: width >= 0, data length >= width * rows.
* Views: views length = 16 * rows; every valid view references
  in-range bytes of its data buffer with a matching prefix; utf8
  views hold valid UTF-8.
* 'ColStruct' rows: every child has at least rows rows (row @i@ of
  the struct is row @i@ of each child).
* Lists and maps: offsets as for var-length, against the child
  length (both keys and values for maps).
* List views: offsets and sizes have rows elements; every valid row
  has @offset >= 0@, @size >= 0@, @offset + size <= child length@.
* 'ColFixedSizeList' size rows: size >= 0, child length >= rows * size.
* Unions: one child index per row (an index into the children, not the
  schema type id), each below the number of children. Dense offsets
  index into the selected child. Sparse children have at least rows rows.
* 'ColDictionary' id indices values: indices is an integer 'ColPrim'
  (8 integer tags) at its wire width carrying the row validity; every
  valid key is in @[0, length values)@ once resolved (the reader holds
  unresolved dictionaries with an empty placeholder until
  'Arrow.Column.resolveDictionaryColumn').
* 'ColRunEndEncoded' offset length runEnds values: runEnds is a
  non-null 'ColPrim' of tag 'PInt16', 'PInt32' or 'PInt64', strictly
  increasing and positive; logical row @i@ is the value of the first
  run whose end is greater than @offset + i@; the last run end is at
  least @offset + length@; values has at least as many rows as runs.
-}
data ColumnArray
  = ColNull {-# UNPACK #-} !Int
  | forall a. ColPrim !(PrimType a) !(Maybe Validity) !(VS.Vector a)
  | ColBool !(Maybe Validity) !Bitmap
  | ColUtf8 !(Maybe Validity) !(VS.Vector Int32) !ByteString
  | ColBinary !(Maybe Validity) !(VS.Vector Int32) !ByteString
  | ColLargeUtf8 !(Maybe Validity) !(VS.Vector Int64) !ByteString
  | ColLargeBinary !(Maybe Validity) !(VS.Vector Int64) !ByteString
  | -- | Width, rows, validity, data.
    ColFixedSizeBinary {-# UNPACK #-} !Int {-# UNPACK #-} !Int !(Maybe Validity) !ByteString
  | -- | Validity, the 16-byte views, the variadic data buffers.
    ColUtf8View !(Maybe Validity) !ByteString !(V.Vector ByteString)
  | ColBinaryView !(Maybe Validity) !ByteString !(V.Vector ByteString)
  | -- | Rows, validity, named children.
    ColStruct {-# UNPACK #-} !Int !(Maybe Validity) !(V.Vector (Text, ColumnArray))
  | ColList !(Maybe Validity) !(VS.Vector Int32) !ColumnArray
  | ColLargeList !(Maybe Validity) !(VS.Vector Int64) !ColumnArray
  | -- | Validity, offsets, sizes, child.
    ColListView !(Maybe Validity) !(VS.Vector Int32) !(VS.Vector Int32) !ColumnArray
  | ColLargeListView !(Maybe Validity) !(VS.Vector Int64) !(VS.Vector Int64) !ColumnArray
  | -- | List size, rows, validity, child.
    ColFixedSizeList {-# UNPACK #-} !Int {-# UNPACK #-} !Int !(Maybe Validity) !ColumnArray
  | -- | Validity, offsets, keys, values.
    ColMap !(Maybe Validity) !(VS.Vector Int32) !ColumnArray !ColumnArray
  | -- | Child index per row, offset per row, children.
    ColDenseUnion !(VS.Vector Int8) !(VS.Vector Int32) !(V.Vector ColumnArray)
  | -- | Child index per row, children.
    ColSparseUnion !(VS.Vector Int8) !(V.Vector ColumnArray)
  | -- | Dictionary id, indices (integer 'ColPrim' with the row validity), values.
    ColDictionary {-# UNPACK #-} !Int64 !ColumnArray !ColumnArray
  | -- | Logical offset, logical length, run ends, values.
    ColRunEndEncoded {-# UNPACK #-} !Int {-# UNPACK #-} !Int !ColumnArray !ColumnArray


{- | Every byte region a column references (buffers, children,
dictionary values), for retention checks and tests. Regions are the
whole underlying buffers, not trimmed to the logical range.
-}
columnBuffers :: ColumnArray -> [ByteString]
columnBuffers = \case
  ColNull _ -> []
  ColPrim t v xs -> withPrim t (vb v ++ [storableToBytes xs])
  ColBool v b -> vb v ++ [bitmapBytes b]
  ColUtf8 v o d -> vb v ++ [storableToBytes o, d]
  ColBinary v o d -> vb v ++ [storableToBytes o, d]
  ColLargeUtf8 v o d -> vb v ++ [storableToBytes o, d]
  ColLargeBinary v o d -> vb v ++ [storableToBytes o, d]
  ColFixedSizeBinary _ _ v d -> vb v ++ [d]
  ColUtf8View v views bufs -> vb v ++ views : V.toList bufs
  ColBinaryView v views bufs -> vb v ++ views : V.toList bufs
  ColStruct _ v cs -> vb v ++ concatMap (columnBuffers . snd) (V.toList cs)
  ColList v o c -> vb v ++ storableToBytes o : columnBuffers c
  ColLargeList v o c -> vb v ++ storableToBytes o : columnBuffers c
  ColListView v o s c -> vb v ++ storableToBytes o : storableToBytes s : columnBuffers c
  ColLargeListView v o s c -> vb v ++ storableToBytes o : storableToBytes s : columnBuffers c
  ColFixedSizeList _ _ v c -> vb v ++ columnBuffers c
  ColMap v o k x -> vb v ++ storableToBytes o : columnBuffers k ++ columnBuffers x
  ColDenseUnion t o cs -> storableToBytes t : storableToBytes o : concatMap columnBuffers (V.toList cs)
  ColSparseUnion t cs -> storableToBytes t : concatMap columnBuffers (V.toList cs)
  ColDictionary _ ix vals -> columnBuffers ix ++ columnBuffers vals
  ColRunEndEncoded _ _ re vals -> columnBuffers re ++ columnBuffers vals
  where
    vb = maybe [] (\x -> [bitmapBytes (validityBits x)])


-- | Byte-swap groups within each element (big-endian bodies).
swapPattern :: PrimType a -> [Int]
swapPattern = \case
  PInt8 -> []
  PUInt8 -> []
  PIntervalDayTime -> [4, 4]
  PIntervalMonthDayNano -> [4, 4, 8]
  t -> [primWidth t]


-- ============================================================
-- Structural validators
-- ============================================================

-- | Offset widths of var-length and list arrays, with their C kernels.
class (Storable o, Integral o, Bounded o, Show o) => Offset o where
  cOffsetsCheck :: Ptr o -> Int -> Int64 -> IO Int
  cUtf8Boundaries :: Ptr o -> Int -> Ptr Word8 -> Int -> IO Int
  cRebaseOffsets :: Ptr o -> Ptr o -> Int -> Int64 -> IO ()
  cTakeOffsets :: Ptr o -> Ptr o -> Ptr Int -> Int -> IO Int
  cTakeBytes :: Ptr Word8 -> Ptr o -> Ptr Word8 -> Ptr Int -> Int -> IO ()
  cListViewCheck :: Ptr o -> Ptr o -> Int -> Ptr Word8 -> Int -> Int64 -> IO Int


instance Offset Int32 where
  cOffsetsCheck = K.offsetsCheckI32
  cUtf8Boundaries = K.utf8BoundariesI32
  cRebaseOffsets = K.rebaseOffsetsI32
  cTakeOffsets = K.takeOffsetsI32
  cTakeBytes = K.takeBytesI32
  cListViewCheck = K.listViewCheckI32
  {-# INLINE cOffsetsCheck #-}
  {-# INLINE cUtf8Boundaries #-}
  {-# INLINE cRebaseOffsets #-}
  {-# INLINE cTakeOffsets #-}
  {-# INLINE cTakeBytes #-}
  {-# INLINE cListViewCheck #-}


instance Offset Int64 where
  cOffsetsCheck = K.offsetsCheckI64
  cUtf8Boundaries = K.utf8BoundariesI64
  cRebaseOffsets = K.rebaseOffsetsI64
  cTakeOffsets = K.takeOffsetsI64
  cTakeBytes = K.takeBytesI64
  cListViewCheck = K.listViewCheckI64
  {-# INLINE cOffsetsCheck #-}
  {-# INLINE cUtf8Boundaries #-}
  {-# INLINE cRebaseOffsets #-}
  {-# INLINE cTakeOffsets #-}
  {-# INLINE cTakeBytes #-}
  {-# INLINE cListViewCheck #-}


kernelResult :: String -> String -> Int -> Either String ()
kernelResult what msg r
  | r < 0 = Right ()
  | otherwise = Left (what ++ ": " ++ msg ++ " at index " ++ show r)
{-# INLINE kernelResult #-}


{- | Offsets of a var-length or list array with @limit@ addressable
elements: at least one offset, @offs[0] >= 0@, monotonic,
@offs[last] <= limit@. One C pass.
-}
validateOffsets :: Offset o => String -> Int -> VS.Vector o -> Either String ()
validateOffsets what limit offs
  | VS.null offs = Left (what ++ ": offsets buffer is empty (an array of n rows needs n + 1 offsets)")
  | otherwise =
      kernelResult what "invalid offset" $
        unsafeDupablePerformIO $
          VS.unsafeWith offs $
            \p -> cOffsetsCheck p (VS.length offs) (fromIntegral limit)


{- | UTF-8 data of a utf8 array whose offsets already passed
'validateOffsets' against the data length: the referenced range is
valid UTF-8 and every offset is on a character boundary.
-}
validateUtf8 :: Offset o => String -> VS.Vector o -> ByteString -> Either String ()
validateUtf8 what offs dat
  | VS.null offs = Right ()
  | not (BS.isValidUtf8 (BSU.unsafeTake (e - s) (BSU.unsafeDrop s dat))) = Left (what ++ ": data is not valid UTF-8")
  | otherwise =
      kernelResult what "offset inside a UTF-8 character" $
        unsafeDupablePerformIO $
          VS.unsafeWith offs $ \po -> withBytesPtr dat $ \pd ->
            cUtf8Boundaries po (VS.length offs) pd (BS.length dat)
  where
    !s = fromIntegral (VS.unsafeHead offs)
    !e = fromIntegral (VS.unsafeLast offs)


-- | Run an action on a validity's byte pointer and bit offset (null pointer when all valid).
withValidityPtr :: Maybe Validity -> (Ptr Word8 -> Int -> IO r) -> IO r
withValidityPtr Nothing k = k nullPtr 0
withValidityPtr (Just v) k =
  let !b = validityBits v
  in withBytesPtr (bitmapBytes b) $ \p -> k p (bitmapOffset b)
{-# INLINE withValidityPtr #-}


{- | Dictionary keys: an integer 'ColPrim' whose valid keys are all in
@[0, dictLength)@. Null keys are not checked.
-}
validateKeys :: String -> ColumnArray -> Int -> Either String ()
validateKeys what col dictLen = case col of
  ColPrim t v xs -> case t of
    PInt8 -> go K.keysInRangeI8 v xs
    PInt16 -> go K.keysInRangeI16 v xs
    PInt32 -> go K.keysInRangeI32 v xs
    PInt64 -> go K.keysInRangeI64 v xs
    PUInt8 -> go K.keysInRangeU8 v xs
    PUInt16 -> go K.keysInRangeU16 v xs
    PUInt32 -> go K.keysInRangeU32 v xs
    PUInt64 -> go K.keysInRangeU64 v xs
    _ -> notInt
  _ -> notInt
  where
    notInt = Left (what ++ ": dictionary indices must be an integer column")
    go :: Storable a => (Ptr a -> Int -> Ptr Word8 -> Int -> Int64 -> IO Int) -> Maybe Validity -> VS.Vector a -> Either String ()
    go k v xs =
      kernelResult what ("dictionary key outside a dictionary of " ++ show dictLen ++ " values") $
        unsafeDupablePerformIO $
          VS.unsafeWith xs $
            \p -> withValidityPtr v $ \pv vo -> k p (VS.length xs) pv vo (fromIntegral dictLen)


{- | Run ends: a non-null 'PInt16' / 'PInt32' / 'PInt64' column,
positive and strictly increasing, the last at least @minEnd@.
-}
validateRunEnds :: String -> ColumnArray -> Int -> Either String ()
validateRunEnds what col minEnd = case col of
  ColPrim t Nothing xs -> case t of
    PInt16 -> go K.runEndsCheckI16 xs
    PInt32 -> go K.runEndsCheckI32 xs
    PInt64 -> go K.runEndsCheckI64 xs
    _ -> bad
  ColPrim _ (Just _) _ -> Left (what ++ ": run ends must not contain nulls")
  _ -> bad
  where
    bad = Left (what ++ ": run ends must be an int16, int32 or int64 column")
    go :: Storable a => (Ptr a -> Int -> Int64 -> IO Int) -> VS.Vector a -> Either String ()
    go k xs =
      kernelResult what ("run ends must be positive, strictly increasing and reach " ++ show minEnd) $
        unsafeDupablePerformIO $
          VS.unsafeWith xs $
            \p -> k p (VS.length xs) (fromIntegral minEnd)


-- | List-view offsets and sizes: per valid row, in range of the child.
validateListView :: Offset o => String -> Maybe Validity -> VS.Vector o -> VS.Vector o -> Int -> Either String ()
validateListView what v offs sizes childLen
  | VS.length offs /= VS.length sizes = Left (what ++ ": offsets and sizes differ in length")
  | otherwise =
      kernelResult what "list view outside its child" $
        unsafeDupablePerformIO $
          VS.unsafeWith offs $ \po -> VS.unsafeWith sizes $ \ps -> withValidityPtr v $ \pv vo ->
            cListViewCheck po ps (VS.length offs) pv vo (fromIntegral childLen)


-- | Dense union child indices and offsets against the child lengths.
validateDenseUnion :: String -> VS.Vector Int8 -> VS.Vector Int32 -> VS.Vector Int64 -> Either String ()
validateDenseUnion what types offs childLens
  | VS.length types /= VS.length offs = Left (what ++ ": type ids and offsets differ in length")
  | otherwise =
      kernelResult what "union child index or offset out of range" $
        unsafeDupablePerformIO $
          VS.unsafeWith types $ \pt -> VS.unsafeWith offs $ \po -> VS.unsafeWith childLens $ \pl ->
            K.denseUnionCheck pt po (VS.length types) pl (VS.length childLens)


-- | Sparse union child indices below the number of children.
validateSparseUnionTypes :: String -> VS.Vector Int8 -> Int -> Either String ()
validateSparseUnionTypes what types nChildren =
  kernelResult what "union child index out of range" $
    unsafeDupablePerformIO $
      VS.unsafeWith types $
        \pt -> K.keysInRangeI8 pt (VS.length types) nullPtr 0 (fromIntegral nChildren)


{- | Binary or utf8 views: 16 bytes per row, every valid view
references in-range bytes with a matching prefix; utf8 views hold
valid UTF-8.
-}
validateViews :: String -> Bool -> Int -> Maybe Validity -> ByteString -> V.Vector ByteString -> Either String ()
validateViews what utf8 rows v views bufs
  | BS.length views < rows * 16 = Left (what ++ ": views buffer too small for " ++ show rows ++ " rows")
  | otherwise =
      kernelResult what "invalid view" $
        unsafeDupablePerformIO $
          withBytesPtr views $ \pviews -> withValidityPtr v $ \pv vo ->
            withAll (V.toList bufs) [] $ \ptrs ->
              allocaArray nb $ \pptrs -> allocaArray nb $ \plens -> do
                pokeArray pptrs (reverse ptrs)
                pokeArray plens (map (fromIntegral . BS.length) (V.toList bufs) :: [Int64])
                K.viewRefsCheck pviews rows pv vo pptrs plens nb utf8
  where
    !nb = V.length bufs
    withAll :: [ByteString] -> [Ptr Word8] -> ([Ptr Word8] -> IO r) -> IO r
    withAll [] acc k = k acc
    withAll (b : rest) acc k = withBytesPtr b $ \p -> withAll rest (castPtr p : acc) k


-- ============================================================
-- Offset helpers
-- ============================================================

-- | A row's range in a list-like column's child: start and length.
data ChildRange = ChildRange
  { childStart :: {-# UNPACK #-} !Int
  , childLength :: {-# UNPACK #-} !Int
  }
  deriving stock (Show, Eq)


-- | Offsets shifted to start at 0 (fresh memory), and the child range they covered.
rebaseToZero :: forall o. Offset o => VS.Vector o -> (VS.Vector o, ChildRange)
rebaseToZero o
  | VS.null o = (VS.singleton 0, ChildRange 0 0)
  | otherwise =
      let !o0 = VS.unsafeHead o
          !n = VS.length o
          bytes = createAligned (n * sizeOf o0) $ \dst ->
            VS.unsafeWith o $ \src -> cRebaseOffsets (castPtr dst) src n (negate (fromIntegral o0))
      in (unsafeBytesToStorable bytes, ChildRange (fromIntegral o0) (fromIntegral (VS.unsafeLast o - o0)))