packages feed

wireform-arrow (empty) → 0.1.0.0

raw patch · 49 files changed

+23256/−0 lines, 49 filesdep +basedep +bytestringdep +containersbinary-added

Dependencies added: base, bytestring, containers, criterion, crypton, deepseq, directory, filepath, hedgehog, lz4-hs, primitive, process, ram, sydtest, template-haskell, temporary, text, vector, wireform-arrow, wireform-columnar-core, wireform-core, wireform-derive, wireform-flatbuffers, zstd

Files

+ LICENSE view
@@ -0,0 +1,26 @@+Copyright 2026 Ian Duncan++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++1.  Redistributions of source code must retain the above copyright notice,+    this list of conditions and the following disclaimer.++2.  Redistributions in binary form must reproduce the above copyright notice,+    this list of conditions and the following disclaimer in the documentation+    and/or other materials provided with the distribution.++3.  Neither the name of the copyright holder nor the names of its contributors+    may be used to endorse or promote products derived from this software+    without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ bench/Bench.hs view
@@ -0,0 +1,450 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TemplateHaskell #-}++{- | Criterion harness for wireform-arrow IPC encode and decode.++Three summaries consume these reports (see @scripts/bench-manifest.json@):++* @arrow-encode-decode@: per workload at 100k rows through "Arrow.Stream":+  @encode/<workload>@ (one allocation for the whole stream),+  @encode lazy/<workload>@ (lazy chunks aliasing the column buffers),+  @decode/<workload>@ (zero-copy columns aliasing the input) and+  @decode + toVector/<workload>@ (decode, convert every column with the+  @to*Vector@ conversions to "Arrow.Vector" vectors, then 'VG.force'+  each into fresh buffers: the result owns its data and keeps nothing+  of the input alive, the same work as arrow-rs's @to_owned_column@+  copying into @Vec<Option<T>>@, @String@s and nested @Vec@s; the+  conversions alone alias the decoded buffers and cost O(rows) or less).+* @arrow-encode-decode-small@: the same four groups as a 100-row batch,+  suffixed @ (100 rows)@.+* @arrow-api-paths@: the mixed 6-column table at 100k rows through+  each public entry point (top-level benches named after the cell).++Every input (columns, encoded bytes, records) is built in 'env' with the+column builders, so only the codec call is timed. Results are forced with+'nf'. Run with a large nursery (the stanza bakes in @-A64m@): the decoders+allocate little, but the materializing rows and the typed paths are+dominated by minor GCs at the default 4 MB nursery.+-}+module Main (main) where++import Arrow.Column+import Arrow.File qualified as AF+import Arrow.Read.Columns (defaultReadOptions)+import Arrow.Record (Table, decodeTable, encodeTable)+import Arrow.Record.TH (deriveTable)+import Arrow.Stream (+  decodeArrowFile,+  decodeArrowStream,+  defaultWriteOptions,+  encodeArrowFile,+  encodeArrowFileLazy,+  encodeArrowStream,+  encodeArrowStreamLazy,+ )+import Arrow.Types (+  ArrowType (..),+  DictionaryEncoding (..),+  Field (..),+  Precision (..),+  Schema,+  defaultField,+  defaultLeafField,+  defaultSchema,+ )+import Arrow.Vector qualified as AV+import Arrow.Write qualified as AW+import Control.DeepSeq (NFData)+import Control.Monad.ST (ST, runST)+import Criterion.Main+import Data.ByteString (ByteString)+import Data.ByteString.Lazy qualified as BL+import Data.Int (Int32, Int64)+import Data.Text (Text)+import Data.Vector qualified as V+import Data.Vector.Generic qualified as VG+import Data.Vector.Storable qualified as VS+import Foreign.Storable (Storable)+import GHC.Generics (Generic)+++-- | One benchmark input: a schema plus a single record batch.+data Workload = Workload+  { wlName :: !String+  , wlSchema :: !Schema+  , wlBatch :: !(V.Vector ColumnArray)+  }+  deriving stock (Generic)+  deriving anyclass (NFData)+++bigRows :: Int+bigRows = 100000+++smallRows :: Int+smallRows = 100+++-- | A short label per row, cycling through a fixed pool so string+-- lengths stay realistic (5 to 12 bytes) without per-row 'show'.+labelPool :: V.Vector Text+labelPool =+  V.fromList+    [ "alpha", "bravo", "charlie", "delta", "echo", "foxtrot"+    , "golf", "hotel", "india", "juliett", "kilo", "lima"+    , "mike", "november", "oscar", "papa-quebec"+    ]+++label :: Int -> Text+label i = V.unsafeIndex labelPool (i `mod` V.length labelPool)+++-- | Every tenth row is null (roughly 10% nulls).+nullRow :: Int -> Bool+nullRow i = i `mod` 10 == 0+++loopN :: Int -> (Int -> ST s ()) -> ST s ()+loopN n f = go 0+ where+  go !i+    | i >= n = pure ()+    | otherwise = f i >> go (i + 1)+++primCol :: Storable a => PrimType a -> Int -> (Int -> Maybe a) -> ColumnArray+primCol ty n f = runST $ do+  b <- newPrimBuilder ty n+  loopN n (appendPrimMaybe b . f)+  freezeBuilder b+++textCol :: Int -> (Int -> Maybe Text) -> ColumnArray+textCol n f = runST $ do+  b <- newUtf8Builder n+  loopN n (appendTextMaybe b . f)+  freezeBuilder b+++boolCol :: Int -> ColumnArray+boolCol n = runST $ do+  b <- newBoolBuilder n+  loopN n (appendBool b . even)+  freezeBuilder b+++int64Col, doubleCol, nullableInt64Col, utf8Col, nullableUtf8Col :: Int -> ColumnArray+int64Col n = primCol PInt64 n (\i -> Just (fromIntegral i * 7919))+doubleCol n = primCol PDouble n (\i -> Just (fromIntegral i * 1.25))+nullableInt64Col n = primCol PInt64 n (\i -> if nullRow i then Nothing else Just (fromIntegral i))+utf8Col n = textCol n (Just . label)+nullableUtf8Col n = textCol n (\i -> if nullRow i then Nothing else Just (label i))+++-- | Setup-time construction failure is a harness bug; fail loudly.+built :: String -> Either String a -> a+built what = either (\e -> error (what <> ": " <> e)) id+++i64, f64, i32, utf8 :: ArrowType+i64 = AInt 64 True+f64 = AFloatingPoint DoublePrecision+i32 = AInt 32 True+utf8 = AUtf8+++single :: String -> Field -> ColumnArray -> Workload+single name f c = Workload name (defaultSchema (V.singleton f)) (V.singleton c)+++workloads :: Int -> [Workload]+workloads n =+  [ single "int64" (defaultLeafField "v" False i64) (int64Col n)+  , single "double" (defaultLeafField "v" False f64) (doubleCol n)+  , single "nullable int64" (defaultLeafField "v" True i64) (nullableInt64Col n)+  , single "utf8" (defaultLeafField "v" False utf8) (utf8Col n)+  , single "nullable utf8" (defaultLeafField "v" True utf8) (nullableUtf8Col n)+  , mixed n+  , listInt32 n+  , struct3 n+  , dictUtf8 n+  ]+++-- | Workload names, known without building the data (criterion needs the+-- names outside 'env').+workloadNames :: [String]+workloadNames =+  [ "int64", "double", "nullable int64", "utf8", "nullable utf8"+  , "mixed 6-col", "list<int32>", "struct<int32,double,bool>", "dictionary<utf8>"+  ]+++mixedSchema :: Schema+mixedSchema =+  defaultSchema $+    V.fromList+      [ defaultLeafField "tradeId" False i64+      , defaultLeafField "tradePrice" False f64+      , defaultLeafField "tradeQty" True i64+      , defaultLeafField "tradeSymbol" False utf8+      , defaultLeafField "tradeNote" True utf8+      , defaultLeafField "tradeSettled" False ABool+      ]+++mixed :: Int -> Workload+mixed n =+  Workload+    "mixed 6-col"+    mixedSchema+    ( V.fromList+        [ int64Col n+        , doubleCol n+        , nullableInt64Col n+        , utf8Col n+        , nullableUtf8Col n+        , boolCol n+        ]+    )+++-- | Four int32 elements per row.+listInt32 :: Int -> Workload+listInt32 n =+  single+    "list<int32>"+    (defaultField "v" False AList (V.singleton (defaultLeafField "item" False i32)))+    ( built "list<int32>" $+        mkList+          Nothing+          (VS.generate (n + 1) (\i -> fromIntegral (i * 4)))+          (primCol PInt32 (n * 4) (Just . fromIntegral))+    )+++struct3 :: Int -> Workload+struct3 n =+  single+    "struct<int32,double,bool>"+    ( defaultField "v" False AStruct $+        V.fromList+          [ defaultLeafField "a" False i32+          , defaultLeafField "b" False f64+          , defaultLeafField "c" False ABool+          ]+    )+    ( built "struct" $+        mkStruct n Nothing $+          V.fromList+            [ ("a", primCol PInt32 n (Just . fromIntegral))+            , ("b", doubleCol n)+            , ("c", boolCol n)+            ]+    )+++-- | Sixteen distinct values, int32 keys cycling over them.+dictUtf8 :: Int -> Workload+dictUtf8 n =+  single+    "dictionary<utf8>"+    ( (defaultLeafField "v" False utf8)+        { fieldDictionary = Just (DictionaryEncoding 0 i32 False)+        }+    )+    ( built "dictionary<utf8>" $+        mkDictionary+          0+          (primCol PInt32 n (\i -> Just (fromIntegral (i `mod` V.length labelPool))))+          (fromTexts labelPool)+    )+++{- | A decoded column converted to "Arrow.Vector" vectors and detached+from the input, one constructor per shape the workloads use.+-}+data Owned+  = OInt32 !(AV.Vector (Maybe Int32))+  | OInt64 !(AV.Vector (Maybe Int64))+  | ODouble !(AV.Vector (Maybe Double))+  | OText !(AV.Vector (Maybe Text))+  | OBool !(AV.Vector (Maybe Bool))+  | OListInt32 !(AV.Vector (Maybe (AV.Vector (Maybe Int32))))+  | OStruct !(V.Vector Owned)+  deriving stock (Generic)+  deriving anyclass (NFData)+++{- | The "toVector" half of the decode + toVector rows: convert (which+aliases the decoded buffers), then 'VG.force' into fresh buffers so the+result owns its data, as arrow-rs's @to_owned_column@ does with its+@Vec@s and @String@s.+-}+toOwned :: ColumnArray -> Either String Owned+toOwned c+  | Just p <- asPrim PInt64 c = Right (OInt64 (VG.force (toMaybeVector p)))+  | Just p <- asPrim PInt32 c = Right (OInt32 (VG.force (toMaybeVector p)))+  | Just p <- asPrim PDouble c = Right (ODouble (VG.force (toMaybeVector p)))+  | Just _ <- asUtf8 c = OText . VG.force <$> toTextVector c+  | Just _ <- asBool c = OBool . VG.force <$> toBoolVector c+toOwned c = case c of+  ColList _ _ child+    | Just _ <- asPrim PInt32 child -> OListInt32 . VG.force <$> toListVector int32s c+  -- The workloads build struct children exactly as long as the struct.+  ColStruct _ _ fs -> OStruct <$> traverse (toOwned . snd) fs+  -- String dictionaries convert their values once; the owned rows share+  -- one copy of the values instead of one string per row.+  ColDictionary _ _ vals | Just _ <- asUtf8 vals -> OText . VG.force <$> toTextVector c+  ColDictionary {} -> expandDictionary c >>= toOwned+  _ -> Left ("toOwned: unsupported column " <> columnTag c)+ where+  int32s ch = maybe (Left ("toOwned: list child " <> columnTag ch)) (Right . toMaybeVector) (asPrim PInt32 ch)+++decodeToVector :: ByteString -> Either String [V.Vector Owned]+decodeToVector bs = do+  (_, batches) <- decodeArrowStream defaultReadOptions bs+  traverse (traverse toOwned) batches+++-- | Derived record matching 'mixedSchema' column for column.+data Trade = Trade+  { tradeId :: !Int64+  , tradePrice :: !Double+  , tradeQty :: !(Maybe Int64)+  , tradeSymbol :: !Text+  , tradeNote :: !(Maybe Text)+  , tradeSettled :: !Bool+  }+  deriving stock (Generic)+  deriving anyclass (NFData)++-- Close the declaration group so 'deriveTable' can reify 'Trade'.+$(pure [])+++tradeTable :: Table Trade+tradeTable = $(deriveTable ''Trade)+++trades :: Int -> V.Vector Trade+trades n =+  V.generate n $ \i ->+    Trade+      { tradeId = fromIntegral i * 7919+      , tradePrice = fromIntegral i * 1.25+      , tradeQty = if nullRow i then Nothing else Just (fromIntegral i)+      , tradeSymbol = label i+      , tradeNote = if nullRow i then Nothing else Just (label i)+      , tradeSettled = even i+      }+++typedDecode :: ByteString -> Either String (V.Vector Trade)+typedDecode bs = case decodeArrowStream defaultReadOptions bs of+  Left err -> Left err+  Right (sch, [cols]) -> decodeTable tradeTable sch cols+  Right (_, batches) -> Left ("expected one batch, got " <> show (length batches))+++typedEncode :: V.Vector Trade -> Either String ByteString+typedEncode ts =+  let (sch, cols) = encodeTable tradeTable ts+   in encodeArrowStream defaultWriteOptions sch [cols]+++encodeW :: Workload -> Either String ByteString+encodeW w = encodeArrowStream defaultWriteOptions (wlSchema w) [wlBatch w]+++-- | 'nf' on a lazy 'BL.ByteString' forces the chunk spine, which is the+-- whole cost of the lazy encoder (the chunks alias the column buffers).+encodeLazyW :: Workload -> Either String BL.ByteString+encodeLazyW w = encodeArrowStreamLazy defaultWriteOptions (wlSchema w) [wlBatch w]+++-- | Fail fast at setup time rather than timing an error path.+checked :: String -> Either String a -> IO a+checked what = either (\e -> fail (what <> ": " <> e)) pure+++-- | Build the workload in 'env' and check it encodes.+workloadEnv :: Int -> String -> IO Workload+workloadEnv n name = case filter ((== name) . wlName) (workloads n) of+  [w] -> checked name (encodeW w) >> pure w+  _ -> fail ("unknown workload " <> name)+++-- | Encoded bytes of the workload, checked to decode and materialize.+decodeEnv :: Int -> String -> IO ByteString+decodeEnv n name = do+  w <- workloadEnv n name+  bs <- checked name (encodeW w)+  _ <- checked name (decodeToVector bs)+  pure bs+++codecGroups :: String -> Int -> [Benchmark]+codecGroups suffix n =+  [ bgroup ("encode" <> suffix) (map (encodeBench encodeW) workloadNames)+  , bgroup ("encode lazy" <> suffix) (map (encodeBench encodeLazyW) workloadNames)+  , bgroup ("decode" <> suffix) (map (decodeBench (decodeArrowStream defaultReadOptions)) workloadNames)+  , bgroup ("decode + toVector" <> suffix) (map (decodeBench decodeToVector) workloadNames)+  ]+ where+  encodeBench :: NFData b => (Workload -> b) -> String -> Benchmark+  encodeBench f name = env (workloadEnv n name) $ \w -> bench name $ nf f w+  decodeBench :: NFData b => (ByteString -> b) -> String -> Benchmark+  decodeBench f name = env (decodeEnv n name) $ \bs -> bench name $ nf f bs+++apiPaths :: [Benchmark]+apiPaths =+  [ env mixedEnv $ \w ->+      bench "stream encode (Arrow.Stream)" $ nf encodeW w+  , env mixedEnv $ \w ->+      bench "stream encode lazy (Arrow.Stream)" $ nf encodeLazyW w+  , env mixedEnv $ \w ->+      bench "stream encode (Arrow.Write)" $+        nf (AW.writeArrowStream (wlSchema w)) (V.singleton (wlBatch w))+  , env (checked "stream decode" (encodeW (mixed bigRows))) $ \bs ->+      bench "stream decode (Arrow.Stream)" $ nf (decodeArrowStream defaultReadOptions) bs+  , env mixedEnv $ \w ->+      bench "file encode (Arrow.Stream)" $+        nf (encodeArrowFile defaultWriteOptions (wlSchema w)) [wlBatch w]+  , env mixedEnv $ \w ->+      bench "file encode lazy (Arrow.Stream)" $+        nf (encodeArrowFileLazy defaultWriteOptions (wlSchema w)) [wlBatch w]+  , env specFileBytes $ \bs ->+      bench "file decode (Arrow.Stream)" $ nf (decodeArrowFile defaultReadOptions) bs+  , env specFileBytes $ \bs ->+      bench "file read (Arrow.File)" $+        nf (either (error . ("file read: " <>)) id . AF.readArrowFileColumns defaultReadOptions) bs+  , env (pure (trades bigRows)) $ \ts ->+      bench "typed encode (Arrow.Record)" $ nf typedEncode ts+  , env typedBytes $ \bs ->+      bench "typed decode (Arrow.Record)" $ nf typedDecode bs+  ]+ where+  mixedEnv = workloadEnv bigRows "mixed 6-col"+  specFileBytes = do+    w <- mixedEnv+    bs <- checked "file encode" (encodeArrowFile defaultWriteOptions (wlSchema w) [wlBatch w])+    _ <- checked "file decode" (decodeArrowFile defaultReadOptions bs)+    pure bs+  typedBytes = do+    bs <- checked "typed encode" (typedEncode (trades bigRows))+    _ <- checked "typed decode" (typedDecode bs)+    pure bs+++main :: IO ()+main =+  defaultMain $+    codecGroups "" bigRows+      <> codecGroups (" (" <> show smallRows <> " rows)") smallRows+      <> apiPaths
+ src/Arrow/Column.hs view
@@ -0,0 +1,2651 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE UnboxedTuples #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+-- The Eq, Show and NFData instances of 'ColumnArray' (defined in+-- "Arrow.Column.Internal") live here because they are written with+-- the accessors and row model of this module.+{-# OPTIONS_GHC -Wno-orphans #-}++{- | Arrow columns as Arrow buffers.++A 'ColumnArray' holds Arrow-native buffers: storable vectors for+fixed-width values and offsets, LSB-first bitmaps for validity and+booleans, and 'ByteString' regions for variable-length data. A column+decoded from IPC aliases the input bytes (zero copy) and keeps them+alive; 'copyColumn' detaches a column (for example a small slice of a+large input) by copying only the bytes it references.++Nullability is per value, not per type: every array with a validity+slot carries @Maybe Validity@, 'Nothing' meaning no nulls (a validity+whose null count is 0 is normalised to 'Nothing').++Construction is checked. The patterns exported here only match; build+columns with the builders ("Arrow.Column.Builder", re-exported), the+@from*@ conversions, or the validating @mk*@ constructors. The raw+constructors live in "Arrow.Column.Internal". Every column that user+code can build therefore satisfies the invariants listed there, which+is what lets the @*At@ accessors read without re-validating.++'Eq' is logical and O(n): two columns are equal when they have the+same type, the same length, the same validity per row, and equal+values in valid rows. Null slots, bit offsets, offset bases and+dictionary layouts do not matter; floating point compares bit+patterns. 'Show' prints logical rows.+-}+module Arrow.Column (+  -- * Columns+  ColumnArray,+  pattern ColNull,+  pattern ColPrim,+  pattern ColInt8,+  pattern ColInt16,+  pattern ColInt32,+  pattern ColInt64,+  pattern ColUInt8,+  pattern ColUInt16,+  pattern ColUInt32,+  pattern ColUInt64,+  pattern ColFloat16,+  pattern ColFloat,+  pattern ColDouble,+  pattern ColDate32,+  pattern ColDate64,+  pattern ColTime32,+  pattern ColTime64,+  pattern ColTimestamp,+  pattern ColDuration,+  pattern ColIntervalYearMonth,+  pattern ColIntervalDayTime,+  pattern ColIntervalMonthDayNano,+  pattern ColDecimal32,+  pattern ColDecimal64,+  pattern ColDecimal128,+  pattern ColDecimal256,+  pattern ColBool,+  pattern ColUtf8,+  pattern ColBinary,+  pattern ColLargeUtf8,+  pattern ColLargeBinary,+  pattern ColFixedSizeBinary,+  pattern ColUtf8View,+  pattern ColBinaryView,+  pattern ColStruct,+  pattern ColList,+  pattern ColLargeList,+  pattern ColListView,+  pattern ColLargeListView,+  pattern ColFixedSizeList,+  pattern ColMap,+  pattern ColDenseUnion,+  pattern ColSparseUnion,+  pattern ColDictionary,+  pattern ColRunEndEncoded,++  -- * Element tags+  PrimType (..),+  withPrim,+  primWidth,+  SomePrimType (..),+  primTypeFor,+  Offset,++  -- * Buffers+  Bitmap,+  bitmapBytes,+  bitmapOffset,+  bitmapLength,+  mkBitmap,+  emptyBitmap,+  bitAt,+  bitmapSetCount,+  bitmapGenerate,+  bitmapFromBools,+  bitmapToBools,+  Validity,+  validityBits,+  validityNullCount,+  mkValidity,+  validityGenerate,+  validityFromBools,+  isValidAt,+  Float16 (..),+  float16ToDouble,+  IntervalDayTime (..),+  IntervalMonthDayNano (..),+  Decimal128 (..),+  decimal128ToInteger,+  decimal128FromInteger,+  Decimal256 (..),+  decimal256ToInteger,+  decimal256FromInteger,++  -- * Shape+  columnLength,+  nullCount,+  validity,+  columnTag,+  hasValiditySlot,++  -- * Typed views (O(1); index functions INLINE)+  PrimArray (..),+  asPrim,+  primArrayLength,+  primAt,+  unsafePrimAt,+  primValueAt,+  unsafePrimValueAt,+  BytesArray (..),+  Utf8Array (..),+  asUtf8,+  asLargeUtf8,+  asBinary,+  asLargeBinary,+  bytesArrayLength,+  bytesAt,+  unsafeBytesAt,+  textAt,+  unsafeTextAt,+  BoolArray (..),+  asBool,+  boolArrayAt,+  boolAt,+  anyBytesAt,+  anyTextAt,+  ChildRange (..),+  listRange,+  dictKeyAt,++  -- * Conversions to "Arrow.Vector" (aliasing the column; see each for its cost)+  toStorable,+  toMaybeVector,+  toTextVector,+  toBytesVector,+  toBoolVector,+  toListVector,+  copyColumn,++  -- * Construction+  primColumn,+  primColumnV,+  mkPrim,+  fromMaybes,+  fromMaybeVector,+  fromBoolVector,+  fromBools,+  fromMaybeBools,+  fromTexts,+  fromMaybeTexts,+  fromMaybeLargeTexts,+  fromByteStrings,+  fromMaybeByteStrings,+  fromMaybeLargeByteStrings,+  fromMaybeFixedSizeBinary,+  fromMaybeUtf8View,+  fromMaybeBinaryView,+  mkBool,+  mkUtf8,+  mkBinary,+  mkLargeUtf8,+  mkLargeBinary,+  mkFixedSizeBinary,+  mkUtf8View,+  mkBinaryView,+  mkStruct,+  mkList,+  mkLargeList,+  mkListView,+  mkLargeListView,+  mkFixedSizeList,+  mkMap,+  mkDenseUnion,+  mkSparseUnion,+  mkDictionary,+  mkRunEndEncoded,+  emptyColumnFor,+  placeholderColumn,+  fillerColumn,++  -- * Builders+  module Arrow.Column.Builder,++  -- * Row operations+  sliceColumnArray,+  concatColumnArray,+  concatColumnArrays,+  takeColumnArray,+  rebaseRunEnds,++  -- * Dictionaries+  expandDictionary,+  resolveDictionaryColumn,++  -- * Nullability+  maskValidity,+  toNullableColumn,++  -- * Map invariants+  validateMapKeysSorted,+) where++import Arrow.Column.Builder+import Arrow.Column.Internal (ColumnArray)+import Arrow.Column.Internal hiding (ColumnArray (..))+import Arrow.Column.Internal qualified as I+import Arrow.Vector.Internal (FixedWidth)+import Arrow.Vector.Internal qualified as AV+import Arrow.Types (ArrowType (..), DictionaryEncoding (..), Field (..), UnionMode (..))+import Columnar.SIMD qualified as K+import Control.DeepSeq (NFData (..))+import Control.Monad (forM_, when)+import Control.Monad.ST (stToIO)+import Data.Bits (unsafeShiftR, (.&.))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Internal qualified as BSI+import Data.ByteString.Unsafe qualified as BSU+import Data.Int (Int16, Int32, Int64, Int8)+import Data.List (intersperse)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Text.Array qualified as TA+import Data.Text.Foreign qualified as TF+import Data.Type.Equality ((:~:) (..))+import Data.Vector qualified as V+import Data.Vector.Generic qualified as G+import Data.Vector.Storable qualified as VS+import Data.Vector.Storable.Mutable qualified as VSM+import Data.Word (Word16, Word32, Word64, Word8)+import Foreign.ForeignPtr (castForeignPtr)+import GHC.Exts (Int (I#), MutableByteArray#, RealWorld, byteArrayContents#, minusAddr#, unsafeFreezeByteArray#)+import GHC.ForeignPtr (ForeignPtr (..), ForeignPtrContents (..), unsafeWithForeignPtr)+import GHC.IO (IO (..))+import Foreign.Marshal.Utils (copyBytes, fillBytes)+import Foreign.Ptr (Ptr, castPtr, plusPtr)+import Foreign.Storable (Storable (..))+import System.IO.Unsafe (unsafeDupablePerformIO)+++-- ============================================================+-- Patterns+-- ============================================================++-- | An Arrow NULL column of the given length (no buffers; every row null).+pattern ColNull :: Int -> ColumnArray+pattern ColNull n = I.ColNull n+++-- | Any fixed-width column, with its element tag.+pattern ColPrim :: () => forall a. PrimType a -> Maybe Validity -> VS.Vector a -> ColumnArray+pattern ColPrim t v xs <- I.ColPrim t v xs+++pattern ColInt8 :: Maybe Validity -> VS.Vector Int8 -> ColumnArray+pattern ColInt8 v xs <- I.ColPrim PInt8 v xs+++pattern ColInt16 :: Maybe Validity -> VS.Vector Int16 -> ColumnArray+pattern ColInt16 v xs <- I.ColPrim PInt16 v xs+++pattern ColInt32 :: Maybe Validity -> VS.Vector Int32 -> ColumnArray+pattern ColInt32 v xs <- I.ColPrim PInt32 v xs+++pattern ColInt64 :: Maybe Validity -> VS.Vector Int64 -> ColumnArray+pattern ColInt64 v xs <- I.ColPrim PInt64 v xs+++pattern ColUInt8 :: Maybe Validity -> VS.Vector Word8 -> ColumnArray+pattern ColUInt8 v xs <- I.ColPrim PUInt8 v xs+++pattern ColUInt16 :: Maybe Validity -> VS.Vector Word16 -> ColumnArray+pattern ColUInt16 v xs <- I.ColPrim PUInt16 v xs+++pattern ColUInt32 :: Maybe Validity -> VS.Vector Word32 -> ColumnArray+pattern ColUInt32 v xs <- I.ColPrim PUInt32 v xs+++pattern ColUInt64 :: Maybe Validity -> VS.Vector Word64 -> ColumnArray+pattern ColUInt64 v xs <- I.ColPrim PUInt64 v xs+++pattern ColFloat16 :: Maybe Validity -> VS.Vector Float16 -> ColumnArray+pattern ColFloat16 v xs <- I.ColPrim PFloat16 v xs+++pattern ColFloat :: Maybe Validity -> VS.Vector Float -> ColumnArray+pattern ColFloat v xs <- I.ColPrim PFloat v xs+++pattern ColDouble :: Maybe Validity -> VS.Vector Double -> ColumnArray+pattern ColDouble v xs <- I.ColPrim PDouble v xs+++pattern ColDate32 :: Maybe Validity -> VS.Vector Int32 -> ColumnArray+pattern ColDate32 v xs <- I.ColPrim PDate32 v xs+++pattern ColDate64 :: Maybe Validity -> VS.Vector Int64 -> ColumnArray+pattern ColDate64 v xs <- I.ColPrim PDate64 v xs+++pattern ColTime32 :: Maybe Validity -> VS.Vector Int32 -> ColumnArray+pattern ColTime32 v xs <- I.ColPrim PTime32 v xs+++pattern ColTime64 :: Maybe Validity -> VS.Vector Int64 -> ColumnArray+pattern ColTime64 v xs <- I.ColPrim PTime64 v xs+++pattern ColTimestamp :: Maybe Validity -> VS.Vector Int64 -> ColumnArray+pattern ColTimestamp v xs <- I.ColPrim PTimestamp v xs+++pattern ColDuration :: Maybe Validity -> VS.Vector Int64 -> ColumnArray+pattern ColDuration v xs <- I.ColPrim PDuration v xs+++pattern ColIntervalYearMonth :: Maybe Validity -> VS.Vector Int32 -> ColumnArray+pattern ColIntervalYearMonth v xs <- I.ColPrim PIntervalYearMonth v xs+++pattern ColIntervalDayTime :: Maybe Validity -> VS.Vector IntervalDayTime -> ColumnArray+pattern ColIntervalDayTime v xs <- I.ColPrim PIntervalDayTime v xs+++pattern ColIntervalMonthDayNano :: Maybe Validity -> VS.Vector IntervalMonthDayNano -> ColumnArray+pattern ColIntervalMonthDayNano v xs <- I.ColPrim PIntervalMonthDayNano v xs+++-- | Precision, scale, validity, values.+pattern ColDecimal32 :: Int -> Int -> Maybe Validity -> VS.Vector Int32 -> ColumnArray+pattern ColDecimal32 p s v xs <- I.ColPrim (PDecimal32 p s) v xs+++pattern ColDecimal64 :: Int -> Int -> Maybe Validity -> VS.Vector Int64 -> ColumnArray+pattern ColDecimal64 p s v xs <- I.ColPrim (PDecimal64 p s) v xs+++pattern ColDecimal128 :: Int -> Int -> Maybe Validity -> VS.Vector Decimal128 -> ColumnArray+pattern ColDecimal128 p s v xs <- I.ColPrim (PDecimal128 p s) v xs+++pattern ColDecimal256 :: Int -> Int -> Maybe Validity -> VS.Vector Decimal256 -> ColumnArray+pattern ColDecimal256 p s v xs <- I.ColPrim (PDecimal256 p s) v xs+++-- | Validity, values bitmap.+pattern ColBool :: Maybe Validity -> Bitmap -> ColumnArray+pattern ColBool v b <- I.ColBool v b+++-- | Validity, offsets (rows + 1, need not start at 0), UTF-8 data.+pattern ColUtf8 :: Maybe Validity -> VS.Vector Int32 -> ByteString -> ColumnArray+pattern ColUtf8 v o d <- I.ColUtf8 v o d+++pattern ColBinary :: Maybe Validity -> VS.Vector Int32 -> ByteString -> ColumnArray+pattern ColBinary v o d <- I.ColBinary v o d+++pattern ColLargeUtf8 :: Maybe Validity -> VS.Vector Int64 -> ByteString -> ColumnArray+pattern ColLargeUtf8 v o d <- I.ColLargeUtf8 v o d+++pattern ColLargeBinary :: Maybe Validity -> VS.Vector Int64 -> ByteString -> ColumnArray+pattern ColLargeBinary v o d <- I.ColLargeBinary v o d+++-- | Width, rows, validity, data (row @i@ is bytes @[i * width, (i + 1) * width)@).+pattern ColFixedSizeBinary :: Int -> Int -> Maybe Validity -> ByteString -> ColumnArray+pattern ColFixedSizeBinary w n v d <- I.ColFixedSizeBinary w n v d+++-- | Validity, 16-byte views, variadic data buffers.+pattern ColUtf8View :: Maybe Validity -> ByteString -> V.Vector ByteString -> ColumnArray+pattern ColUtf8View v views bufs <- I.ColUtf8View v views bufs+++pattern ColBinaryView :: Maybe Validity -> ByteString -> V.Vector ByteString -> ColumnArray+pattern ColBinaryView v views bufs <- I.ColBinaryView v views bufs+++-- | Rows, validity, named children (each has at least rows rows).+pattern ColStruct :: Int -> Maybe Validity -> V.Vector (Text, ColumnArray) -> ColumnArray+pattern ColStruct n v cs <- I.ColStruct n v cs+++pattern ColList :: Maybe Validity -> VS.Vector Int32 -> ColumnArray -> ColumnArray+pattern ColList v o c <- I.ColList v o c+++pattern ColLargeList :: Maybe Validity -> VS.Vector Int64 -> ColumnArray -> ColumnArray+pattern ColLargeList v o c <- I.ColLargeList v o c+++-- | Validity, offsets, sizes, child.+pattern ColListView :: Maybe Validity -> VS.Vector Int32 -> VS.Vector Int32 -> ColumnArray -> ColumnArray+pattern ColListView v o s c <- I.ColListView v o s c+++pattern ColLargeListView :: Maybe Validity -> VS.Vector Int64 -> VS.Vector Int64 -> ColumnArray -> ColumnArray+pattern ColLargeListView v o s c <- I.ColLargeListView v o s c+++-- | List size, rows, validity, child.+pattern ColFixedSizeList :: Int -> Int -> Maybe Validity -> ColumnArray -> ColumnArray+pattern ColFixedSizeList w n v c <- I.ColFixedSizeList w n v c+++-- | Validity, offsets, keys, values.+pattern ColMap :: Maybe Validity -> VS.Vector Int32 -> ColumnArray -> ColumnArray -> ColumnArray+pattern ColMap v o k x <- I.ColMap v o k x+++-- | Child index per row (not the schema type id), offset per row, children.+pattern ColDenseUnion :: VS.Vector Int8 -> VS.Vector Int32 -> V.Vector ColumnArray -> ColumnArray+pattern ColDenseUnion t o cs <- I.ColDenseUnion t o cs+++pattern ColSparseUnion :: VS.Vector Int8 -> V.Vector ColumnArray -> ColumnArray+pattern ColSparseUnion t cs <- I.ColSparseUnion t cs+++-- | Dictionary id, indices (an integer column at wire width, carrying the row validity), values.+pattern ColDictionary :: Int64 -> ColumnArray -> ColumnArray -> ColumnArray+pattern ColDictionary did ix vals <- I.ColDictionary did ix vals+++-- | Logical offset, logical length, run ends, values.+pattern ColRunEndEncoded :: Int -> Int -> ColumnArray -> ColumnArray -> ColumnArray+pattern ColRunEndEncoded off len re vals <- I.ColRunEndEncoded off len re vals+++{-# COMPLETE+  ColNull+  , ColPrim+  , ColBool+  , ColUtf8+  , ColBinary+  , ColLargeUtf8+  , ColLargeBinary+  , ColFixedSizeBinary+  , ColUtf8View+  , ColBinaryView+  , ColStruct+  , ColList+  , ColLargeList+  , ColListView+  , ColLargeListView+  , ColFixedSizeList+  , ColMap+  , ColDenseUnion+  , ColSparseUnion+  , ColDictionary+  , ColRunEndEncoded+  #-}+++{-# COMPLETE+  ColNull+  , ColInt8+  , ColInt16+  , ColInt32+  , ColInt64+  , ColUInt8+  , ColUInt16+  , ColUInt32+  , ColUInt64+  , ColFloat16+  , ColFloat+  , ColDouble+  , ColDate32+  , ColDate64+  , ColTime32+  , ColTime64+  , ColTimestamp+  , ColDuration+  , ColIntervalYearMonth+  , ColIntervalDayTime+  , ColIntervalMonthDayNano+  , ColDecimal32+  , ColDecimal64+  , ColDecimal128+  , ColDecimal256+  , ColBool+  , ColUtf8+  , ColBinary+  , ColLargeUtf8+  , ColLargeBinary+  , ColFixedSizeBinary+  , ColUtf8View+  , ColBinaryView+  , ColStruct+  , ColList+  , ColLargeList+  , ColListView+  , ColLargeListView+  , ColFixedSizeList+  , ColMap+  , ColDenseUnion+  , ColSparseUnion+  , ColDictionary+  , ColRunEndEncoded+  #-}+++-- ============================================================+-- Shape+-- ============================================================++-- | Row count. O(1).+columnLength :: ColumnArray -> Int+columnLength = \case+  I.ColNull n -> n+  I.ColPrim t _ xs -> withPrim t (VS.length xs)+  I.ColBool _ b -> bitmapLength b+  I.ColUtf8 _ o _ -> offsetRows o+  I.ColBinary _ o _ -> offsetRows o+  I.ColLargeUtf8 _ o _ -> offsetRows o+  I.ColLargeBinary _ o _ -> offsetRows o+  I.ColFixedSizeBinary _ n _ _ -> n+  I.ColUtf8View _ views _ -> BS.length views `quot` 16+  I.ColBinaryView _ views _ -> BS.length views `quot` 16+  I.ColStruct n _ _ -> n+  I.ColList _ o _ -> offsetRows o+  I.ColLargeList _ o _ -> offsetRows o+  I.ColListView _ o _ _ -> VS.length o+  I.ColLargeListView _ o _ _ -> VS.length o+  I.ColFixedSizeList _ n _ _ -> n+  I.ColMap _ o _ _ -> offsetRows o+  I.ColDenseUnion t _ _ -> VS.length t+  I.ColSparseUnion t _ -> VS.length t+  I.ColDictionary _ ix _ -> columnLength ix+  I.ColRunEndEncoded _ n _ _ -> n+++offsetRows :: Storable o => VS.Vector o -> Int+offsetRows o = max 0 (VS.length o - 1)+{-# INLINE offsetRows #-}+++{- | Null rows of the array itself, O(1). A 'ColNull' is all null;+unions and run-end-encoded columns have no validity of their own (0).+-}+nullCount :: ColumnArray -> Int+nullCount = \case+  I.ColNull n -> n+  c -> maybe 0 validityNullCount (validity c)+++-- | The array's own validity ('Nothing' when it has no nulls or no validity slot).+validity :: ColumnArray -> Maybe Validity+validity = \case+  I.ColNull _ -> Nothing+  I.ColPrim _ v _ -> v+  I.ColBool v _ -> v+  I.ColUtf8 v _ _ -> v+  I.ColBinary v _ _ -> v+  I.ColLargeUtf8 v _ _ -> v+  I.ColLargeBinary v _ _ -> v+  I.ColFixedSizeBinary _ _ v _ -> v+  I.ColUtf8View v _ _ -> v+  I.ColBinaryView v _ _ -> v+  I.ColStruct _ v _ -> v+  I.ColList v _ _ -> v+  I.ColLargeList v _ _ -> v+  I.ColListView v _ _ _ -> v+  I.ColLargeListView v _ _ _ -> v+  I.ColFixedSizeList _ _ v _ -> v+  I.ColMap v _ _ _ -> v+  I.ColDenseUnion {} -> Nothing+  I.ColSparseUnion {} -> Nothing+  I.ColDictionary _ ix _ -> validity ix+  I.ColRunEndEncoded {} -> Nothing+++-- | The constructor name (@ColInt64@, @ColUtf8@, ...), for messages.+columnTag :: ColumnArray -> String+columnTag = \case+  I.ColNull _ -> "ColNull"+  I.ColPrim t _ _ -> "Col" ++ primTypeName t+  I.ColBool {} -> "ColBool"+  I.ColUtf8 {} -> "ColUtf8"+  I.ColBinary {} -> "ColBinary"+  I.ColLargeUtf8 {} -> "ColLargeUtf8"+  I.ColLargeBinary {} -> "ColLargeBinary"+  I.ColFixedSizeBinary {} -> "ColFixedSizeBinary"+  I.ColUtf8View {} -> "ColUtf8View"+  I.ColBinaryView {} -> "ColBinaryView"+  I.ColStruct {} -> "ColStruct"+  I.ColList {} -> "ColList"+  I.ColLargeList {} -> "ColLargeList"+  I.ColListView {} -> "ColListView"+  I.ColLargeListView {} -> "ColLargeListView"+  I.ColFixedSizeList {} -> "ColFixedSizeList"+  I.ColMap {} -> "ColMap"+  I.ColDenseUnion {} -> "ColDenseUnion"+  I.ColSparseUnion {} -> "ColSparseUnion"+  I.ColDictionary {} -> "ColDictionary"+  I.ColRunEndEncoded {} -> "ColRunEndEncoded"+++{- | Whether the array has a validity bitmap slot (everything except+'ColNull', unions and run-end-encoded columns).+-}+hasValiditySlot :: ColumnArray -> Bool+hasValiditySlot = \case+  I.ColNull _ -> False+  I.ColDenseUnion {} -> False+  I.ColSparseUnion {} -> False+  I.ColRunEndEncoded {} -> False+  _ -> True+++-- ============================================================+-- Typed views+-- ============================================================++-- | A fixed-width column's validity and values.+data PrimArray a = PrimArray !(Maybe Validity) !(VS.Vector a)+++-- | O(1). 'Nothing' unless the column has this tag (decimal precision and scale are not compared).+asPrim :: PrimType a -> ColumnArray -> Maybe (PrimArray a)+asPrim t = \case+  I.ColPrim t' v xs | Just Refl <- samePrimTag t t' -> Just (PrimArray v xs)+  _ -> Nothing+{-# INLINE asPrim #-}+++primArrayLength :: Storable a => PrimArray a -> Int+primArrayLength (PrimArray _ xs) = VS.length xs+{-# INLINE primArrayLength #-}+++-- | Row @i@: 'Nothing' when null or out of range. A bit test and a load.+primAt :: Storable a => PrimArray a -> Int -> Maybe a+primAt (PrimArray v xs) i+  | i < 0 || i >= VS.length xs = Nothing+  | unsafeIsValidAt v i = Just $! VS.unsafeIndex xs i+  | otherwise = Nothing+{-# INLINE primAt #-}+++-- | 'primAt' without the range check.+unsafePrimAt :: Storable a => PrimArray a -> Int -> Maybe a+unsafePrimAt (PrimArray v xs) i+  | unsafeIsValidAt v i = Just $! VS.unsafeIndex xs i+  | otherwise = Nothing+{-# INLINE unsafePrimAt #-}+++-- | The value slot of row @i@, ignoring validity (arrow-rs @value()@); errors out of range.+primValueAt :: Storable a => PrimArray a -> Int -> a+primValueAt (PrimArray _ xs) i = xs VS.! i+{-# INLINE primValueAt #-}+++unsafePrimValueAt :: Storable a => PrimArray a -> Int -> a+unsafePrimValueAt (PrimArray _ xs) i = VS.unsafeIndex xs i+{-# INLINE unsafePrimValueAt #-}+++-- | O(1) alias of the values; null slots hold unspecified values.+toStorable :: PrimArray a -> VS.Vector a+toStorable (PrimArray _ xs) = xs+++{- | O(1) for the values (the result aliases the column's value buffer+and validity bitmap); a column without nulls gets a fresh all-set+validity bitmap (rows / 8 bytes).+-}+toMaybeVector :: FixedWidth a => PrimArray a -> AV.Vector (Maybe a)+toMaybeVector (PrimArray v xs) = withValidity (VS.length xs) v (AV.fromStorable xs)+{-# INLINE toMaybeVector #-}+++-- | A var-length column's validity, offsets and data.+data BytesArray o = BytesArray !(Maybe Validity) !(VS.Vector o) !ByteString+++{- | A utf8 column: its data is valid UTF-8 at every row boundary, so+'textAt' copies without re-validating. Only 'asUtf8' and+'asLargeUtf8' produce one.+-}+newtype Utf8Array o = Utf8Array (BytesArray o)+++asUtf8 :: ColumnArray -> Maybe (Utf8Array Int32)+asUtf8 = \case+  I.ColUtf8 v o d -> Just (Utf8Array (BytesArray v o d))+  _ -> Nothing+{-# INLINE asUtf8 #-}+++asLargeUtf8 :: ColumnArray -> Maybe (Utf8Array Int64)+asLargeUtf8 = \case+  I.ColLargeUtf8 v o d -> Just (Utf8Array (BytesArray v o d))+  _ -> Nothing+{-# INLINE asLargeUtf8 #-}+++-- | Binary columns, and utf8 columns viewed as bytes.+asBinary :: ColumnArray -> Maybe (BytesArray Int32)+asBinary = \case+  I.ColBinary v o d -> Just (BytesArray v o d)+  I.ColUtf8 v o d -> Just (BytesArray v o d)+  _ -> Nothing+{-# INLINE asBinary #-}+++asLargeBinary :: ColumnArray -> Maybe (BytesArray Int64)+asLargeBinary = \case+  I.ColLargeBinary v o d -> Just (BytesArray v o d)+  I.ColLargeUtf8 v o d -> Just (BytesArray v o d)+  _ -> Nothing+{-# INLINE asLargeBinary #-}+++bytesArrayLength :: Storable o => BytesArray o -> Int+bytesArrayLength (BytesArray _ o _) = offsetRows o+{-# INLINE bytesArrayLength #-}+++-- | Row @i@ as a zero-copy slice; 'Nothing' when null or out of range.+bytesAt :: (Storable o, Integral o) => BytesArray o -> Int -> Maybe ByteString+bytesAt arr@(BytesArray _ o _) i+  | i < 0 || i >= VS.length o - 1 = Nothing+  | otherwise = unsafeBytesAt arr i+{-# INLINE bytesAt #-}+++unsafeBytesAt :: (Storable o, Integral o) => BytesArray o -> Int -> Maybe ByteString+unsafeBytesAt (BytesArray v o d) i+  | unsafeIsValidAt v i =+      let !s = fromIntegral (VS.unsafeIndex o i)+          !e = fromIntegral (VS.unsafeIndex o (i + 1))+      in Just (BSU.unsafeTake (e - s) (BSU.unsafeDrop s d))+  | otherwise = Nothing+{-# INLINE unsafeBytesAt #-}+++-- | Row @i@ copied into a fresh 'Text' (one copy, no re-validation).+textAt :: (Storable o, Integral o) => Utf8Array o -> Int -> Maybe Text+textAt (Utf8Array arr) i = utf8ToText <$> bytesAt arr i+{-# INLINE textAt #-}+++unsafeTextAt :: (Storable o, Integral o) => Utf8Array o -> Int -> Maybe Text+unsafeTextAt (Utf8Array arr) i = utf8ToText <$> unsafeBytesAt arr i+{-# INLINE unsafeTextAt #-}+++-- | Copy bytes known to be valid UTF-8 into a 'Text'.+utf8ToText :: ByteString -> Text+utf8ToText bs+  | BS.null bs = T.empty+  | otherwise = unsafeDupablePerformIO $ withBytesPtr bs $ \p -> TF.fromPtr p (fromIntegral (BS.length bs))+++data BoolArray = BoolArray !(Maybe Validity) !Bitmap+++asBool :: ColumnArray -> Maybe BoolArray+asBool = \case+  I.ColBool v b -> Just (BoolArray v b)+  _ -> Nothing+{-# INLINE asBool #-}+++boolArrayAt :: BoolArray -> Int -> Maybe Bool+boolArrayAt (BoolArray v b) i+  | i < 0 || i >= bitmapLength b = Nothing+  | unsafeIsValidAt v i = if unsafeBitAt b i then justTrue else justFalse+  | otherwise = Nothing+{-# INLINE boolArrayAt #-}+++boolAt :: ColumnArray -> Int -> Maybe Bool+boolAt c i = asBool c >>= \arr -> boolArrayAt arr i+{-# INLINE boolAt #-}+++{- | Row @i@ of any byte-like column (utf8, binary, their large and+view variants, fixed-size binary) as a zero-copy slice.+-}+anyBytesAt :: ColumnArray -> Int -> Maybe ByteString+anyBytesAt c i+  | i < 0 || i >= columnLength c || not (unsafeIsValidAt (validity c) i) = Nothing+  | otherwise = case c of+      I.ColUtf8 _ o d -> Just $! varSlice o d i+      I.ColBinary _ o d -> Just $! varSlice o d i+      I.ColLargeUtf8 _ o d -> Just $! varSlice o d i+      I.ColLargeBinary _ o d -> Just $! varSlice o d i+      I.ColUtf8View _ views bufs -> Just $! viewAt views bufs i+      I.ColBinaryView _ views bufs -> Just $! viewAt views bufs i+      I.ColFixedSizeBinary w _ _ d -> Just $! BSU.unsafeTake w (BSU.unsafeDrop (i * w) d)+      _ -> Nothing++{-# INLINE anyBytesAt #-}++-- | Row @i@ of a utf8, large utf8 or utf8 view column, copied into a 'Text'.+anyTextAt :: ColumnArray -> Int -> Maybe Text+anyTextAt c i = case c of+  I.ColUtf8 {} -> textOf (anyBytesAt c i)+  I.ColLargeUtf8 {} -> textOf (anyBytesAt c i)+  I.ColUtf8View {} -> textOf (anyBytesAt c i)+  _ -> Nothing+  where+    textOf = \case+      Just bs -> Just $! utf8ToText bs+      Nothing -> Nothing++{-# INLINE anyTextAt #-}++-- | Shared closures for boxed bool rows, so materializing bools allocates nothing per row.+justTrue, justFalse :: Maybe Bool+justTrue = Just True+justFalse = Just False+{-# NOINLINE justTrue #-}+{-# NOINLINE justFalse #-}+++varSlice :: (Storable o, Integral o) => VS.Vector o -> ByteString -> Int -> ByteString+varSlice o d i =+  let !s = fromIntegral (VS.unsafeIndex o i)+      !e = fromIntegral (VS.unsafeIndex o (i + 1))+  in BSU.unsafeTake (e - s) (BSU.unsafeDrop s d)+{-# INLINE varSlice #-}+++-- | Little-endian int32 at byte @k@ (no alignment needed).+le32 :: ByteString -> Int -> Int+le32 bs k =+  let b j = fromIntegral (BSU.unsafeIndex bs (k + j)) :: Word32+      !w = b 0 + b 1 * 0x100 + b 2 * 0x10000 + b 3 * 0x1000000+  in fromIntegral (fromIntegral w :: Int32)+{-# INLINE le32 #-}+++-- | The bytes of view @i@ (the views are valid).+viewAt :: ByteString -> V.Vector ByteString -> Int -> ByteString+viewAt views bufs i =+  let !base = i * 16+      !len = le32 views base+  in if len <= 12+       then BSU.unsafeTake len (BSU.unsafeDrop (base + 4) views)+       else+         let !bi = le32 views (base + 8)+             !off = le32 views (base + 12)+         in BSU.unsafeTake len (BSU.unsafeDrop off (V.unsafeIndex bufs bi))+++{- | The child range of row @i@ of a list, large list, list view, map+or fixed-size list column; 'Nothing' when the row is null, out of+range, or the column is not list-like.+-}+listRange :: ColumnArray -> Int -> Maybe ChildRange+listRange c i+  | i < 0 || i >= columnLength c || not (unsafeIsValidAt (validity c) i) = Nothing+  | otherwise = case c of+      I.ColList _ o _ -> Just (offsetRange o i)+      I.ColLargeList _ o _ -> Just (offsetRange o i)+      I.ColMap _ o _ _ -> Just (offsetRange o i)+      I.ColListView _ o s _ -> Just (ChildRange (fromIntegral (VS.unsafeIndex o i)) (fromIntegral (VS.unsafeIndex s i)))+      I.ColLargeListView _ o s _ -> Just (ChildRange (fromIntegral (VS.unsafeIndex o i)) (fromIntegral (VS.unsafeIndex s i)))+      I.ColFixedSizeList w _ _ _ -> Just (ChildRange (i * w) w)+      _ -> Nothing+++offsetRange :: (Storable o, Integral o) => VS.Vector o -> Int -> ChildRange+offsetRange o i =+  let !s = fromIntegral (VS.unsafeIndex o i)+  in ChildRange s (fromIntegral (VS.unsafeIndex o (i + 1)) - s)+{-# INLINE offsetRange #-}+++-- | The dictionary key of row @i@ of a dictionary column ('Nothing' when null or out of range).+dictKeyAt :: ColumnArray -> Int -> Maybe Int+dictKeyAt c i = case c of+  I.ColDictionary _ ix _ | i >= 0 && i < columnLength ix && unsafeIsValidAt (validity ix) i -> Just (keyAt ix i)+  _ -> Nothing++{-# INLINE dictKeyAt #-}++-- | Key @i@ of an integer column (validity ignored); 0 for a non-integer column.+keyAt :: ColumnArray -> Int -> Int+keyAt c i = case c of+  I.ColPrim t _ xs | Just IntegralPrim <- integralPrim t -> fromIntegral (VS.unsafeIndex xs i)+  _ -> 0+{-# INLINE keyAt #-}+++-- ============================================================+-- Conversions+-- ============================================================++{- | Pair a values vector with a column validity: the bitmap is aliased;+no nulls is a fresh all-set bitmap of @n@ bits.+-}+withValidity :: Int -> Maybe Validity -> AV.Vector a -> AV.Vector (Maybe a)+withValidity n v xs = case v of+  Just (Validity (Bitmap bs off _) _) -> AV.V_Maybe (off .&. 7) (VS.unsafeDrop (off `unsafeShiftR` 3) (AV.bytesToBits bs)) xs+  Nothing -> AV.V_Maybe 0 (AV.allSetBits n) xs+{-# INLINE withValidity #-}+++{- | Utf8, large utf8 and utf8 view columns, and dictionaries whose+values are one of those. Rows are 'Text' slices of one text array and+are never re-validated.++* Utf8 and large utf8: O(rows) (the offsets are widened to 'Int').+  Data in GHC heap memory (an in-memory decoded input, or a column+  built by this package) is aliased without a copy; other memory (an+  mmapped file) has its referenced bytes copied once.+* Utf8 views: O(bytes), every valid row copied once into one buffer.+* Dictionaries: the values convert once and the rows share them+  (O(rows) gather).+-}+toTextVector :: ColumnArray -> Either String (AV.Vector (Maybe Text))+toTextVector c = case c of+  I.ColUtf8 v o d -> Right (withValidity (offsetRows o) v (AV.V_Text (textRows o d)))+  I.ColLargeUtf8 v o d -> Right (withValidity (offsetRows o) v (AV.V_Text (textRows o d)))+  I.ColUtf8View v views bufs ->+    let !n = columnLength c+    in Right (withValidity n v (AV.V_Text (bytesToText (packRows n v (viewAt views bufs)))))+  I.ColDictionary _ keys vals -> toTextVector vals >>= gatherDictionary "Arrow.Column.toTextVector" keys+  _ -> Left ("Arrow.Column.toTextVector: not a utf8 column: " ++ columnTag c)+++-- | The rows of a utf8 column over its data, as a text array.+textRows :: (Storable o, Integral o) => VS.Vector o -> ByteString -> AV.VarVec Text+textRows o d+  | VS.length o < 2 = AV.VarVec VS.empty VS.empty TA.empty+  | otherwise = withTextArray d base (top - base) (offsetsVar o)+  where+    !base = fromIntegral (VS.unsafeIndex o 0)+    !top = fromIntegral (VS.unsafeIndex o (VS.length o - 1))+{-# INLINE textRows #-}+++-- | Byte rows whose store is valid UTF-8 at every row boundary, as text rows.+bytesToText :: AV.VarVec ByteString -> AV.VarVec Text+bytesToText (AV.VarVec st en bs) = withTextArray bs 0 (BS.length bs) $ \arr delta ->+  if delta == 0+    then AV.VarVec st en arr+    else AV.VarVec (VS.map (+ delta) st) (VS.map (+ delta) en) arr+++{- | Rows from Arrow offsets: starts and ends are two overlapping slices+of one widened offsets buffer, each offset shifted by @delta@.+-}+offsetsVar :: (Storable o, Integral o) => VS.Vector o -> AV.Store a -> Int -> AV.VarVec a+offsetsVar o store delta =+  let !n = VS.length o - 1+      !offs = VS.generate (n + 1) (\i -> fromIntegral (VS.unsafeIndex o i) + delta)+  in AV.VarVec (VS.unsafeSlice 0 n offs) (VS.unsafeSlice 1 n offs) store+{-# INLINE offsetsVar #-}+++{- | Bytes @[base, base + len)@ of a 'ByteString' as a text array, and+the shift from a data offset to an array offset. A 'ByteString' backed+by GHC heap memory (every in-memory decoded input and every buffer this+package allocates) is aliased with no copy; other memory is copied once.+-}+withTextArray :: ByteString -> Int -> Int -> (TA.Array -> Int -> r) -> r+withTextArray bs@(BSI.BS (ForeignPtr addr contents) _) base len k = case contents of+  PlainPtr mba -> heap mba+  MallocPtr mba _ -> heap mba+  _ -> k (copyToTextArray bs base len) (negate base)+  where+    heap mba = case freezeHeapBytes mba of+      arr@(TA.ByteArray ba) -> k arr (I# (minusAddr# addr (byteArrayContents# ba)))+{-# INLINE withTextArray #-}+++-- | The heap array behind a 'ByteString' (immutable by contract) as a text array.+freezeHeapBytes :: MutableByteArray# RealWorld -> TA.Array+freezeHeapBytes mba = unsafeDupablePerformIO $ IO $ \s -> case unsafeFreezeByteArray# mba s of+  (# s', ba #) -> (# s', TA.ByteArray ba #)+++-- | @len@ bytes of @bs@ from byte @off@, copied into a fresh text array.+copyToTextArray :: ByteString -> Int -> Int -> TA.Array+copyToTextArray bs off len+  | len <= 0 = TA.empty+  | otherwise = unsafeDupablePerformIO $ withBytesPtr bs $ \p -> stToIO $ do+      ma <- TA.new len+      TA.copyFromPointer ma 0 (p `plusPtr` off) len+      TA.unsafeFreeze ma+++{- | Every byte-like column, and dictionaries whose values are byte-like.+Rows are zero-copy 'ByteString' slices.++* Utf8, binary and their large variants: O(rows) (offsets widened),+  the data buffer aliased.+* Fixed-size binary: O(rows), the data buffer aliased.+* Views: O(bytes), every valid row copied once into one buffer.+* Dictionaries: the values convert once and the rows share them+  (O(rows) gather).+-}+toBytesVector :: ColumnArray -> Either String (AV.Vector (Maybe ByteString))+toBytesVector c = case c of+  I.ColUtf8 v o d -> Right (offsetBytes v o d)+  I.ColBinary v o d -> Right (offsetBytes v o d)+  I.ColLargeUtf8 v o d -> Right (offsetBytes v o d)+  I.ColLargeBinary v o d -> Right (offsetBytes v o d)+  I.ColUtf8View v views bufs -> Right (withValidity n v (AV.V_Bytes (packRows n v (viewAt views bufs))))+  I.ColBinaryView v views bufs -> Right (withValidity n v (AV.V_Bytes (packRows n v (viewAt views bufs))))+  I.ColFixedSizeBinary w _ v d ->+    Right (withValidity n v (AV.V_Bytes (AV.VarVec (VS.generate n (* w)) (VS.generate n (\i -> i * w + w)) d)))+  I.ColDictionary _ keys vals -> toBytesVector vals >>= gatherDictionary "Arrow.Column.toBytesVector" keys+  _ -> Left ("Arrow.Column.toBytesVector: not a byte column: " ++ columnTag c)+  where+    !n = columnLength c+    offsetBytes :: (Storable o, Integral o) => Maybe Validity -> VS.Vector o -> ByteString -> AV.Vector (Maybe ByteString)+    offsetBytes v o d+      | VS.length o < 2 = withValidity 0 v (AV.V_Bytes (AV.VarVec VS.empty VS.empty BS.empty))+      | otherwise = withValidity (offsetRows o) v (AV.V_Bytes (offsetsVar o d 0))+    {-# INLINE offsetBytes #-}+++{- | O(bytes): the valid rows given by @at@ copied into one fresh buffer+(null rows are empty).+-}+packRows :: Int -> Maybe Validity -> (Int -> ByteString) -> AV.VarVec ByteString+packRows n v at = unsafeDupablePerformIO $ do+  let lenAt i = if unsafeIsValidAt v i then BS.length (at i) else 0+      total = go 0 0+        where+          go !i !acc+            | i >= n = acc+            | otherwise = go (i + 1) (acc + lenAt i)+  fp <- mallocAligned total+  st <- VSM.unsafeNew n+  en <- VSM.unsafeNew n+  unsafeWithForeignPtr fp $ \dst -> do+    let fill !i !pos+          | i >= n = pure ()+          | unsafeIsValidAt v i = do+              let row = at i+                  !len = BS.length row+              withBytesPtr row $ \src -> copyBytes (dst `plusPtr` pos) src len+              VSM.unsafeWrite st i pos+              VSM.unsafeWrite en i (pos + len)+              fill (i + 1) (pos + len)+          | otherwise = do+              VSM.unsafeWrite st i pos+              VSM.unsafeWrite en i pos+              fill (i + 1) pos+    fill 0 0+  AV.VarVec <$> VS.unsafeFreeze st <*> VS.unsafeFreeze en <*> pure (BSI.BS fp total)+{-# INLINE packRows #-}+++{- | O(1) for the bits (the result aliases the column's bitmaps); a+column without nulls gets a fresh all-set validity bitmap.+-}+toBoolVector :: ColumnArray -> Either String (AV.Vector (Maybe Bool))+toBoolVector c = case c of+  I.ColBool v (Bitmap bs off len) ->+    Right (withValidity len v (AV.V_Bool (off .&. 7) len (VS.unsafeDrop (off `unsafeShiftR` 3) (AV.bytesToBits bs))))+  _ -> Left ("Arrow.Column.toBoolVector: not a bool column: " ++ columnTag c)+++{- | List, large list, list view and fixed-size list columns, O(rows)+plus the child conversion. The child rows the lists reference are+converted once with the given function and every list row is a slice+of that child vector; null rows are 'Nothing'. The function must+return one element per child row it is given.+-}+toListVector :: forall a. AV.Element a => (ColumnArray -> Either String (AV.Vector a)) -> ColumnArray -> Either String (AV.Vector (Maybe (AV.Vector a)))+toListVector conv c = case c of+  I.ColList v o child -> offsets v o child+  I.ColLargeList v o child -> offsets v o child+  I.ColListView v o z child -> views v o z child+  I.ColLargeListView v o z child -> views v o z child+  I.ColFixedSizeList w n v child -> do+    kids <- convert (sliceColumnArray 0 (n * w) child) (n * w)+    Right (withValidity n v (AV.V_List (AV.VarVec (VS.generate n (* w)) (VS.generate n (\i -> i * w + w)) kids)))+  _ -> Left ("Arrow.Column.toListVector: not a list column: " ++ columnTag c)+  where+    convert ch len = do+      kids <- conv ch+      if G.length kids == len+        then Right kids+        else Left ("Arrow.Column.toListVector: the child conversion returned " ++ show (G.length kids) ++ " rows for " ++ show len)+    offsets :: (Storable o, Integral o) => Maybe Validity -> VS.Vector o -> ColumnArray -> Either String (AV.Vector (Maybe (AV.Vector a)))+    offsets v o child+      | VS.length o < 2 = Right (withValidity 0 v (AV.V_List (AV.VarVec VS.empty VS.empty G.empty)))+      | otherwise = do+          let !n = VS.length o - 1+              !base = fromIntegral (VS.unsafeIndex o 0)+              !len = fromIntegral (VS.unsafeIndex o n) - base+          kids <- convert (sliceColumnArray base len child) len+          Right (withValidity n v (AV.V_List (offsetsVar o kids (negate base))))+    {-# INLINE offsets #-}+    -- Null list-view rows may hold any offset and size, so they become empty rows.+    views :: (Storable o, Integral o) => Maybe Validity -> VS.Vector o -> VS.Vector o -> ColumnArray -> Either String (AV.Vector (Maybe (AV.Vector a)))+    views v o z child = do+      kids <- convert child (columnLength child)+      let !n = VS.length o+          st = VS.generate n (\i -> if unsafeIsValidAt v i then fromIntegral (VS.unsafeIndex o i) else 0)+          en = VS.generate n (\i -> if unsafeIsValidAt v i then fromIntegral (VS.unsafeIndex o i) + fromIntegral (VS.unsafeIndex z i) else 0)+      Right (withValidity n v (AV.V_List (AV.VarVec st en kids)))+    {-# INLINE views #-}+++{- | Rows of a dictionary column from its converted values, O(rows):+row @i@ is the value its key selects (sharing the values' store), null+where the key or the selected value is null. Every valid key is checked+against the values first.+-}+gatherDictionary :: forall a. AV.VarElem a => String -> ColumnArray -> AV.Vector (Maybe a) -> Either String (AV.Vector (Maybe a))+gatherDictionary what keys (AV.V_Maybe vo vbits vals) = do+  let !(AV.VarVec vst ven store) = AV.toVar vals+  validateKeys what keys (VS.length vst)+  let go :: (Storable k, Integral k) => Maybe Validity -> VS.Vector k -> AV.Vector (Maybe a)+      go kv xs =+        let !n = VS.length xs+            key i = fromIntegral (VS.unsafeIndex xs i) :: Int+            valid i = unsafeIsValidAt kv i && AV.indexBit vbits (vo + key i)+            !st = VS.generate n (\i -> if valid i then VS.unsafeIndex vst (key i) else 0)+            !en = VS.generate n (\i -> if valid i then VS.unsafeIndex ven (key i) else 0)+        in case bitmapGenerate n valid of+             Bitmap bs _ _ -> AV.V_Maybe 0 (AV.bytesToBits bs) (AV.fromVar (AV.VarVec st en store))+      {-# INLINE go #-}+  case keys of+    I.ColPrim t kv xs -> case t of+      PInt8 -> Right (go kv xs)+      PInt16 -> Right (go kv xs)+      PInt32 -> Right (go kv xs)+      PInt64 -> Right (go kv xs)+      PUInt8 -> Right (go kv xs)+      PUInt16 -> Right (go kv xs)+      PUInt32 -> Right (go kv xs)+      PUInt64 -> Right (go kv xs)+      _ -> notInt+    _ -> notInt+  where+    notInt = Left (what ++ ": dictionary indices must be an integer column")+++-- ============================================================+-- Construction+-- ============================================================++-- | O(1): a fixed-width column without nulls.+primColumn :: PrimType a -> VS.Vector a -> ColumnArray+primColumn t xs = I.ColPrim t Nothing xs+++-- | O(1) length check plus a popcount: values with a validity bitmap (bit set = valid).+primColumnV :: PrimType a -> Bitmap -> VS.Vector a -> Either String ColumnArray+primColumnV t b xs = withPrim t $ do+  _ <- mkBitmap (bitmapBytes b) (bitmapOffset b) (bitmapLength b)+  if bitmapLength b /= VS.length xs+    then Left "Arrow.Column.primColumnV: bitmap length differs from the number of values"+    else Right (I.ColPrim t (mkValidity b) xs)+++-- | A fixed-width column with an optional validity (checked and normalised).+mkPrim :: PrimType a -> Maybe Validity -> VS.Vector a -> Either String ColumnArray+mkPrim t v xs = withPrim t $ I.ColPrim t <$> checkValidity "Arrow.Column.mkPrim" (VS.length xs) v <*> pure xs+++-- | O(n), one pass: 'Nothing' rows become nulls (their slots hold zero bytes).+fromMaybes :: PrimType a -> V.Vector (Maybe a) -> ColumnArray+fromMaybes t vals = withPrim t $+  let !n = V.length vals+      !w = primWidth t+      bytes = createAligned (n * w) $ \p -> do+        fillBytes p 0 (n * w)+        forM_ [0 .. n - 1] $ \i -> case V.unsafeIndex vals i of+          Just x -> pokeElemOff (castPtr p) i x+          Nothing -> pure ()+  in I.ColPrim t (validityGenerate n (\i -> isJustAt vals i)) (unsafeBytesToStorable bytes)+++isJustAt :: V.Vector (Maybe a) -> Int -> Bool+isJustAt v i = case V.unsafeIndex v i of+  Just _ -> True+  Nothing -> False+{-# INLINE isJustAt #-}+++fromBools :: V.Vector Bool -> ColumnArray+fromBools v = I.ColBool Nothing (bitmapFromBools v)+++fromMaybeBools :: V.Vector (Maybe Bool) -> ColumnArray+fromMaybeBools v =+  I.ColBool+    (validityGenerate (V.length v) (isJustAt v))+    (bitmapGenerate (V.length v) (\i -> V.unsafeIndex v i == Just True))+++{- | O(1) plus a popcount: a fixed-width column over the vector's own+buffers (no copy; the inverse of 'toMaybeVector'). Null slots keep+whatever values the vector holds there.+-}+fromMaybeVector :: PrimType a -> AV.Vector (Maybe a) -> ColumnArray+fromMaybeVector t (AV.V_Maybe off bits xs) = withPrim t $+  let !vals = AV.storableVector xs+  in I.ColPrim t (mkValidity (Bitmap (AV.bitsToBytes bits) off (VS.length vals))) vals+++-- | O(1) plus a popcount: a bool column over the vector's own bitmaps (the inverse of 'toBoolVector').+fromBoolVector :: AV.Vector (Maybe Bool) -> ColumnArray+fromBoolVector (AV.V_Maybe off bits (AV.V_Bool bo n bs)) =+  I.ColBool (mkValidity (Bitmap (AV.bitsToBytes bits) off n)) (Bitmap (AV.bitsToBytes bs) bo n)++++-- | Utf8 column, two passes (lengths, then one copy per row). Errors past 2^31 - 1 bytes.+fromTexts :: V.Vector Text -> ColumnArray+fromTexts = fromMaybeTexts . V.map Just+++fromMaybeTexts :: V.Vector (Maybe Text) -> ColumnArray+fromMaybeTexts = varFromMaybes I.ColUtf8 TF.lengthWord8 TF.unsafeCopyToPtr+++fromMaybeLargeTexts :: V.Vector (Maybe Text) -> ColumnArray+fromMaybeLargeTexts = varFromMaybes I.ColLargeUtf8 TF.lengthWord8 TF.unsafeCopyToPtr+++-- | Construct a binary column from non-null rows, copying their bytes.+fromByteStrings :: V.Vector ByteString -> ColumnArray+fromByteStrings vals = unsafeDupablePerformIO $ do+  -- Check the sum before narrowing to Arrow's signed 32-bit offsets.+  let !n = V.length vals+      !total = V.foldl' addLength 0 vals+      addLength !acc bytes+        | BS.length bytes > fromIntegral (maxBound :: Int32) - acc =+            errorWithoutStackTrace "Arrow.Column: var-length data exceeds the 32-bit offset range (use the large variant)"+        | otherwise = acc + BS.length bytes+  offsets <- VSM.unsafeNew (n + 1)+  fp <- mallocAligned total+  unsafeWithForeignPtr fp $ \dst -> do+    let fill !i !pos+          | i == n = VSM.unsafeWrite offsets n (fromIntegral pos)+          | otherwise = do+              VSM.unsafeWrite offsets i (fromIntegral pos)+              let bytes = V.unsafeIndex vals i+              copyBS bytes (dst `plusPtr` pos)+              fill (i + 1) (pos + BS.length bytes)+    fill 0 0+  offs <- VS.unsafeFreeze offsets+  pure (I.ColBinary Nothing offs (BSI.BS fp total))+++fromMaybeByteStrings :: V.Vector (Maybe ByteString) -> ColumnArray+fromMaybeByteStrings = varFromMaybes I.ColBinary BS.length copyBS+++fromMaybeLargeByteStrings :: V.Vector (Maybe ByteString) -> ColumnArray+fromMaybeLargeByteStrings = varFromMaybes I.ColLargeBinary BS.length copyBS+++copyBS :: ByteString -> Ptr Word8 -> IO ()+copyBS bs dst = withBytesPtr bs $ \src -> copyBytes dst src (BS.length bs)+++varFromMaybes+  :: forall o x+   . Offset o+  => (Maybe Validity -> VS.Vector o -> ByteString -> ColumnArray)+  -> (x -> Int)+  -> (x -> Ptr Word8 -> IO ())+  -> V.Vector (Maybe x)+  -> ColumnArray+varFromMaybes mk len copy vals = unsafeDupablePerformIO $ do+  let !n = V.length vals+      !limit = fromInteger (min (toInteger (maxBound :: o)) (toInteger (maxBound :: Int))) :: Int+      addLength !acc Nothing = acc+      addLength !acc (Just value)+        | len value > limit - acc =+            errorWithoutStackTrace "Arrow.Column: var-length data exceeds the offset range (use the large variant for 32-bit columns)"+        | otherwise = acc + len value+      !total = V.foldl' addLength 0 vals+  offsets <- VSM.unsafeNew (n + 1)+  fp <- mallocAligned total+  unsafeWithForeignPtr fp $ \dst -> do+    let fill !i !pos+          | i == n = VSM.unsafeWrite offsets n (fromIntegral pos)+          | otherwise = do+              VSM.unsafeWrite offsets i (fromIntegral pos)+              case V.unsafeIndex vals i of+                Nothing -> fill (i + 1) pos+                Just value -> do+                  copy value (dst `plusPtr` pos)+                  fill (i + 1) (pos + len value)+    fill 0 0+  offs <- VS.unsafeFreeze offsets+  pure (mk (validityGenerate n (isJustAt vals)) offs (BSI.BS fp total))+{-# INLINE varFromMaybes #-}+++-- | Fixed-size binary rows; every present row must have the given width.+fromMaybeFixedSizeBinary :: Int -> V.Vector (Maybe ByteString) -> Either String ColumnArray+fromMaybeFixedSizeBinary w vals+  | w < 0 = Left "Arrow.Column.fromMaybeFixedSizeBinary: negative width"+  | V.any (maybe False ((/= w) . BS.length)) vals =+      Left ("Arrow.Column.fromMaybeFixedSizeBinary: a row is not " ++ show w ++ " bytes")+  | otherwise =+      let !n = V.length vals+          dat = createAligned (n * w) $ \p -> do+            fillBytes p 0 (n * w)+            forM_ [0 .. n - 1] $ \i -> forM_ (V.unsafeIndex vals i) $ \bs -> copyBS bs (p `plusPtr` (i * w))+      in Right (I.ColFixedSizeBinary w n (validityGenerate n (isJustAt vals)) dat)+++-- | Utf8 view column: short strings inline, long ones in one data buffer.+fromMaybeUtf8View :: V.Vector (Maybe Text) -> ColumnArray+fromMaybeUtf8View vals = buildViews I.ColUtf8View (V.map (fmap textBytes) vals)+  where+    textBytes t = createAligned (TF.lengthWord8 t) (TF.unsafeCopyToPtr t)+++fromMaybeBinaryView :: V.Vector (Maybe ByteString) -> ColumnArray+fromMaybeBinaryView = buildViews I.ColBinaryView+++buildViews :: (Maybe Validity -> ByteString -> V.Vector ByteString -> ColumnArray) -> V.Vector (Maybe ByteString) -> ColumnArray+buildViews mk vals+  | outOfLine > fromIntegral (maxBound :: Int32) =+      errorWithoutStackTrace "Arrow.Column: view data exceeds 2^31 - 1 bytes"+  | otherwise = mk (validityGenerate n (isJustAt vals)) views (if outOfLine == 0 then V.empty else V.singleton dat)+  where+    !n = V.length vals+    long = maybe 0 (\b -> if BS.length b > 12 then BS.length b else 0)+    !outOfLine = V.sum (V.map long vals)+    dat = createAligned outOfLine $ \p ->+      V.foldM'_ (\off mb -> case mb of+        Just b | BS.length b > 12 -> copyBS b (p `plusPtr` off) >> pure (off + BS.length b)+        _ -> pure off) 0 vals+    views = createAligned (n * 16) $ \p -> do+      fillBytes p 0 (n * 16)+      V.ifoldM'_ (\off i mb -> case mb of+        Nothing -> pure off+        Just b -> do+          let !base = p `plusPtr` (i * 16)+              !len = BS.length b+          pokeByteOff base 0 (fromIntegral len :: Int32)+          if len <= 12+            then copyBS b (base `plusPtr` 4) >> pure off+            else do+              copyBS (BS.take 4 b) (base `plusPtr` 4)+              pokeByteOff base 8 (0 :: Int32)+              pokeByteOff base 12 (fromIntegral off :: Int32)+              pure (off + len)) 0 vals+++mkBool :: Maybe Validity -> Bitmap -> Either String ColumnArray+mkBool v b = do+  _ <- mkBitmap (bitmapBytes b) (bitmapOffset b) (bitmapLength b)+  v' <- checkValidity "Arrow.Column.mkBool" (bitmapLength b) v+  Right (I.ColBool v' b)+++mkVar :: Offset o => String -> Bool -> (Maybe Validity -> VS.Vector o -> ByteString -> ColumnArray) -> Maybe Validity -> VS.Vector o -> ByteString -> Either String ColumnArray+mkVar what utf8 mk v offs dat = do+  validateOffsets what (BS.length dat) offs+  when utf8 (validateUtf8 what offs dat)+  v' <- checkValidity what (offsetRows offs) v+  Right (mk v' offs dat)+++-- | Validates offsets, UTF-8 and character boundaries (C kernels).+mkUtf8 :: Maybe Validity -> VS.Vector Int32 -> ByteString -> Either String ColumnArray+mkUtf8 = mkVar "Arrow.Column.mkUtf8" True I.ColUtf8+++mkBinary :: Maybe Validity -> VS.Vector Int32 -> ByteString -> Either String ColumnArray+mkBinary = mkVar "Arrow.Column.mkBinary" False I.ColBinary+++mkLargeUtf8 :: Maybe Validity -> VS.Vector Int64 -> ByteString -> Either String ColumnArray+mkLargeUtf8 = mkVar "Arrow.Column.mkLargeUtf8" True I.ColLargeUtf8+++mkLargeBinary :: Maybe Validity -> VS.Vector Int64 -> ByteString -> Either String ColumnArray+mkLargeBinary = mkVar "Arrow.Column.mkLargeBinary" False I.ColLargeBinary+++-- | Width, rows, validity, data.+mkFixedSizeBinary :: Int -> Int -> Maybe Validity -> ByteString -> Either String ColumnArray+mkFixedSizeBinary w n v dat+  | w < 0 || n < 0 = Left "Arrow.Column.mkFixedSizeBinary: negative width or row count"+  | toInteger (BS.length dat) < toInteger w * toInteger n = Left "Arrow.Column.mkFixedSizeBinary: data shorter than width * rows"+  | otherwise = (\v' -> I.ColFixedSizeBinary w n v' dat) <$> checkValidity "Arrow.Column.mkFixedSizeBinary" n v+++mkViews :: String -> Bool -> (Maybe Validity -> ByteString -> V.Vector ByteString -> ColumnArray) -> Maybe Validity -> ByteString -> V.Vector ByteString -> Either String ColumnArray+mkViews what utf8 mk v views bufs+  | BS.length views `rem` 16 /= 0 = Left (what ++ ": views buffer is not a multiple of 16 bytes")+  | otherwise = do+      let !rows = BS.length views `quot` 16+      v' <- checkValidity what rows v+      validateViews what utf8 rows v' views bufs+      Right (mk v' views bufs)+++mkUtf8View :: Maybe Validity -> ByteString -> V.Vector ByteString -> Either String ColumnArray+mkUtf8View = mkViews "Arrow.Column.mkUtf8View" True I.ColUtf8View+++mkBinaryView :: Maybe Validity -> ByteString -> V.Vector ByteString -> Either String ColumnArray+mkBinaryView = mkViews "Arrow.Column.mkBinaryView" False I.ColBinaryView+++-- | Rows, validity, children (each with at least rows rows).+mkStruct :: Int -> Maybe Validity -> V.Vector (Text, ColumnArray) -> Either String ColumnArray+mkStruct n v cs+  | n < 0 = Left "Arrow.Column.mkStruct: negative row count"+  | V.any ((< n) . columnLength . snd) cs = Left "Arrow.Column.mkStruct: a child has fewer rows than the struct"+  | otherwise = (\v' -> I.ColStruct n v' cs) <$> checkValidity "Arrow.Column.mkStruct" n v+++mkListLike :: Offset o => String -> (Maybe Validity -> VS.Vector o -> ColumnArray -> ColumnArray) -> Maybe Validity -> VS.Vector o -> ColumnArray -> Either String ColumnArray+mkListLike what mk v offs child = do+  validateOffsets what (columnLength child) offs+  v' <- checkValidity what (offsetRows offs) v+  Right (mk v' offs child)+++mkList :: Maybe Validity -> VS.Vector Int32 -> ColumnArray -> Either String ColumnArray+mkList = mkListLike "Arrow.Column.mkList" I.ColList+++mkLargeList :: Maybe Validity -> VS.Vector Int64 -> ColumnArray -> Either String ColumnArray+mkLargeList = mkListLike "Arrow.Column.mkLargeList" I.ColLargeList+++-- | Validity, offsets, sizes, child.+mkListView :: Maybe Validity -> VS.Vector Int32 -> VS.Vector Int32 -> ColumnArray -> Either String ColumnArray+mkListView v o s c = do+  v' <- checkValidity "Arrow.Column.mkListView" (VS.length o) v+  validateListView "Arrow.Column.mkListView" v' o s (columnLength c)+  Right (I.ColListView v' o s c)+++mkLargeListView :: Maybe Validity -> VS.Vector Int64 -> VS.Vector Int64 -> ColumnArray -> Either String ColumnArray+mkLargeListView v o s c = do+  v' <- checkValidity "Arrow.Column.mkLargeListView" (VS.length o) v+  validateListView "Arrow.Column.mkLargeListView" v' o s (columnLength c)+  Right (I.ColLargeListView v' o s c)+++-- | List size, rows, validity, child (at least rows * size rows).+mkFixedSizeList :: Int -> Int -> Maybe Validity -> ColumnArray -> Either String ColumnArray+mkFixedSizeList w n v c+  | w < 0 || n < 0 = Left "Arrow.Column.mkFixedSizeList: negative size or row count"+  | toInteger (columnLength c) < toInteger w * toInteger n = Left "Arrow.Column.mkFixedSizeList: child shorter than rows * size"+  | otherwise = (\v' -> I.ColFixedSizeList w n v' c) <$> checkValidity "Arrow.Column.mkFixedSizeList" n v+++-- | Validity, offsets, keys, values (offsets checked against both children).+mkMap :: Maybe Validity -> VS.Vector Int32 -> ColumnArray -> ColumnArray -> Either String ColumnArray+mkMap v offs k x = do+  validateOffsets "Arrow.Column.mkMap" (min (columnLength k) (columnLength x)) offs+  v' <- checkValidity "Arrow.Column.mkMap" (offsetRows offs) v+  Right (I.ColMap v' offs k x)+++-- | Child index per row, offset per row, children.+mkDenseUnion :: VS.Vector Int8 -> VS.Vector Int32 -> V.Vector ColumnArray -> Either String ColumnArray+mkDenseUnion t o cs = do+  validateDenseUnion "Arrow.Column.mkDenseUnion" t o (VS.fromList (map (fromIntegral . columnLength) (V.toList cs)))+  Right (I.ColDenseUnion t o cs)+++-- | Child index per row, children (each with at least as many rows as the union).+mkSparseUnion :: VS.Vector Int8 -> V.Vector ColumnArray -> Either String ColumnArray+mkSparseUnion t cs+  | V.any ((< VS.length t) . columnLength) cs = Left "Arrow.Column.mkSparseUnion: a child has fewer rows than the union"+  | otherwise = do+      validateSparseUnionTypes "Arrow.Column.mkSparseUnion" t (V.length cs)+      Right (I.ColSparseUnion t cs)+++-- | Dictionary id, indices (an integer column), values. Checks every valid key (C kernel).+mkDictionary :: Int64 -> ColumnArray -> ColumnArray -> Either String ColumnArray+mkDictionary did ix vals = do+  validateKeys "Arrow.Column.mkDictionary" ix (columnLength vals)+  Right (I.ColDictionary did ix vals)+++{- | Run ends (an int16, int32 or int64 column, positive and strictly+increasing) and values (one row per run). The logical length is the+last run end.+-}+mkRunEndEncoded :: ColumnArray -> ColumnArray -> Either String ColumnArray+mkRunEndEncoded re vals = do+  validateRunEnds "Arrow.Column.mkRunEndEncoded" re 0+  let !runs = columnLength re+      !len = if runs == 0 then 0 else keyAt re (runs - 1)+  if columnLength vals < runs+    then Left "Arrow.Column.mkRunEndEncoded: fewer values than runs"+    else Right (I.ColRunEndEncoded 0 len re vals)+++emptyPrim :: PrimType a -> ColumnArray+emptyPrim t = withPrim t (I.ColPrim t Nothing VS.empty)+++offsets0 :: Storable o => Num o => VS.Vector o+offsets0 = VS.singleton 0+++{- | The zero-row column of a field's type, recursively for children,+with dictionary encoding (indices of the declared index type and a+'placeholderColumn' as values).+-}+emptyColumnFor :: Field -> Either String ColumnArray+emptyColumnFor f = case fieldDictionary f of+  Just de -> do+    vals <- placeholderColumn f+    ix <- case primTypeFor (deIndexType de) of+      Just (SomePrimType t) | Just IntegralPrim <- integralPrim t -> Right (emptyPrim t)+      _ -> Left ("Arrow.Column: dictionary index type must be an integer type, got " ++ show (deIndexType de))+    Right (I.ColDictionary (deId de) ix vals)+  Nothing -> case fieldType f of+    ANull -> Right (I.ColNull 0)+    AInt w _+      | Just (SomePrimType t) <- primTypeFor (fieldType f) -> Right (emptyPrim t)+      | otherwise -> Left ("Arrow.Column: unsupported integer bit width " ++ show w)+    ty | Just (SomePrimType t) <- primTypeFor ty -> Right (emptyPrim t)+    ABool -> Right (I.ColBool Nothing emptyBitmap)+    AUtf8 -> Right (I.ColUtf8 Nothing offsets0 BS.empty)+    ABinary -> Right (I.ColBinary Nothing offsets0 BS.empty)+    ALargeUtf8 -> Right (I.ColLargeUtf8 Nothing offsets0 BS.empty)+    ALargeBinary -> Right (I.ColLargeBinary Nothing offsets0 BS.empty)+    AFixedSizeBinary w -> Right (I.ColFixedSizeBinary w 0 Nothing BS.empty)+    AUtf8View -> Right (I.ColUtf8View Nothing BS.empty V.empty)+    ABinaryView -> Right (I.ColBinaryView Nothing BS.empty V.empty)+    AStruct -> I.ColStruct 0 Nothing <$> V.mapM (\c -> (,) (fieldName c) <$> emptyColumnFor c) (fieldChildren f)+    AList -> I.ColList Nothing offsets0 <$> onlyChild+    ALargeList -> I.ColLargeList Nothing offsets0 <$> onlyChild+    AListView -> I.ColListView Nothing VS.empty VS.empty <$> onlyChild+    ALargeListView -> I.ColLargeListView Nothing VS.empty VS.empty <$> onlyChild+    AFixedSizeList n -> I.ColFixedSizeList n 0 Nothing <$> onlyChild+    AMap _ -> case V.toList (fieldChildren f) of+      [entries] | [kf, vf] <- V.toList (fieldChildren entries) -> I.ColMap Nothing offsets0 <$> emptyColumnFor kf <*> emptyColumnFor vf+      _ -> Left "Arrow.Column: map field must have one entries struct child with key and value"+    AUnion Dense _ -> I.ColDenseUnion VS.empty VS.empty <$> V.mapM emptyColumnFor (fieldChildren f)+    AUnion Sparse _ -> I.ColSparseUnion VS.empty <$> V.mapM emptyColumnFor (fieldChildren f)+    ARunEndEncoded -> case V.toList (fieldChildren f) of+      [ref, vf] -> I.ColRunEndEncoded 0 0 <$> emptyColumnFor ref {fieldNullable = False} <*> emptyColumnFor vf+      _ -> Left "Arrow.Column: RunEndEncoded field must have exactly two children (run_ends, values)"+    ty -> Left ("Arrow.Column: no column for type " ++ show ty)+  where+    onlyChild = case V.toList (fieldChildren f) of+      [c] -> emptyColumnFor c+      _ -> Left ("Arrow.Column: " ++ show (fieldType f) ++ " field must have exactly one child")+++{- | The empty values column of a dictionary field, held by an+unresolved dictionary column until 'resolveDictionaryColumn'.+-}+placeholderColumn :: Field -> Either String ColumnArray+placeholderColumn f = emptyColumnFor f {fieldNullable = False, fieldDictionary = Nothing}+++zeroBytes :: Int -> ByteString+zeroBytes n = createAligned n (\p -> fillBytes p 0 n)+++zeroStorable :: forall a. Storable a => Int -> VS.Vector a+zeroStorable n = unsafeBytesToStorable (zeroBytes (n * sizeOf (undefined :: a)))+++{- | @n@ rows with the shape of the given column (same tag, widths,+decimal parameters, child shapes and dictionary id), every row valid+and zero or empty: lists are empty, unions select their first child,+a run-end-encoded column is one run, a 'ColNull' stays null. Used+where a column needs a row count but its rows carry no meaning (the+children under null struct rows, the expansion of null rows over an+empty dictionary).+-}+fillerColumn :: Int -> ColumnArray -> ColumnArray+fillerColumn n0 col = case col of+  I.ColNull _ -> I.ColNull n+  I.ColPrim t _ _ -> withPrim t (I.ColPrim t Nothing (zeroStorable n))+  I.ColBool _ _ -> I.ColBool Nothing (Bitmap (zeroBytes ((n + 7) `unsafeShiftR` 3)) 0 n)+  I.ColUtf8 {} -> I.ColUtf8 Nothing (zeroStorable (n + 1)) BS.empty+  I.ColBinary {} -> I.ColBinary Nothing (zeroStorable (n + 1)) BS.empty+  I.ColLargeUtf8 {} -> I.ColLargeUtf8 Nothing (zeroStorable (n + 1)) BS.empty+  I.ColLargeBinary {} -> I.ColLargeBinary Nothing (zeroStorable (n + 1)) BS.empty+  I.ColFixedSizeBinary w _ _ _ -> I.ColFixedSizeBinary w n Nothing (zeroBytes (n * w))+  I.ColUtf8View {} -> I.ColUtf8View Nothing (zeroBytes (n * 16)) V.empty+  I.ColBinaryView {} -> I.ColBinaryView Nothing (zeroBytes (n * 16)) V.empty+  I.ColStruct _ _ cs -> I.ColStruct n Nothing (V.map (fmap (fillerColumn n)) cs)+  I.ColList _ _ c -> I.ColList Nothing (zeroStorable (n + 1)) (empty c)+  I.ColLargeList _ _ c -> I.ColLargeList Nothing (zeroStorable (n + 1)) (empty c)+  I.ColListView _ _ _ c -> I.ColListView Nothing (zeroStorable n) (zeroStorable n) (empty c)+  I.ColLargeListView _ _ _ c -> I.ColLargeListView Nothing (zeroStorable n) (zeroStorable n) (empty c)+  I.ColFixedSizeList w _ _ c -> I.ColFixedSizeList w n Nothing (fillerColumn (n * w) c)+  I.ColMap _ _ k x -> I.ColMap Nothing (zeroStorable (n + 1)) (empty k) (empty x)+  I.ColDenseUnion _ _ cs+    | n == 0 || V.null cs -> I.ColDenseUnion VS.empty VS.empty (V.map empty cs)+    | otherwise -> I.ColDenseUnion (zeroStorable n) (zeroStorable n) (V.imap (\i c -> if i == 0 then fillerColumn 1 c else empty c) cs)+  I.ColSparseUnion _ cs -> I.ColSparseUnion (zeroStorable n) (V.map (fillerColumn n) cs)+  I.ColDictionary did ix vals+    | n == 0 -> I.ColDictionary did (fillerColumn 0 ix) vals+    | columnLength vals == 0 -> I.ColDictionary did (fillerColumn n ix) (fillerColumn 1 vals)+    | otherwise -> I.ColDictionary did (fillerColumn n ix) vals+  I.ColRunEndEncoded _ _ re vals+    | n == 0 -> I.ColRunEndEncoded 0 0 (fillerColumn 0 re) (empty vals)+    | otherwise -> I.ColRunEndEncoded 0 n (singleRunEnd re n) (fillerColumn 1 vals)+  where+    !n = max 0 n0+    empty = sliceColumnArray 0 0+++-- | A one-element run-end column of the same tag holding @n@.+singleRunEnd :: ColumnArray -> Int -> ColumnArray+singleRunEnd re n = case re of+  I.ColPrim t _ _ | Just IntegralPrim <- integralPrim t -> I.ColPrim t Nothing (VS.singleton (fromIntegral n))+  _ -> re+++-- ============================================================+-- Slicing+-- ============================================================++{- | Rows @[start, start + len)@, clamped to the column (a negative+start becomes 0, a length past the end is truncated). O(1) for flat+and offset columns (bit offsets move, offsets are not rebased, the+child is kept), O(fields) for structs and sparse unions (children are+sliced lazily), O(log runs) for run-end-encoded columns. Slicing a+validity recounts its nulls with a popcount.+-}+sliceColumnArray :: Int -> Int -> ColumnArray -> ColumnArray+sliceColumnArray !start0 !len0 col =+  let !n = columnLength col+      !start = min n (max 0 start0)+      !len = max 0 (min len0 (n - start))+  in if len == n && start == 0+       then col+       else sliceRows start len col+++-- | Slice rows @[s, s + l)@; the caller has clamped the window.+sliceRows :: Int -> Int -> ColumnArray -> ColumnArray+sliceRows !s !l = \case+  I.ColNull _ -> I.ColNull l+  I.ColPrim t v xs -> withPrim t (I.ColPrim t (sv v) (VS.slice s l xs))+  I.ColBool v b -> I.ColBool (sv v) (sliceBitmap s l b)+  I.ColUtf8 v o d -> I.ColUtf8 (sv v) (VS.slice s (l + 1) o) d+  I.ColBinary v o d -> I.ColBinary (sv v) (VS.slice s (l + 1) o) d+  I.ColLargeUtf8 v o d -> I.ColLargeUtf8 (sv v) (VS.slice s (l + 1) o) d+  I.ColLargeBinary v o d -> I.ColLargeBinary (sv v) (VS.slice s (l + 1) o) d+  I.ColFixedSizeBinary w _ v d -> I.ColFixedSizeBinary w l (sv v) (BSU.unsafeDrop (s * w) d)+  I.ColUtf8View v views bufs -> I.ColUtf8View (sv v) (sliceViews views) bufs+  I.ColBinaryView v views bufs -> I.ColBinaryView (sv v) (sliceViews views) bufs+  I.ColStruct _ v cs -> I.ColStruct l (sv v) (V.map (fmap (sliceColumnArray s l)) cs)+  I.ColList v o c -> I.ColList (sv v) (VS.slice s (l + 1) o) c+  I.ColLargeList v o c -> I.ColLargeList (sv v) (VS.slice s (l + 1) o) c+  I.ColListView v o z c -> I.ColListView (sv v) (VS.slice s l o) (VS.slice s l z) c+  I.ColLargeListView v o z c -> I.ColLargeListView (sv v) (VS.slice s l o) (VS.slice s l z) c+  I.ColFixedSizeList w _ v c -> I.ColFixedSizeList w l (sv v) (sliceColumnArray (s * w) (l * w) c)+  I.ColMap v o k x -> I.ColMap (sv v) (VS.slice s (l + 1) o) k x+  I.ColDenseUnion t o cs -> I.ColDenseUnion (VS.slice s l t) (VS.slice s l o) cs+  I.ColSparseUnion t cs -> I.ColSparseUnion (VS.slice s l t) (V.map (sliceColumnArray s l) cs)+  I.ColDictionary did ix vals -> I.ColDictionary did (sliceColumnArray s l ix) vals+  I.ColRunEndEncoded off _ re vals+    | l == 0 -> I.ColRunEndEncoded 0 0 (sliceColumnArray 0 0 re) (sliceColumnArray 0 0 vals)+    | otherwise ->+        let !lo = off + s+            !i = physicalRun re lo+            !j = physicalRun re (lo + l - 1)+            !k = j - i + 1+        in I.ColRunEndEncoded lo l (sliceColumnArray i k re) (sliceColumnArray i k vals)+  where+    sv = sliceValidity s l+    sliceViews = BSU.unsafeTake (l * 16) . BSU.unsafeDrop (s * 16)+++{- | Index of the run holding logical position @p@ (already offset):+the first run whose end is greater than @p@. Binary search.+-}+physicalRun :: ColumnArray -> Int -> Int+physicalRun re p = go 0 (columnLength re)+  where+    go !lo !hi+      | lo >= hi = lo+      | otherwise =+          let !mid = (lo + hi) `unsafeShiftR` 1+          in if keyAt re mid > p then go lo mid else go (mid + 1) hi+++-- ============================================================+-- Detaching+-- ============================================================++{- | A copy that references only fresh memory: each buffer's logical+range is copied (offsets rebased to 0, bitmaps to bit offset 0), so+the result no longer keeps the decoded input alive. O(referenced bytes).+-}+copyColumn :: ColumnArray -> ColumnArray+copyColumn = \case+  I.ColNull n -> I.ColNull n+  I.ColPrim t v xs -> withPrim t (I.ColPrim t (copyValidity v) (copyStorable xs))+  I.ColBool v b -> I.ColBool (copyValidity v) (copyBitmap b)+  I.ColUtf8 v o d -> copyVar I.ColUtf8 v o d+  I.ColBinary v o d -> copyVar I.ColBinary v o d+  I.ColLargeUtf8 v o d -> copyVar I.ColLargeUtf8 v o d+  I.ColLargeBinary v o d -> copyVar I.ColLargeBinary v o d+  I.ColFixedSizeBinary w n v d -> I.ColFixedSizeBinary w n (copyValidity v) (freshBytes (BSU.unsafeTake (w * n) d))+  I.ColUtf8View v views bufs -> I.ColUtf8View (copyValidity v) (freshBytes views) (V.map freshBytes bufs)+  I.ColBinaryView v views bufs -> I.ColBinaryView (copyValidity v) (freshBytes views) (V.map freshBytes bufs)+  I.ColStruct n v cs -> I.ColStruct n (copyValidity v) (forceElems (V.map (\(nm, c) -> let !c' = copyColumn (sliceColumnArray 0 n c) in (nm, c')) cs))+  I.ColList v o c -> copyList I.ColList v o c+  I.ColLargeList v o c -> copyList I.ColLargeList v o c+  I.ColListView v o z c -> I.ColListView (copyValidity v) (copyStorable o) (copyStorable z) (copyColumn c)+  I.ColLargeListView v o z c -> I.ColLargeListView (copyValidity v) (copyStorable o) (copyStorable z) (copyColumn c)+  I.ColFixedSizeList w n v c -> I.ColFixedSizeList w n (copyValidity v) (copyColumn (sliceColumnArray 0 (w * n) c))+  I.ColMap v o k x ->+    let !(o', r) = rebaseToZero o+    in I.ColMap (copyValidity v) o' (copyColumn (sliceRange r k)) (copyColumn (sliceRange r x))+  I.ColDenseUnion t o cs -> I.ColDenseUnion (copyStorable t) (copyStorable o) (forceElems (V.map copyColumn cs))+  I.ColSparseUnion t cs -> I.ColSparseUnion (copyStorable t) (forceElems (V.map (copyColumn . sliceColumnArray 0 (VS.length t)) cs))+  I.ColDictionary did ix vals -> I.ColDictionary did (copyColumn ix) (copyColumn vals)+  I.ColRunEndEncoded off n re vals -> I.ColRunEndEncoded off n (copyColumn re) (copyColumn (sliceColumnArray 0 (columnLength re) vals))+  where+    copyVar :: Offset o => (Maybe Validity -> VS.Vector o -> ByteString -> ColumnArray) -> Maybe Validity -> VS.Vector o -> ByteString -> ColumnArray+    copyVar mk v o d =+      let !(o', ChildRange s l) = rebaseToZero o+      in mk (copyValidity v) o' (freshBytes (BSU.unsafeTake l (BSU.unsafeDrop s d)))+    copyList :: Offset o => (Maybe Validity -> VS.Vector o -> ColumnArray -> ColumnArray) -> Maybe Validity -> VS.Vector o -> ColumnArray -> ColumnArray+    copyList mk v o c =+      let !(o', r) = rebaseToZero o+      in mk (copyValidity v) o' (copyColumn (sliceRange r c))+++sliceRange :: ChildRange -> ColumnArray -> ColumnArray+sliceRange (ChildRange s l) = sliceColumnArray s l+++freshBytes :: ByteString -> ByteString+freshBytes bs = createAligned (BS.length bs) (copyBS bs)+++-- ============================================================+-- Concatenation+-- ============================================================++-- | @concatColumnArrays [a, b]@.+concatColumnArray :: ColumnArray -> ColumnArray -> Either String ColumnArray+concatColumnArray a b = concatColumnArrays [a, b]+++{- | Append columns of the same type (same tag, decimal parameters,+widths, struct field names, union arity, dictionary index type and+run-end type). Per buffer: one allocation of the summed size, a+memcpy per input, bit copies for bitmaps at any bit offset, and one+rebase pass per input for offsets. Var-length data and list children+keep only the referenced ranges. Dictionaries with identical values+(same buffers or logically equal) concatenate their keys; otherwise+the values are concatenated and later keys shifted, which must fit+the key width. Offsets that overflow their width are rejected.+-}+concatColumnArrays :: [ColumnArray] -> Either String ColumnArray+concatColumnArrays = \case+  [] -> Left "Arrow.Column.concatColumnArrays: no columns"+  [c] -> Right c+  cols@(c0 : _) -> case c0 of+    I.ColNull _ -> I.ColNull . sum <$> traverse (\case I.ColNull n -> Right n; c -> mismatch c) cols+    I.ColPrim t _ _ -> withPrim t $ do+      parts <- traverse (primPart t) cols+      Right (I.ColPrim t (concatValidity (map (\(v, xs) -> (VS.length xs, v)) parts)) (concatStorable (map snd parts)))+    I.ColBool {} -> do+      parts <- traverse (\case I.ColBool v b -> Right (v, b); c -> mismatch c) cols+      Right (I.ColBool (concatValidity (map (\(v, b) -> (bitmapLength b, v)) parts)) (concatBitmaps (map snd parts)))+    I.ColUtf8 {} -> traverse (\case I.ColUtf8 v o d -> Right (v, o, d); c -> mismatch c) cols >>= concatVar I.ColUtf8+    I.ColBinary {} -> traverse (\case I.ColBinary v o d -> Right (v, o, d); c -> mismatch c) cols >>= concatVar I.ColBinary+    I.ColLargeUtf8 {} -> traverse (\case I.ColLargeUtf8 v o d -> Right (v, o, d); c -> mismatch c) cols >>= concatVar I.ColLargeUtf8+    I.ColLargeBinary {} -> traverse (\case I.ColLargeBinary v o d -> Right (v, o, d); c -> mismatch c) cols >>= concatVar I.ColLargeBinary+    I.ColFixedSizeBinary w _ _ _ -> do+      parts <- traverse (\case I.ColFixedSizeBinary w' n v d | w' == w -> Right (n, v, d); c -> mismatch c) cols+      Right $+        I.ColFixedSizeBinary+          w+          (sum (map fst3 parts))+          (concatValidity (map (\(n, v, _) -> (n, v)) parts))+          (concatBytes (map (\(n, _, d) -> BSU.unsafeTake (n * w) d) parts))+    I.ColUtf8View {} -> traverse (\case I.ColUtf8View v vs bs -> Right (v, vs, bs); c -> mismatch c) cols >>= concatViews I.ColUtf8View+    I.ColBinaryView {} -> traverse (\case I.ColBinaryView v vs bs -> Right (v, vs, bs); c -> mismatch c) cols >>= concatViews I.ColBinaryView+    I.ColStruct _ _ cs0 -> do+      parts <- traverse (\case I.ColStruct n v cs | V.map fst cs == V.map fst cs0 -> Right (n, v, cs); c -> mismatch c) cols+      children <-+        V.generateM (V.length cs0) $ \j ->+          (,) (fst (V.unsafeIndex cs0 j))+            <$> concatColumnArrays (map (\(n, _, cs) -> sliceColumnArray 0 n (snd (V.unsafeIndex cs j))) parts)+      Right (I.ColStruct (sum (map fst3 parts)) (concatValidity (map (\(n, v, _) -> (n, v)) parts)) children)+    I.ColList {} -> traverse (\case I.ColList v o c -> Right (v, o, c); c -> mismatch c) cols >>= concatLists I.ColList+    I.ColLargeList {} -> traverse (\case I.ColLargeList v o c -> Right (v, o, c); c -> mismatch c) cols >>= concatLists I.ColLargeList+    I.ColListView {} -> traverse (\case I.ColListView v o z c -> Right (v, o, z, c); c -> mismatch c) cols >>= concatListViews I.ColListView+    I.ColLargeListView {} -> traverse (\case I.ColLargeListView v o z c -> Right (v, o, z, c); c -> mismatch c) cols >>= concatListViews I.ColLargeListView+    I.ColFixedSizeList w _ _ _ -> do+      parts <- traverse (\case I.ColFixedSizeList w' n v c | w' == w -> Right (n, v, c); c -> mismatch c) cols+      child <- concatColumnArrays (map (\(n, _, c) -> sliceColumnArray 0 (n * w) c) parts)+      Right (I.ColFixedSizeList w (sum (map fst3 parts)) (concatValidity (map (\(n, v, _) -> (n, v)) parts)) child)+    I.ColMap {} -> do+      parts <- traverse (\case I.ColMap v o k x -> Right (v, o, k, x); c -> mismatch c) cols+      (offs, ranges) <- concatOffsets (map (\(_, o, _, _) -> o) parts)+      keys <- concatColumnArrays (zipWith (\r (_, _, k, _) -> sliceRange r k) ranges parts)+      vals <- concatColumnArrays (zipWith (\r (_, _, _, x) -> sliceRange r x) ranges parts)+      Right (I.ColMap (concatValidity (map (\(v, o, _, _) -> (offsetRows o, v)) parts)) offs keys vals)+    I.ColDenseUnion _ _ cs0 -> do+      parts <- traverse (\case I.ColDenseUnion t o cs | V.length cs == V.length cs0 -> Right (t, o, cs); c -> mismatch c) cols+      let !k = V.length cs0+      children <- V.generateM k $ \j -> concatColumnArrays (map (\(_, _, cs) -> V.unsafeIndex cs j) parts)+      when (V.any (\c -> columnLength c > fromIntegral (maxBound :: Int32)) children) $+        Left "Arrow.Column.concatColumnArrays: dense union child exceeds 32-bit offsets"+      let bases = scanl (\acc (_, _, cs) -> zipWith (+) acc (map columnLength (V.toList cs))) (replicate k 0) parts+          shifted =+            zipWith+              ( \base (t, o, _) ->+                  let !bv = VS.fromList (map fromIntegral base) :: VS.Vector Int32+                  in VS.zipWith (\ti oi -> oi + VS.unsafeIndex bv (fromIntegral ti)) t o+              )+              bases+              parts+      Right (I.ColDenseUnion (concatStorable (map fst3 parts)) (concatStorable shifted) children)+    I.ColSparseUnion _ cs0 -> do+      parts <- traverse (\case I.ColSparseUnion t cs | V.length cs == V.length cs0 -> Right (t, cs); c -> mismatch c) cols+      children <-+        V.generateM (V.length cs0) $ \j ->+          concatColumnArrays (map (\(t, cs) -> sliceColumnArray 0 (VS.length t) (V.unsafeIndex cs j)) parts)+      Right (I.ColSparseUnion (concatStorable (map fst parts)) children)+    I.ColDictionary did ix0 vals0 -> do+      parts <- traverse (\case I.ColDictionary _ ix vals | sameShape ix ix0 -> Right (ix, vals); c -> mismatch c) cols+      if all (\(_, vals) -> identical vals vals0 || vals == vals0) parts+        then (\ix -> I.ColDictionary did ix vals0) <$> concatColumnArrays (map fst parts)+        else do+          vals <- concatColumnArrays (map snd parts)+          let bases = scanl (+) 0 (map (columnLength . snd) parts)+          shifted <- sequence (zipWith (\base (ix, _) -> shiftKeys base ix) bases parts)+          ix <- concatColumnArrays shifted+          Right (I.ColDictionary did ix vals)+    I.ColRunEndEncoded _ _ re0 _ -> do+      parts <- traverse (\case I.ColRunEndEncoded off n re vals | sameShape re re0 -> Right (off, n, re, vals); c -> mismatch c) cols+      concatRuns re0 parts+  where+    mismatch :: ColumnArray -> Either String b+    mismatch c = Left ("Arrow.Column.concatColumnArrays: incompatible columns " ++ columnTag c ++ " and the first column")+++fst3 :: (a, b, c) -> a+fst3 (a, _, _) = a+++primPart :: PrimType a -> ColumnArray -> Either String (Maybe Validity, VS.Vector a)+primPart t = \case+  I.ColPrim t' v xs | samePrimType t t', Just Refl <- samePrimTag t t' -> Right (v, xs)+  c -> Left ("Arrow.Column.concatColumnArrays: incompatible columns " ++ columnTag c ++ " and Col" ++ primTypeName t)+++-- | One aligned allocation holding the bytes in order.+concatBytes :: [ByteString] -> ByteString+concatBytes bss =+  createAligned (sum (map BS.length bss)) $ \dst ->+    let go !_ [] = pure ()+        go !off (b : rest) = copyBS b (dst `plusPtr` off) >> go (off + BS.length b) rest+    in go 0 bss+++concatStorable :: Storable a => [VS.Vector a] -> VS.Vector a+concatStorable = unsafeBytesToStorable . concatBytes . map storableToBytes+++{- | Offsets of consecutive list-like pieces as one offsets vector+starting at 0 (one rebase pass per piece), and each piece's child range.+-}+concatOffsets :: forall o. Offset o => [VS.Vector o] -> Either String (VS.Vector o, [ChildRange])+concatOffsets offss+  | toInteger total > toInteger (maxBound :: o) = Left "Arrow.Column.concatColumnArrays: offsets overflow their width"+  | otherwise = Right (unsafeBytesToStorable bytes, ranges)+  where+    ranges = map (\o -> ChildRange (fromIntegral (VS.unsafeHead o)) (fromIntegral (VS.unsafeLast o - VS.unsafeHead o))) offss+    !total = sum (map childLength ranges)+    !rows = sum (map offsetRows offss)+    !w = sizeOf (0 :: o)+    bytes = createAligned ((rows + 1) * w) $ \dst -> do+      pokeElemOff (castPtr dst) 0 (0 :: o)+      let go !_ !_ [] = pure ()+          go !pos !base (o : rest) = do+            let !m = offsetRows o+                !o0 = VS.unsafeHead o+            VS.unsafeWith o $ \src ->+              cRebaseOffsets (castPtr (dst `plusPtr` ((pos + 1) * w)) :: Ptr o) (src `plusPtr` w) m (fromIntegral base - fromIntegral o0)+            go (pos + m) (base + fromIntegral (VS.unsafeLast o - o0) :: Int) rest+      go 0 0 offss+++concatVar :: Offset o => (Maybe Validity -> VS.Vector o -> ByteString -> ColumnArray) -> [(Maybe Validity, VS.Vector o, ByteString)] -> Either String ColumnArray+concatVar mk parts = do+  (offs, ranges) <- concatOffsets (map (\(_, o, _) -> o) parts)+  let dat = concatBytes (zipWith (\(ChildRange s l) (_, _, d) -> BSU.unsafeTake l (BSU.unsafeDrop s d)) ranges parts)+  Right (mk (concatValidity (map (\(v, o, _) -> (offsetRows o, v)) parts)) offs dat)+++concatLists :: Offset o => (Maybe Validity -> VS.Vector o -> ColumnArray -> ColumnArray) -> [(Maybe Validity, VS.Vector o, ColumnArray)] -> Either String ColumnArray+concatLists mk parts = do+  (offs, ranges) <- concatOffsets (map (\(_, o, _) -> o) parts)+  child <- concatColumnArrays (zipWith (\r (_, _, c) -> sliceRange r c) ranges parts)+  Right (mk (concatValidity (map (\(v, o, _) -> (offsetRows o, v)) parts)) offs child)+++concatListViews+  :: forall o+   . Offset o+  => (Maybe Validity -> VS.Vector o -> VS.Vector o -> ColumnArray -> ColumnArray)+  -> [(Maybe Validity, VS.Vector o, VS.Vector o, ColumnArray)]+  -> Either String ColumnArray+concatListViews mk parts = do+  child <- concatColumnArrays (map (\(_, _, _, c) -> c) parts)+  when (toInteger (columnLength child) > toInteger (maxBound :: o)) $+    Left "Arrow.Column.concatColumnArrays: list-view child exceeds the offset width"+  let bases = scanl (+) 0 (map (\(_, _, _, c) -> columnLength c) parts)+      shifted = zipWith (\base (_, o, _, _) -> VS.map (+ fromIntegral base) o) bases parts+  Right $+    mk+      (concatValidity (map (\(v, o, _, _) -> (VS.length o, v)) parts))+      (concatStorable shifted)+      (concatStorable (map (\(_, _, z, _) -> z) parts))+      child+++-- | Views of later pieces point past the variadic buffers of earlier pieces.+concatViews :: (Maybe Validity -> ByteString -> V.Vector ByteString -> ColumnArray) -> [(Maybe Validity, ByteString, V.Vector ByteString)] -> Either String ColumnArray+concatViews mk parts =+  let bases = scanl (+) 0 (map (\(_, _, bs) -> V.length bs) parts)+      rows = map (\(_, vs, _) -> BS.length vs `quot` 16) parts+      views = createAligned (16 * sum rows) $ \dst ->+        let go !_ [] = pure ()+            go !pos ((base, (_, vs, _)) : rest) = do+              let !m = BS.length vs `quot` 16+              copyBS (BSU.unsafeTake (m * 16) vs) (dst `plusPtr` (pos * 16))+              when (base /= 0) $ forM_ [0 .. m - 1] $ \i -> do+                let !p = dst `plusPtr` ((pos + i) * 16)+                len <- peekByteOff p 0 :: IO Int32+                when (len > 12) $ do+                  bi <- peekByteOff p 8 :: IO Int32+                  pokeByteOff p 8 (bi + fromIntegral base)+              go (pos + m) rest+        in go 0 (zip bases parts)+  in Right (mk (concatValidity (zip rows (map (\(v, _, _) -> v) parts))) views (V.concat (map (\(_, _, bs) -> bs) parts)))+++-- | Add @base@ to every valid key (null slots become 0); 'Left' if a key would overflow its width.+shiftKeys :: Int -> ColumnArray -> Either String ColumnArray+shiftKeys base col+  | base == 0 = Right col+  | otherwise = case col of+      I.ColPrim t v xs | Just IntegralPrim <- integralPrim t ->+        let ok i x = not (unsafeIsValidAt v i) || toInteger x + toInteger base <= toInteger (maxBound `asTypeOf` x)+        in if VS.and (VS.imap ok xs)+             then Right (I.ColPrim t v (VS.imap (\i x -> if unsafeIsValidAt v i then x + fromIntegral base else 0) xs))+             else Left "Arrow.Column.concatColumnArrays: combined dictionary does not fit the key width"+      _ -> Left "Arrow.Column.concatColumnArrays: dictionary indices must be an integer column"+++-- | Run ends of a run-end column as a list of Ints.+runEndList :: ColumnArray -> [Int]+runEndList re = map (keyAt re) [0 .. columnLength re - 1]+++{- | The IPC layout of a run-end-encoded column: logical offset 0 and+run ends rebased so the last one is exactly the length (one pass over+the runs in the window). Any other column is returned unchanged.+-}+rebaseRunEnds :: ColumnArray -> Either String ColumnArray+rebaseRunEnds col = case col of+  I.ColRunEndEncoded off n re vals+    | off == 0 && (if runs == 0 then n == 0 else keyAt re (runs - 1) == n) && columnLength vals == runs -> Right col+    | otherwise -> concatRuns re [(off, n, re, vals)]+    where+      !runs = columnLength re+  _ -> Right col+++{- | Concatenate run-end-encoded pieces: each piece's runs are cut to+its logical window, rebased, and shifted past the earlier pieces.+-}+concatRuns :: ColumnArray -> [(Int, Int, ColumnArray, ColumnArray)] -> Either String ColumnArray+concatRuns re0 parts = case re0 of+  I.ColPrim t _ _ | Just IntegralPrim <- integralPrim t -> do+    let windows =+          map+            ( \(off, n, re, vals) ->+                if n == 0+                  then ([], sliceColumnArray 0 0 vals)+                  else+                    let !i = physicalRun re off+                        !j = physicalRun re (off + n - 1)+                        cut = map (\e -> min e (off + n) - off) (take (j - i + 1) (drop i (runEndList re)))+                    in (cut, sliceColumnArray i (j - i + 1) vals)+            )+            parts+        lens = map (\(_, n, _, _) -> n) parts+        bases = scanl (+) 0 lens+        ends = concat (zipWith (\b (es, _) -> map (+ b) es) bases windows)+        !total = sum lens+    when (toInteger total > toInteger (maxBound `asTypeOf` VS.head (vecOf t re0))) $+      Left "Arrow.Column.concatColumnArrays: run end overflows its integer width"+    vals <- concatColumnArrays (map snd windows)+    Right (I.ColRunEndEncoded 0 total (I.ColPrim t Nothing (VS.fromList (map fromIntegral ends))) vals)+  _ -> Left "Arrow.Column.concatColumnArrays: run ends must be an integer column"+  where+    vecOf :: PrimType a -> ColumnArray -> VS.Vector a+    vecOf t c = case c of+      I.ColPrim t' _ xs | Just Refl <- samePrimTag t t' -> xs+      _ -> withPrim t VS.empty+++-- ============================================================+-- Gathering+-- ============================================================++{- | Rows by index (repeats allowed); any index outside the column is+a 'Left'. Fixed-width buffers and bitmaps go through C gather+kernels; var-length columns take two passes (offsets, then bytes);+lists gather their child ranges; dictionaries, dense unions and list+views gather only their own buffers and share the rest.+-}+takeColumnArray :: VS.Vector Int -> ColumnArray -> Either String ColumnArray+takeColumnArray ix col+  | VS.any (\i -> i < 0 || i >= n) ix =+      Left ("Arrow.Column.takeColumnArray: index out of range for a column of " ++ show n ++ " rows")+  | otherwise = takeRows ix col+  where+    !n = columnLength col+++-- | 'takeColumnArray' after the range check.+takeRows :: VS.Vector Int -> ColumnArray -> Either String ColumnArray+takeRows ix col = case col of+  I.ColNull _ -> Right (I.ColNull k)+  I.ColPrim t v xs -> withPrim t (Right (I.ColPrim t (tv v) (gatherStorable ix xs)))+  I.ColBool v b -> Right (I.ColBool (tv v) (fst (takeBitmap ix b)))+  I.ColUtf8 v o d -> (\(o', d') -> I.ColUtf8 (tv v) o' d') <$> takeVar ix o d+  I.ColBinary v o d -> (\(o', d') -> I.ColBinary (tv v) o' d') <$> takeVar ix o d+  I.ColLargeUtf8 v o d -> (\(o', d') -> I.ColLargeUtf8 (tv v) o' d') <$> takeVar ix o d+  I.ColLargeBinary v o d -> (\(o', d') -> I.ColLargeBinary (tv v) o' d') <$> takeVar ix o d+  I.ColFixedSizeBinary w _ v d -> Right (I.ColFixedSizeBinary w k (tv v) (gatherBytes w ix d))+  I.ColUtf8View v views bufs -> Right (I.ColUtf8View (tv v) (gatherBytes 16 ix views) bufs)+  I.ColBinaryView v views bufs -> Right (I.ColBinaryView (tv v) (gatherBytes 16 ix views) bufs)+  I.ColStruct _ v cs -> I.ColStruct k (tv v) <$> V.mapM (traverse (takeRows ix)) cs+  I.ColList v o c -> takeList I.ColList v o c+  I.ColLargeList v o c -> takeList I.ColLargeList v o c+  I.ColListView v o z c -> Right (I.ColListView (tv v) (gatherStorable ix o) (gatherStorable ix z) c)+  I.ColLargeListView v o z c -> Right (I.ColLargeListView (tv v) (gatherStorable ix o) (gatherStorable ix z) c)+  I.ColFixedSizeList w _ v c ->+    I.ColFixedSizeList w k (tv v) <$> takeRows (VS.generate (k * w) (\j -> VS.unsafeIndex ix (j `quot` w) * w + j `rem` w)) c+  I.ColMap v o keys vals -> do+    (o', childIx) <- takeOffsetRanges ix o+    I.ColMap (tv v) o' <$> takeRows childIx keys <*> takeRows childIx vals+  I.ColDenseUnion t o cs -> Right (I.ColDenseUnion (gatherStorable ix t) (gatherStorable ix o) cs)+  I.ColSparseUnion t cs -> I.ColSparseUnion (gatherStorable ix t) <$> V.mapM (takeRows ix) cs+  I.ColDictionary did keys vals -> (\keys' -> I.ColDictionary did keys' vals) <$> takeRows ix keys+  I.ColRunEndEncoded off _ re vals -> takeRuns ix off re vals+  where+    !k = VS.length ix+    tv = takeValidity ix+    takeList :: Offset o => (Maybe Validity -> VS.Vector o -> ColumnArray -> ColumnArray) -> Maybe Validity -> VS.Vector o -> ColumnArray -> Either String ColumnArray+    takeList mk v o c = do+      (o', childIx) <- takeOffsetRanges ix o+      mk (tv v) o' <$> takeRows childIx c+++gatherStorable :: forall a. Storable a => VS.Vector Int -> VS.Vector a -> VS.Vector a+gatherStorable ix xs = unsafeBytesToStorable (gatherBytes (sizeOf (undefined :: a)) ix (storableToBytes xs))+++-- | Gather @w@-byte elements by index into one fresh buffer.+gatherBytes :: Int -> VS.Vector Int -> ByteString -> ByteString+gatherBytes w ix src =+  createAligned (n * w) $ \dst -> withBytesPtr src $ \ps -> VS.unsafeWith ix $ \pix -> case w of+    1 -> K.gather1 dst ps pix n+    2 -> K.gather2 dst ps pix n+    4 -> K.gather4 dst ps pix n+    8 -> K.gather8 dst ps pix n+    16 -> K.gather16 dst ps pix n+    _ -> forM_ [0 .. n - 1] $ \j -> copyBytes (dst `plusPtr` (j * w)) (ps `plusPtr` (VS.unsafeIndex ix j * w)) w+  where+    !n = VS.length ix+++takeVar :: forall o. Offset o => VS.Vector Int -> VS.Vector o -> ByteString -> Either String (VS.Vector o, ByteString)+takeVar ix offs dat = unsafeDupablePerformIO $ do+  let !n = VS.length ix+  ofp <- mallocAligned ((n + 1) * sizeOf (0 :: o))+  total <- unsafeWithForeignPtr ofp $ \po -> VS.unsafeWith offs $ \ps -> VS.unsafeWith ix $ \pix ->+    cTakeOffsets (castPtr po) ps pix n+  if total < 0+    then pure (Left "Arrow.Column.takeColumnArray: gathered data exceeds the offset width")+    else do+      let bytes = createAligned total $ \dst -> withBytesPtr dat $ \pd -> VS.unsafeWith offs $ \ps -> VS.unsafeWith ix $ \pix ->+            cTakeBytes dst ps pd pix n+      pure (Right (VS.unsafeFromForeignPtr0 (castForeignPtr ofp) (n + 1), bytes))+++-- | New offsets (from 0) for the selected rows, and the child indices they cover.+takeOffsetRanges :: forall o. Offset o => VS.Vector Int -> VS.Vector o -> Either String (VS.Vector o, VS.Vector Int)+takeOffsetRanges ix offs = unsafeDupablePerformIO $ do+  let !n = VS.length ix+  ofp <- mallocAligned ((n + 1) * sizeOf (0 :: o))+  total <- unsafeWithForeignPtr ofp $ \po -> VS.unsafeWith offs $ \ps -> VS.unsafeWith ix $ \pix ->+    cTakeOffsets (castPtr po) ps pix n+  if total < 0+    then pure (Left "Arrow.Column.takeColumnArray: gathered child exceeds the offset width")+    else do+      let childIx = VS.create $ do+            mv <- VSM.unsafeNew total+            let go !j !pos+                  | j >= n = pure ()+                  | otherwise = do+                      let !r = VS.unsafeIndex ix j+                          !s = fromIntegral (VS.unsafeIndex offs r) :: Int+                          !e = fromIntegral (VS.unsafeIndex offs (r + 1)) :: Int+                      forM_ [0 .. e - s - 1] $ \q -> VSM.unsafeWrite mv (pos + q) (s + q)+                      go (j + 1) (pos + e - s)+            go 0 0+            pure mv+      pure (Right (VS.unsafeFromForeignPtr0 (castForeignPtr ofp) (n + 1), childIx))+++-- | Gather a run-end-encoded column: one run per change of physical run.+takeRuns :: VS.Vector Int -> Int -> ColumnArray -> ColumnArray -> Either String ColumnArray+takeRuns ix off re vals = case re of+  I.ColPrim t _ xs | Just IntegralPrim <- integralPrim t -> do+    let phys = map (\i -> physicalRun re (off + i)) (VS.toList ix)+        grouped = groupRuns phys+        ends = scanl1 (+) (map snd grouped)+        !k = VS.length ix+    when (toInteger k > toInteger (maxBound `asTypeOf` VS.head xs)) $+      Left "Arrow.Column.takeColumnArray: run end overflows its integer width"+    vals' <- takeRows (VS.fromList (map fst grouped)) vals+    Right (I.ColRunEndEncoded 0 k (I.ColPrim t Nothing (VS.fromList (map fromIntegral ends))) vals')+  _ -> Left "Arrow.Column.takeColumnArray: run ends must be an integer column"+  where+    groupRuns :: [Int] -> [(Int, Int)]+    groupRuns = foldr step []+    step p ((q, c) : rest) | p == q = (q, c + 1) : rest+    step p acc = (p, 1) : acc+++-- ============================================================+-- Dictionaries+-- ============================================================++{- | Replace a dictionary column by the values its keys select, with+the key validity applied; any other column is returned unchanged. A+column whose rows are all null may reference an empty dictionary; it+expands to all-null rows of the value type.+-}+expandDictionary :: ColumnArray -> Either String ColumnArray+expandDictionary = \case+  I.ColDictionary _ keys vals -> do+    let !n = columnLength keys+        !kv = validity keys+    dense <-+      if columnLength vals == 0+        then+          if nullCount keys == n+            then Right (fillerColumn n vals)+            else Left "Arrow.Column.expandDictionary: dictionary key out of range for an empty dictionary"+        else takeColumnArray (VS.generate n (\i -> if unsafeIsValidAt kv i then keyAt keys i else 0)) vals+    maskValidity kv dense+  col -> Right col+++{- | Replace the values of every dictionary column (at any depth) with+the dictionary registered for its id, checking every valid key against+it (C kernel). A column with no valid key keeps its placeholder when+the id is unknown; otherwise an unknown id is a 'Left'. Values the+lookup returns are used as they are (dictionaries nested in them must+already be resolved).+-}+resolveDictionaryColumn :: (Int64 -> Maybe ColumnArray) -> ColumnArray -> Either String ColumnArray+resolveDictionaryColumn lookupVals = go+  where+    go col = case col of+      I.ColDictionary did keys _ -> case lookupVals did of+        Nothing+          | nullCount keys == columnLength keys -> Right col+          | otherwise -> Left ("Arrow.Column: no dictionary batch for dictionary id " ++ show did)+        Just vals -> do+          validateKeys ("Arrow.Column: dictionary id " ++ show did) keys (columnLength vals)+          Right (I.ColDictionary did keys vals)+      I.ColStruct n v cs -> I.ColStruct n v <$> V.mapM (traverse go) cs+      I.ColList v o c -> I.ColList v o <$> go c+      I.ColLargeList v o c -> I.ColLargeList v o <$> go c+      I.ColListView v o z c -> I.ColListView v o z <$> go c+      I.ColLargeListView v o z c -> I.ColLargeListView v o z <$> go c+      I.ColFixedSizeList w n v c -> I.ColFixedSizeList w n v <$> go c+      I.ColMap v o k x -> I.ColMap v o <$> go k <*> go x+      I.ColDenseUnion t o cs -> I.ColDenseUnion t o <$> V.mapM go cs+      I.ColSparseUnion t cs -> I.ColSparseUnion t <$> V.mapM go cs+      I.ColRunEndEncoded off n re vals -> I.ColRunEndEncoded off n re <$> go vals+      _ -> Right col+++-- ============================================================+-- Nullability+-- ============================================================++{- | Null out the rows whose mask bit is clear (rows already null stay+null). The mask must have the column's length. Columns without a+validity slot (unions, run-end-encoded) accept only 'Nothing'; a+'ColNull' is unchanged.+-}+maskValidity :: Maybe Validity -> ColumnArray -> Either String ColumnArray+maskValidity Nothing col = Right col+maskValidity m@(Just mv) col+  | bitmapLength (validityBits mv) /= columnLength col =+      Left "Arrow.Column.maskValidity: mask length differs from the column length"+  | otherwise = case col of+      I.ColNull _ -> Right col+      I.ColPrim t v xs -> Right (I.ColPrim t (a v) xs)+      I.ColBool v b -> Right (I.ColBool (a v) b)+      I.ColUtf8 v o d -> Right (I.ColUtf8 (a v) o d)+      I.ColBinary v o d -> Right (I.ColBinary (a v) o d)+      I.ColLargeUtf8 v o d -> Right (I.ColLargeUtf8 (a v) o d)+      I.ColLargeBinary v o d -> Right (I.ColLargeBinary (a v) o d)+      I.ColFixedSizeBinary w n v d -> Right (I.ColFixedSizeBinary w n (a v) d)+      I.ColUtf8View v views bufs -> Right (I.ColUtf8View (a v) views bufs)+      I.ColBinaryView v views bufs -> Right (I.ColBinaryView (a v) views bufs)+      I.ColStruct n v cs -> Right (I.ColStruct n (a v) cs)+      I.ColList v o c -> Right (I.ColList (a v) o c)+      I.ColLargeList v o c -> Right (I.ColLargeList (a v) o c)+      I.ColListView v o z c -> Right (I.ColListView (a v) o z c)+      I.ColLargeListView v o z c -> Right (I.ColLargeListView (a v) o z c)+      I.ColFixedSizeList w n v c -> Right (I.ColFixedSizeList w n (a v) c)+      I.ColMap v o k x -> Right (I.ColMap (a v) o k x)+      I.ColDictionary did keys vals -> (\keys' -> I.ColDictionary did keys' vals) <$> maskValidity m keys+      I.ColDenseUnion {} -> noSlot+      I.ColSparseUnion {} -> noSlot+      I.ColRunEndEncoded {} -> noSlot+  where+    a = andValidity m+    noSlot = Left ("Arrow.Column.maskValidity: " ++ columnTag col ++ " has no validity bitmap")+++{- | Every column with a validity slot can hold nulls, so this is the+identity on those (and on 'ColNull'); unions and run-end-encoded+columns are rejected because they have no validity of their own.+-}+toNullableColumn :: ColumnArray -> Either String ColumnArray+toNullableColumn col = case col of+  I.ColDenseUnion {} -> noSlot+  I.ColSparseUnion {} -> noSlot+  I.ColRunEndEncoded {} -> noSlot+  _ -> Right col+  where+    noSlot = Left ("Arrow.Column: " ++ columnTag col ++ " has no validity bitmap and cannot be made nullable")+++-- ============================================================+-- Map invariants+-- ============================================================++{- | Check that a 'ColMap' column satisfies the @keysSorted@ promise of+its 'AMap' field: within every non-null entry the keys are non-null+and non-decreasing.++Keys compare by value: integers, dates, times, timestamps and+durations numerically; floating point (half floats included)+numerically with @-0 == 0@ and every NaN equal to every other NaN+and greater than any number; booleans with @False < True@; strings by+code point; binary (fixed-size and views included) as unsigned bytes;+decimals as signed integers; dictionary keys by their resolved+values. Interval, nested, union, run-end-encoded and null key columns+have no order and are rejected, as is any column that is not a map.+-}+validateMapKeysSorted :: ColumnArray -> Either String ()+validateMapKeysSorted = \case+  I.ColMap v offs keys _ -> do+    (cmp, isNull) <- keyOrder keys+    let !nk = columnLength keys+        !nEntries = offsetRows offs+        entry !i+          | i >= nEntries = Right ()+          | otherwise = do+              let !s = fromIntegral (VS.unsafeIndex offs i) :: Int+                  !e = fromIntegral (VS.unsafeIndex offs (i + 1)) :: Int+              if s < 0 || e < s || e > nk+                then Left ("Arrow.Column.validateMapKeysSorted: entry " ++ show i ++ " has offsets outside the key column")+                else+                  if unsafeIsValidAt v i+                    then keysIn i s e s >> entry (i + 1)+                    else entry (i + 1)+        keysIn i s e !j+          | j >= e = Right ()+          | isNull j = Left ("Arrow.Column.validateMapKeysSorted: entry " ++ show i ++ " has a null key")+          | j > s && cmp (j - 1) j == GT =+              Left ("Arrow.Column.validateMapKeysSorted: entry " ++ show i ++ " key " ++ show (j - s) ++ " is smaller than the key before it")+          | otherwise = keysIn i s e (j + 1)+    entry 0+  c -> Left ("Arrow.Column.validateMapKeysSorted: expected a map column, got " ++ columnTag c)+++-- | Row comparator and null test for an orderable key column.+keyOrder :: ColumnArray -> Either String (Int -> Int -> Ordering, Int -> Bool)+keyOrder col = case col of+  I.ColPrim t v xs -> case t of+    PInt8 -> ordered v xs+    PInt16 -> ordered v xs+    PInt32 -> ordered v xs+    PInt64 -> ordered v xs+    PUInt8 -> ordered v xs+    PUInt16 -> ordered v xs+    PUInt32 -> ordered v xs+    PUInt64 -> ordered v xs+    PDate32 -> ordered v xs+    PDate64 -> ordered v xs+    PTime32 -> ordered v xs+    PTime64 -> ordered v xs+    PTimestamp -> ordered v xs+    PDuration -> ordered v xs+    PDecimal32 _ _ -> ordered v xs+    PDecimal64 _ _ -> ordered v xs+    PDecimal128 _ _ -> ordered v xs+    PDecimal256 _ _ -> ordered v xs+    PFloat16 -> by v (\i j -> compareFloating (float16ToDouble (VS.unsafeIndex xs i)) (float16ToDouble (VS.unsafeIndex xs j)))+    PFloat -> by v (\i j -> compareFloating (VS.unsafeIndex xs i) (VS.unsafeIndex xs j))+    PDouble -> by v (\i j -> compareFloating (VS.unsafeIndex xs i) (VS.unsafeIndex xs j))+    _ -> unordered+  I.ColBool v b -> by v (\i j -> compare (unsafeBitAt b i) (unsafeBitAt b j))+  I.ColUtf8 {} -> bytes+  I.ColBinary {} -> bytes+  I.ColLargeUtf8 {} -> bytes+  I.ColLargeBinary {} -> bytes+  I.ColUtf8View {} -> bytes+  I.ColBinaryView {} -> bytes+  I.ColFixedSizeBinary {} -> bytes+  I.ColDictionary _ keys vals -> do+    (cmp, isNull) <- keyOrder vals+    let !nv = columnLength vals+        kv = validity keys+    if all (\i -> not (unsafeIsValidAt kv i) || (keyAt keys i >= 0 && keyAt keys i < nv)) [0 .. columnLength keys - 1]+      then+        Right+          ( \i j -> cmp (keyAt keys i) (keyAt keys j)+          , \i -> not (unsafeIsValidAt kv i) || isNull (keyAt keys i)+          )+      else Left "Arrow.Column.validateMapKeysSorted: dictionary key index outside its dictionary (is the dictionary resolved?)"+  _ -> unordered+  where+    unordered = Left ("Arrow.Column.validateMapKeysSorted: map keys of type " ++ columnTag col ++ " have no defined order")+    v0 = validity col+    by v cmp = Right (cmp, \i -> not (unsafeIsValidAt v i))+    ordered :: (Storable a, Ord a) => Maybe Validity -> VS.Vector a -> Either String (Int -> Int -> Ordering, Int -> Bool)+    ordered v xs = by v (\i j -> compare (VS.unsafeIndex xs i) (VS.unsafeIndex xs j))+    bytes = by v0 (\i j -> compare (rowBytes col i) (rowBytes col j))+++-- | Bytes of row @i@ of a byte-like column, ignoring validity.+rowBytes :: ColumnArray -> Int -> ByteString+rowBytes c i = case c of+  I.ColUtf8 _ o d -> varSlice o d i+  I.ColBinary _ o d -> varSlice o d i+  I.ColLargeUtf8 _ o d -> varSlice o d i+  I.ColLargeBinary _ o d -> varSlice o d i+  I.ColUtf8View _ views bufs -> viewAt views bufs i+  I.ColBinaryView _ views bufs -> viewAt views bufs i+  I.ColFixedSizeBinary w _ _ d -> BSU.unsafeTake w (BSU.unsafeDrop (i * w) d)+  _ -> BS.empty+++-- | Numeric order with @-0 == 0@; NaNs equal to each other and above every number.+compareFloating :: RealFloat a => a -> a -> Ordering+compareFloating x y = case (isNaN x, isNaN y) of+  (True, True) -> EQ+  (True, False) -> GT+  (False, True) -> LT+  (False, False) -> compare x y+++-- ============================================================+-- Eq, Show, NFData+-- ============================================================++{- | Structural, O(columns): every child column and validity is forced,+recursively. A column stores no lazy row thunks, so this is the whole+normal form; row-level behaviour needs the accessors.+-}+instance NFData ColumnArray where+  rnf = \case+    I.ColNull n -> n `seq` ()+    I.ColPrim t v xs -> withPrim t (rnf v `seq` xs `seq` ())+    I.ColBool v b -> rnf v `seq` rnf b+    I.ColUtf8 v o d -> rnf v `seq` rnf o `seq` rnf d+    I.ColBinary v o d -> rnf v `seq` rnf o `seq` rnf d+    I.ColLargeUtf8 v o d -> rnf v `seq` rnf o `seq` rnf d+    I.ColLargeBinary v o d -> rnf v `seq` rnf o `seq` rnf d+    I.ColFixedSizeBinary _ _ v d -> rnf v `seq` rnf d+    I.ColUtf8View v views bufs -> rnf v `seq` rnf views `seq` rnf bufs+    I.ColBinaryView v views bufs -> rnf v `seq` rnf views `seq` rnf bufs+    I.ColStruct _ v cs -> rnf v `seq` V.foldl' (\() (n, c) -> rnf n `seq` rnf c) () cs+    I.ColList v _ c -> rnf v `seq` rnf c+    I.ColLargeList v _ c -> rnf v `seq` rnf c+    I.ColListView v _ _ c -> rnf v `seq` rnf c+    I.ColLargeListView v _ _ c -> rnf v `seq` rnf c+    I.ColFixedSizeList _ _ v c -> rnf v `seq` rnf c+    I.ColMap v _ k x -> rnf v `seq` rnf k `seq` rnf x+    I.ColDenseUnion _ _ cs -> V.foldl' (\() c -> rnf c) () cs+    I.ColSparseUnion _ cs -> V.foldl' (\() c -> rnf c) () cs+    I.ColDictionary _ k x -> rnf k `seq` rnf x+    I.ColRunEndEncoded _ _ re vals -> rnf re `seq` rnf vals+++-- | Logical, O(n); see the module header.+instance Eq ColumnArray where+  a == b =+    sameShape a b+      && columnLength a == columnLength b+      && (identical a b || rowsEqual a b)+++-- | Same type: tag, parameters, widths, field names, arity, recursively.+sameShape :: ColumnArray -> ColumnArray -> Bool+sameShape a b = case (a, b) of+  (I.ColNull _, I.ColNull _) -> True+  (I.ColPrim t _ _, I.ColPrim t' _ _) -> samePrimType t t'+  (I.ColBool {}, I.ColBool {}) -> True+  (I.ColUtf8 {}, I.ColUtf8 {}) -> True+  (I.ColBinary {}, I.ColBinary {}) -> True+  (I.ColLargeUtf8 {}, I.ColLargeUtf8 {}) -> True+  (I.ColLargeBinary {}, I.ColLargeBinary {}) -> True+  (I.ColFixedSizeBinary w _ _ _, I.ColFixedSizeBinary w' _ _ _) -> w == w'+  (I.ColUtf8View {}, I.ColUtf8View {}) -> True+  (I.ColBinaryView {}, I.ColBinaryView {}) -> True+  (I.ColStruct _ _ cs, I.ColStruct _ _ cs') ->+    V.length cs == V.length cs' && V.and (V.zipWith (\(n, x) (n', y) -> n == n' && sameShape x y) cs cs')+  (I.ColList _ _ c, I.ColList _ _ c') -> sameShape c c'+  (I.ColLargeList _ _ c, I.ColLargeList _ _ c') -> sameShape c c'+  (I.ColListView _ _ _ c, I.ColListView _ _ _ c') -> sameShape c c'+  (I.ColLargeListView _ _ _ c, I.ColLargeListView _ _ _ c') -> sameShape c c'+  (I.ColFixedSizeList w _ _ c, I.ColFixedSizeList w' _ _ c') -> w == w' && sameShape c c'+  (I.ColMap _ _ k x, I.ColMap _ _ k' x') -> sameShape k k' && sameShape x x'+  (I.ColDenseUnion _ _ cs, I.ColDenseUnion _ _ cs') -> sameChildren cs cs'+  (I.ColSparseUnion _ cs, I.ColSparseUnion _ cs') -> sameChildren cs cs'+  (I.ColDictionary did k x, I.ColDictionary did' k' x') -> did == did' && sameShape k k' && sameShape x x'+  (I.ColRunEndEncoded _ _ r x, I.ColRunEndEncoded _ _ r' x') -> sameShape r r' && sameShape x x'+  _ -> False+  where+    sameChildren cs cs' = V.length cs == V.length cs' && V.and (V.zipWith sameShape cs cs')+++-- | Buffer identity of flat columns: same memory, same window.+identical :: ColumnArray -> ColumnArray -> Bool+identical a b = case (a, b) of+  (I.ColPrim t v xs, I.ColPrim t' v' ys) ->+    samePrimType t t' && sameV v v' && withPrim t (withPrim t' (sameBS (storableToBytes xs) (storableToBytes ys)))+  (I.ColBool v x, I.ColBool v' y) -> sameV v v' && sameBitmap x y+  (I.ColUtf8 v o d, I.ColUtf8 v' o' d') -> sameV v v' && sameBS (storableToBytes o) (storableToBytes o') && sameBS d d'+  (I.ColBinary v o d, I.ColBinary v' o' d') -> sameV v v' && sameBS (storableToBytes o) (storableToBytes o') && sameBS d d'+  (I.ColLargeUtf8 v o d, I.ColLargeUtf8 v' o' d') -> sameV v v' && sameBS (storableToBytes o) (storableToBytes o') && sameBS d d'+  (I.ColLargeBinary v o d, I.ColLargeBinary v' o' d') -> sameV v v' && sameBS (storableToBytes o) (storableToBytes o') && sameBS d d'+  _ -> False+  where+    sameBS (BSI.BS p l) (BSI.BS p' l') = p == p' && l == l'+    sameBitmap (Bitmap x o l) (Bitmap y o' l') = sameBS x y && o == o' && l == l'+    sameV Nothing Nothing = True+    sameV (Just (Validity x _)) (Just (Validity y _)) = sameBitmap x y+    sameV _ _ = False+++-- | Every row equal (shapes and lengths already agree).+rowsEqual :: ColumnArray -> ColumnArray -> Bool+rowsEqual a b = case (a, b) of+  (I.ColPrim t Nothing xs, I.ColPrim t' Nothing ys) ->+    withPrim t (withPrim t' (storableToBytes xs == storableToBytes ys))+  _ -> all (\i -> rowEq a i b i) [0 .. columnLength a - 1]+++-- | Whether row @i@ is logically valid (unions and run-end columns defer to their child rows).+rowValid :: ColumnArray -> Int -> Bool+rowValid c i = case c of+  I.ColNull _ -> False+  -- A valid key that selects a null value is a null row (arrow-rs logical nulls).+  I.ColDictionary _ k x ->+    unsafeIsValidAt (validity k) i+      && let !ki = keyAt k i in ki < 0 || ki >= columnLength x || rowValid x ki+  _ -> unsafeIsValidAt (validity c) i+++-- | Row @i@ of @a@ equals row @j@ of @b@ (same shape).+rowEq :: ColumnArray -> Int -> ColumnArray -> Int -> Bool+rowEq a i b j =+  let !va = rowValid a i+      !vb = rowValid b j+  in if va /= vb then False else not va || valueEq a i b j+++valueEq :: ColumnArray -> Int -> ColumnArray -> Int -> Bool+valueEq a i b j = case (a, b) of+  (I.ColPrim t _ xs, I.ColPrim t' _ ys) ->+    let !w = primWidth t+    in withPrim t (withPrim t' (slot w (storableToBytes xs) i == slot w (storableToBytes ys) j))+  (I.ColBool _ x, I.ColBool _ y) -> unsafeBitAt x i == unsafeBitAt y j+  (I.ColStruct _ _ cs, I.ColStruct _ _ cs') -> V.and (V.zipWith (\(_, x) (_, y) -> rowEq x i y j) cs cs')+  (I.ColList _ o c, I.ColList _ o' c') -> rangesEq c (offsetRange o i) c' (offsetRange o' j)+  (I.ColLargeList _ o c, I.ColLargeList _ o' c') -> rangesEq c (offsetRange o i) c' (offsetRange o' j)+  (I.ColListView {}, I.ColListView {}) -> viaRanges+  (I.ColLargeListView {}, I.ColLargeListView {}) -> viaRanges+  (I.ColFixedSizeList w _ _ c, I.ColFixedSizeList _ _ _ c') -> rangesEq c (ChildRange (i * w) w) c' (ChildRange (j * w) w)+  (I.ColMap _ o k x, I.ColMap _ o' k' x') ->+    let r = offsetRange o i+        r' = offsetRange o' j+    in rangesEq k r k' r' && rangesEq x r x' r'+  (I.ColDenseUnion t o cs, I.ColDenseUnion t' o' cs') ->+    let !ci = fromIntegral (VS.unsafeIndex t i)+    in ci == (fromIntegral (VS.unsafeIndex t' j) :: Int)+         && rowEq (V.unsafeIndex cs ci) (fromIntegral (VS.unsafeIndex o i)) (V.unsafeIndex cs' ci) (fromIntegral (VS.unsafeIndex o' j))+  (I.ColSparseUnion t cs, I.ColSparseUnion t' cs') ->+    let !ci = fromIntegral (VS.unsafeIndex t i)+    in ci == (fromIntegral (VS.unsafeIndex t' j) :: Int) && rowEq (V.unsafeIndex cs ci) i (V.unsafeIndex cs' ci) j+  (I.ColDictionary _ k x, I.ColDictionary _ k' x') ->+    let !ki = keyAt k i+        !kj = keyAt k' j+    in if ki >= 0 && ki < columnLength x && kj >= 0 && kj < columnLength x'+         then rowEq x ki x' kj+         else ki == kj+  (I.ColRunEndEncoded off _ r x, I.ColRunEndEncoded off' _ r' x') ->+    rowEq x (physicalRun r (off + i)) x' (physicalRun r' (off' + j))+  _ -> rowBytes a i == rowBytes b j+  where+    slot w bs k = BSU.unsafeTake w (BSU.unsafeDrop (k * w) bs)+    viaRanges = case (listRange a i, listRange b j) of+      (Just r, Just r') -> rangesEq (listChild a) r (listChild b) r'+      _ -> False+++listChild :: ColumnArray -> ColumnArray+listChild = \case+  I.ColListView _ _ _ c -> c+  I.ColLargeListView _ _ _ c -> c+  c -> c+++rangesEq :: ColumnArray -> ChildRange -> ColumnArray -> ChildRange -> Bool+rangesEq c (ChildRange s l) c' (ChildRange s' l') =+  l == l' && all (\q -> rowEq c (s + q) c' (s' + q)) [0 .. l - 1]+++-- | @Tag [row, row, ...]@ with nulls as @null@; parameters after the tag.+instance Show ColumnArray where+  showsPrec d c =+    showParen (d > 10) $+      showString (columnTag c) . params . showChar ' ' . showRows c+    where+      params = case c of+        I.ColPrim (PDecimal32 p s) _ _ -> showChar ' ' . shows p . showChar ' ' . shows s+        I.ColPrim (PDecimal64 p s) _ _ -> showChar ' ' . shows p . showChar ' ' . shows s+        I.ColPrim (PDecimal128 p s) _ _ -> showChar ' ' . shows p . showChar ' ' . shows s+        I.ColPrim (PDecimal256 p s) _ _ -> showChar ' ' . shows p . showChar ' ' . shows s+        I.ColFixedSizeBinary w _ _ _ -> showChar ' ' . shows w+        I.ColFixedSizeList w _ _ _ -> showChar ' ' . shows w+        I.ColDictionary did _ _ -> showChar ' ' . shows did+        _ -> id+++showRows :: ColumnArray -> ShowS+showRows c = showChar '[' . foldr (.) id (intersperse (showChar ',') (map (showRowAt c) [0 .. columnLength c - 1])) . showChar ']'+++showRowAt :: ColumnArray -> Int -> ShowS+showRowAt c i+  | not (rowValid c i) = showString "null"+  | otherwise = case c of+      I.ColPrim t _ xs -> withPrim t (shows (VS.unsafeIndex xs i))+      I.ColBool _ b -> shows (unsafeBitAt b i)+      I.ColUtf8 {} -> shows (utf8ToText (rowBytes c i))+      I.ColLargeUtf8 {} -> shows (utf8ToText (rowBytes c i))+      I.ColUtf8View {} -> shows (utf8ToText (rowBytes c i))+      I.ColStruct _ _ cs ->+        showChar '{'+          . foldr (.) id (intersperse (showString ", ") (map (\(n, x) -> showString (T.unpack n) . showString ": " . showRowAt x i) (V.toList cs)))+          . showChar '}'+      I.ColMap _ o k x ->+        let ChildRange s l = offsetRange o i+        in showChar '{'+             . foldr (.) id (intersperse (showString ", ") (map (\q -> showRowAt k (s + q) . showString " => " . showRowAt x (s + q)) [0 .. l - 1]))+             . showChar '}'+      I.ColDenseUnion t o cs ->+        let !ci = fromIntegral (VS.unsafeIndex t i)+        in showChar '<' . shows ci . showString ": " . showRowAt (V.unsafeIndex cs ci) (fromIntegral (VS.unsafeIndex o i)) . showChar '>'+      I.ColSparseUnion t cs ->+        let !ci = fromIntegral (VS.unsafeIndex t i)+        in showChar '<' . shows ci . showString ": " . showRowAt (V.unsafeIndex cs ci) i . showChar '>'+      I.ColDictionary _ k x ->+        let !ki = keyAt k i+        in if ki >= 0 && ki < columnLength x then showRowAt x ki else showChar '#' . shows ki+      I.ColRunEndEncoded off _ r x -> showRowAt x (physicalRun r (off + i))+      _ -> case listRange c i of+        Just (ChildRange s l) ->+          let child = case c of+                I.ColList _ _ x -> x+                I.ColLargeList _ _ x -> x+                I.ColFixedSizeList _ _ _ x -> x+                x -> listChild x+          in showChar '[' . foldr (.) id (intersperse (showChar ',') (map (\q -> showRowAt child (s + q)) [0 .. l - 1])) . showChar ']'+        Nothing -> shows (rowBytes c i)+++-- | Evaluate every element of a boxed vector to WHNF.+forceElems :: V.Vector a -> V.Vector a+forceElems v = V.foldl' (flip seq) () v `seq` v
+ src/Arrow/Column/Buffer.hs view
@@ -0,0 +1,660 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-}++{- | Low-level Arrow buffers: LSB-first bitmaps, validity, the element+types that have no Haskell primitive, and pinned 64-byte aligned+allocation.++This module exports raw constructors. 'Bitmap' and 'Validity' carry+invariants (documented on each type) that the accessors in+"Arrow.Column" rely on; code outside the reader, the writer, the+builders and the smart constructors should build them through+'mkBitmap', 'bitmapGenerate', 'mkValidity' and friends, which+"Arrow.Column" re-exports with the types kept abstract.+-}+module Arrow.Column.Buffer (+  -- * Bitmaps+  Bitmap (..),+  bitmapBytes,+  bitmapOffset,+  bitmapLength,+  mkBitmap,+  emptyBitmap,+  bitAt,+  unsafeBitAt,+  sliceBitmap,+  bitmapSetCount,+  bitmapGenerate,+  bitmapFromBools,+  bitmapToBools,+  copyBitmap,+  concatBitmaps,+  takeBitmap,++  -- * Validity+  Validity (..),+  validityBits,+  validityNullCount,+  mkValidity,+  validityGenerate,+  validityFromBools,+  checkValidity,+  isValidAt,+  unsafeIsValidAt,+  sliceValidity,+  concatValidity,+  takeValidity,+  andValidity,+  copyValidity,++  -- * Element types+  Float16 (..),+  float16ToDouble,+  IntervalDayTime (..),+  IntervalMonthDayNano (..),+  Decimal128 (..),+  decimal128ToInteger,+  decimal128FromInteger,+  Decimal256 (..),+  decimal256ToInteger,+  decimal256FromInteger,++  -- * Pinned buffers+  mallocAligned,+  createAligned,+  storableToBytes,+  bytesToStorable,+  unsafeBytesToStorable,+  copyStorable,+  swapCopy,+  swapViewCopy,+  withBytesPtr,+  withStorablePtr,+) where++import Columnar.SIMD (andBits, copyBits, gatherBits, popCountBits)+import Control.DeepSeq (NFData (..))+import Data.Bits (complement, shiftL, shiftR, unsafeShiftL, unsafeShiftR, (.&.), (.|.))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Internal qualified as BSI+import Data.ByteString.Unsafe qualified as BSU+import Data.Int (Int32, Int64)+import Data.Vector qualified as V+import Data.Vector.Storable qualified as VS+import Data.Word (Word16, Word64, Word8)+import Foreign.ForeignPtr (ForeignPtr, castForeignPtr)+import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)+import Foreign.Marshal.Utils (copyBytes, fillBytes)+import Foreign.Ptr (Ptr, plusPtr, ptrToWordPtr)+import Foreign.Storable (Storable (..))+import GHC.ForeignPtr (mallocPlainForeignPtrAlignedBytes, unsafeWithForeignPtr)+import System.IO.Unsafe (unsafeDupablePerformIO)+++-- ============================================================+-- Pinned buffers+-- ============================================================++-- | Fresh pinned memory, 64-byte aligned (the arrow-rs allocation alignment).+mallocAligned :: Int -> IO (ForeignPtr Word8)+mallocAligned n = mallocPlainForeignPtrAlignedBytes (max 0 n) 64+{-# INLINE mallocAligned #-}+++-- | A 64-byte aligned 'ByteString' of @n@ bytes filled by the action.+createAligned :: Int -> (Ptr Word8 -> IO ()) -> ByteString+createAligned !n fill = unsafeDupablePerformIO $ do+  fp <- mallocAligned n+  unsafeWithForeignPtr fp fill+  pure $! BSI.BS fp (max 0 n)+{-# INLINE createAligned #-}+++-- | O(1): view a storable vector's memory as bytes (no copy).+storableToBytes :: forall a. Storable a => VS.Vector a -> ByteString+storableToBytes v =+  let !(fp, n) = VS.unsafeToForeignPtr0 v+  in BSI.BS (castForeignPtr fp) (n * sizeOf (undefined :: a))+{-# INLINE storableToBytes #-}+++{- | View bytes as a storable vector of @length / sizeOf a@ elements.+O(1) when the bytes are aligned for @a@; otherwise one copy into a+fresh aligned buffer (arrow-rs @align_buffers@). Trailing bytes that+do not fill an element are ignored.+-}+bytesToStorable :: forall a. Storable a => ByteString -> VS.Vector a+bytesToStorable bs@(BSI.BS fp len)+  | aligned = unsafeBytesToStorable bs+  | otherwise = unsafeBytesToStorable (createAligned (n * sz) (\dst -> unsafeWithForeignPtr fp (\src -> copyBytes dst src (n * sz))))+  where+    !sz = sizeOf (undefined :: a)+    !n = len `quot` sz+    !aligned = fromIntegral (ptrToWordPtr (unsafeForeignPtrToPtr fp)) `rem` alignment (undefined :: a) == (0 :: Int)+{-# INLINE bytesToStorable #-}+++{- | O(1) alias of bytes as a storable vector. The caller guarantees+the bytes are aligned for @a@.+-}+unsafeBytesToStorable :: forall a. Storable a => ByteString -> VS.Vector a+unsafeBytesToStorable (BSI.BS fp len) = VS.unsafeFromForeignPtr0 (castForeignPtr fp) (len `quot` sizeOf (undefined :: a))+{-# INLINE unsafeBytesToStorable #-}+++-- | A detached copy of a storable vector in fresh aligned memory.+copyStorable :: forall a. Storable a => VS.Vector a -> VS.Vector a+copyStorable v = unsafeBytesToStorable (copyBytesFresh (storableToBytes v))+++copyBytesFresh :: ByteString -> ByteString+copyBytesFresh bs = createAligned (BS.length bs) $ \dst -> withBytesPtr bs $ \src -> copyBytes dst src (BS.length bs)+++-- | Run an action on the address of a 'ByteString'\'s first byte.+withBytesPtr :: ByteString -> (Ptr Word8 -> IO r) -> IO r+withBytesPtr (BSI.BS fp _) = unsafeWithForeignPtr fp+{-# INLINE withBytesPtr #-}+++-- | Run an action on the address of a storable vector's first element.+withStorablePtr :: Storable a => VS.Vector a -> (Ptr a -> IO r) -> IO r+withStorablePtr = VS.unsafeWith+{-# INLINE withStorablePtr #-}+++-- ============================================================+-- Bitmaps+-- ============================================================++{- | LSB-first bitmap: logical bit @i@ is bit @(offset + i) mod 8@ of+byte @(offset + i) div 8@. Slicing changes only the offset and the+length (arrow-rs @BooleanBuffer@).++Invariant: @offset >= 0@, @length >= 0@ and the bytes cover+@offset + length@ bits.+-}+data Bitmap = Bitmap !ByteString {-# UNPACK #-} !Int {-# UNPACK #-} !Int+++instance NFData Bitmap where+  rnf = (`seq` ())+++-- | Logical equality: same length, same bits (offsets ignored).+instance Eq Bitmap where+  a == b =+    bitmapLength a == bitmapLength b+      && all (\i -> unsafeBitAt a i == unsafeBitAt b i) [0 .. bitmapLength a - 1]+++instance Show Bitmap where+  showsPrec d b =+    showParen (d > 10) $+      showString "bitmap " . shows (map (\x -> if x then '1' else '0') (bitmapToBools b))+++-- | The underlying bytes (they may extend past the logical range).+bitmapBytes :: Bitmap -> ByteString+bitmapBytes (Bitmap bs _ _) = bs+{-# INLINE bitmapBytes #-}+++-- | Offset of logical bit 0, in bits.+bitmapOffset :: Bitmap -> Int+bitmapOffset (Bitmap _ o _) = o+{-# INLINE bitmapOffset #-}+++-- | Number of logical bits.+bitmapLength :: Bitmap -> Int+bitmapLength (Bitmap _ _ n) = n+{-# INLINE bitmapLength #-}+++-- | Checked constructor: bytes, bit offset, bit length.+mkBitmap :: ByteString -> Int -> Int -> Either String Bitmap+mkBitmap bs off len+  | off < 0 || len < 0 = Left "Arrow.Column.Buffer.mkBitmap: negative offset or length"+  | toInteger (BS.length bs) * 8 < toInteger off + toInteger len =+      Left "Arrow.Column.Buffer.mkBitmap: bytes do not cover offset + length bits"+  | otherwise = Right (Bitmap bs off len)+++emptyBitmap :: Bitmap+emptyBitmap = Bitmap BS.empty 0 0+++-- | Bit @i@; errors when @i@ is outside @[0, length)@ (like 'V.!').+bitAt :: Bitmap -> Int -> Bool+bitAt b i+  | i < 0 || i >= bitmapLength b = errorWithoutStackTrace ("Arrow.Column.Buffer.bitAt: index " ++ show i ++ " out of range")+  | otherwise = unsafeBitAt b i+{-# INLINE bitAt #-}+++-- | Bit @i@ without a range check.+unsafeBitAt :: Bitmap -> Int -> Bool+unsafeBitAt (Bitmap bs off _) i =+  let !k = off + i+  in (BSU.unsafeIndex bs (k `unsafeShiftR` 3) `unsafeShiftR` (k .&. 7)) .&. 1 /= 0+{-# INLINE unsafeBitAt #-}+++-- | Bits @[s, s + l)@; the caller keeps the window inside the bitmap.+sliceBitmap :: Int -> Int -> Bitmap -> Bitmap+sliceBitmap s l (Bitmap bs off _) = Bitmap bs (off + s) l+{-# INLINE sliceBitmap #-}+++-- | Number of set bits (C popcount over the bit range).+bitmapSetCount :: Bitmap -> Int+bitmapSetCount (Bitmap bs off len)+  | len <= 0 = 0+  | otherwise = unsafeDupablePerformIO $ withBytesPtr bs $ \p -> popCountBits p off len+++{- | A fresh bitmap of @n@ bits whose bit @i@ is @f i@. One pass, eight+calls per output byte; @f@ inlines into the loop.+-}+bitmapGenerate :: Int -> (Int -> Bool) -> Bitmap+bitmapGenerate !n0 f =+  let !n = max 0 n0+      !nbytes = (n + 7) `unsafeShiftR` 3+      bs = createAligned nbytes $ \p ->+        let byteLoop !j+              | j >= nbytes = pure ()+              | otherwise = do+                  let !base = j `unsafeShiftL` 3+                      !hi = min 8 (n - base)+                      bitLoop !k !acc+                        | k >= hi = acc+                        | f (base + k) = bitLoop (k + 1) (acc .|. (1 `unsafeShiftL` k))+                        | otherwise = bitLoop (k + 1) acc+                  pokeByteOff p j (bitLoop 0 (0 :: Word8))+                  byteLoop (j + 1)+        in byteLoop 0+  in Bitmap bs 0 n+{-# INLINE bitmapGenerate #-}+++bitmapFromBools :: V.Vector Bool -> Bitmap+bitmapFromBools v = bitmapGenerate (V.length v) (V.unsafeIndex v)+++bitmapToBools :: Bitmap -> [Bool]+bitmapToBools b = map (unsafeBitAt b) [0 .. bitmapLength b - 1]+++-- | A detached copy starting at bit offset 0.+copyBitmap :: Bitmap -> Bitmap+copyBitmap (Bitmap bs off len) =+  let !nbytes = (len + 7) `unsafeShiftR` 3+  in Bitmap+       ( createAligned nbytes $ \dst -> do+           fillBytes dst 0 nbytes+           withBytesPtr bs $ \src -> copyBits dst 0 src off len+       )+       0+       len+++-- | Concatenate bitmaps into one fresh buffer (one allocation, bit copies).+concatBitmaps :: [Bitmap] -> Bitmap+concatBitmaps bms =+  let !total = sum (map bitmapLength bms)+      !nbytes = (total + 7) `unsafeShiftR` 3+      bs = createAligned nbytes $ \dst -> do+        fillBytes dst 0 nbytes+        let go !_ [] = pure ()+            go !pos (Bitmap b o l : rest) = do+              withBytesPtr b $ \src -> copyBits dst pos src o l+              go (pos + l) rest+        go 0 bms+  in Bitmap bs 0 total+++-- | Gather bits by index (indices already range-checked). Returns the set count too.+takeBitmap :: VS.Vector Int -> Bitmap -> (Bitmap, Int)+takeBitmap ix (Bitmap bs off _) =+  let !n = VS.length ix+      !nbytes = (n + 7) `unsafeShiftR` 3+  in unsafeDupablePerformIO $ do+       fp <- mallocAligned nbytes+       set <- unsafeWithForeignPtr fp $ \dst ->+         withBytesPtr bs $ \src -> VS.unsafeWith ix $ \pix -> gatherBits dst src off pix n+       pure (Bitmap (BSI.BS fp nbytes) 0 n, set)+++-- | Set bits @[pos, pos + r)@ of a zero-initialised buffer.+setBitRange :: Ptr Word8 -> Int -> Int -> IO ()+setBitRange p !pos !r+  | r <= 0 = pure ()+  | otherwise = do+      let !end = pos + r+          !b0 = pos `unsafeShiftR` 3+          !b1 = (end - 1) `unsafeShiftR` 3+          orByte j m = do+            x <- peekByteOff p j :: IO Word8+            pokeByteOff p j (x .|. m)+          lowMask k = (1 `unsafeShiftL` k) - 1 :: Word8+      if b0 == b1+        then orByte b0 (lowMask ((end - 1) .&. 7 + 1) .&. complement (lowMask (pos .&. 7)))+        else do+          orByte b0 (complement (lowMask (pos .&. 7)))+          fillBytes (p `plusPtr` (b0 + 1)) 0xFF (b1 - b0 - 1)+          orByte b1 (lowMask ((end - 1) .&. 7 + 1))+++-- ============================================================+-- Validity+-- ============================================================++{- | A validity bitmap with its null count (arrow-rs @NullBuffer@).+Bit set means the row is valid.++Invariant: @nullCount > 0@ and equals the number of clear bits.+A column without nulls carries @Nothing@ instead (normalised by+every constructor in this package).+-}+data Validity = Validity !Bitmap {-# UNPACK #-} !Int+  deriving stock (Eq)+++instance NFData Validity where+  rnf = (`seq` ())+++instance Show Validity where+  showsPrec d (Validity b n) =+    showParen (d > 10) $ showString "validity " . showsPrec 11 n . showChar ' ' . showsPrec 11 b+++validityBits :: Validity -> Bitmap+validityBits (Validity b _) = b+{-# INLINE validityBits #-}+++validityNullCount :: Validity -> Int+validityNullCount (Validity _ n) = n+{-# INLINE validityNullCount #-}+++-- | Count nulls (C popcount) and normalise: no nulls is 'Nothing'.+mkValidity :: Bitmap -> Maybe Validity+mkValidity b =+  let !nulls = bitmapLength b - bitmapSetCount b+  in if nulls == 0 then Nothing else Just (Validity b nulls)+++-- | @mkValidity . bitmapGenerate n@.+validityGenerate :: Int -> (Int -> Bool) -> Maybe Validity+validityGenerate n f = mkValidity (bitmapGenerate n f)+{-# INLINE validityGenerate #-}+++validityFromBools :: V.Vector Bool -> Maybe Validity+validityFromBools = mkValidity . bitmapFromBools+++{- | Validate a caller-supplied validity for an array of @rows@ rows:+the bitmap covers its range, its length is @rows@, and its null count+matches the bits. Normalises a validity with no nulls to 'Nothing'.+-}+checkValidity :: String -> Int -> Maybe Validity -> Either String (Maybe Validity)+checkValidity _ _ Nothing = Right Nothing+checkValidity what rows (Just (Validity b@(Bitmap bs off len) declared))+  | off < 0 || len < 0 || toInteger (BS.length bs) * 8 < toInteger off + toInteger len =+      Left (what ++ ": validity bytes do not cover the bitmap")+  | len /= rows = Left (what ++ ": validity has " ++ show len ++ " bits for " ++ show rows ++ " rows")+  | otherwise =+      let !nulls = len - bitmapSetCount b+      in if nulls /= declared+           then Left (what ++ ": validity null count " ++ show declared ++ " does not match the bitmap (" ++ show nulls ++ ")")+           else Right (if nulls == 0 then Nothing else Just (Validity b nulls))+++{- | Whether row @i@ is valid. 'Nothing' is all valid. Errors when @i@+is outside the bitmap (like 'V.!').+-}+isValidAt :: Maybe Validity -> Int -> Bool+isValidAt Nothing _ = True+isValidAt (Just (Validity b _)) i = bitAt b i+{-# INLINE isValidAt #-}+++-- | 'isValidAt' without the range check: one load, shift and mask.+unsafeIsValidAt :: Maybe Validity -> Int -> Bool+unsafeIsValidAt Nothing _ = True+unsafeIsValidAt (Just (Validity b _)) i = unsafeBitAt b i+{-# INLINE unsafeIsValidAt #-}+++-- | Rows @[s, s + l)@ (caller clamps); recounts nulls with popcount.+sliceValidity :: Int -> Int -> Maybe Validity -> Maybe Validity+sliceValidity _ _ Nothing = Nothing+sliceValidity s l (Just (Validity b _)) = mkValidity (sliceBitmap s l b)+{-# INLINE sliceValidity #-}+++{- | Concatenate the validity of consecutive pieces, each given with+its row count. All-'Nothing' stays 'Nothing'; otherwise one fresh+bitmap (bit copies, all-valid runs set in bulk).+-}+concatValidity :: [(Int, Maybe Validity)] -> Maybe Validity+concatValidity pieces+  | all (\(_, v) -> null v) pieces = Nothing+  | otherwise =+      let !total = sum (map fst pieces)+          !nulls = sum (map (\(_, v) -> maybe 0 validityNullCount v) pieces)+          !nbytes = (total + 7) `unsafeShiftR` 3+          bs = createAligned nbytes $ \dst -> do+            fillBytes dst 0 nbytes+            let go !_ [] = pure ()+                go !pos ((r, v) : rest) = do+                  case v of+                    Nothing -> setBitRange dst pos r+                    Just (Validity (Bitmap b o _) _) -> withBytesPtr b $ \src -> copyBits dst pos src o r+                  go (pos + r) rest+            go 0 pieces+      in if nulls == 0 then Nothing else Just (Validity (Bitmap bs 0 total) nulls)+++-- | Gather validity by index (indices already range-checked).+takeValidity :: VS.Vector Int -> Maybe Validity -> Maybe Validity+takeValidity _ Nothing = Nothing+takeValidity ix (Just (Validity b _)) =+  let !(b', set) = takeBitmap ix b+      !nulls = VS.length ix - set+  in if nulls == 0 then Nothing else Just (Validity b' nulls)+++-- | Row-wise AND of two validities of the same length.+andValidity :: Maybe Validity -> Maybe Validity -> Maybe Validity+andValidity Nothing v = v+andValidity v Nothing = v+andValidity (Just (Validity (Bitmap ba oa la) _)) (Just (Validity (Bitmap bb ob _) _)) =+  let !nbytes = (la + 7) `unsafeShiftR` 3+  in unsafeDupablePerformIO $ do+       fp <- mallocAligned nbytes+       set <- unsafeWithForeignPtr fp $ \dst ->+         withBytesPtr ba $ \pa -> withBytesPtr bb $ \pb -> andBits dst pa oa pb ob la+       let !nulls = la - set+       pure (if nulls == 0 then Nothing else Just (Validity (Bitmap (BSI.BS fp nbytes) 0 la) nulls))+++-- | Detached copy at bit offset 0.+copyValidity :: Maybe Validity -> Maybe Validity+copyValidity = fmap (\(Validity b n) -> Validity (copyBitmap b) n)+++-- ============================================================+-- Element types+-- ============================================================++-- | IEEE 754 binary16, kept as its bit pattern.+newtype Float16 = Float16 Word16+  deriving stock (Show)+  deriving newtype (Eq, Ord, Storable)+++-- | binary16 to 'Double' (exact).+float16ToDouble :: Float16 -> Double+float16ToDouble (Float16 w) =+  let !sign = if w .&. 0x8000 /= 0 then -1 else 1+      !ex = fromIntegral ((w `shiftR` 10) .&. 0x1f) :: Int+      !mant = fromIntegral (w .&. 0x3ff) :: Double+  in case ex of+       0 -> sign * mant * 2 ** (-24)+       31 -> if mant == 0 then sign * (1 / 0) else 0 / 0+       _ -> sign * (1 + mant / 1024) * 2 ^^ (ex - 15)+++-- | Arrow @INTERVAL(DAY_TIME)@: days, milliseconds (8 bytes).+data IntervalDayTime = IntervalDayTime {-# UNPACK #-} !Int32 {-# UNPACK #-} !Int32+  deriving stock (Show, Eq, Ord)+++instance Storable IntervalDayTime where+  sizeOf _ = 8+  alignment _ = 4+  peek p = IntervalDayTime <$> peekByteOff p 0 <*> peekByteOff p 4+  poke p (IntervalDayTime d m) = pokeByteOff p 0 d >> pokeByteOff p 4 m+  {-# INLINE peek #-}+  {-# INLINE poke #-}+++-- | Arrow @INTERVAL(MONTH_DAY_NANO)@: months, days, nanoseconds (16 bytes).+data IntervalMonthDayNano = IntervalMonthDayNano {-# UNPACK #-} !Int32 {-# UNPACK #-} !Int32 {-# UNPACK #-} !Int64+  deriving stock (Show, Eq, Ord)+++instance Storable IntervalMonthDayNano where+  sizeOf _ = 16+  alignment _ = 8+  peek p = IntervalMonthDayNano <$> peekByteOff p 0 <*> peekByteOff p 4 <*> peekByteOff p 8+  poke p (IntervalMonthDayNano m d ns) = pokeByteOff p 0 m >> pokeByteOff p 4 d >> pokeByteOff p 8 ns+  {-# INLINE peek #-}+  {-# INLINE poke #-}+++-- | 128-bit little-endian two's complement decimal payload: low word, high word.+data Decimal128 = Decimal128 {-# UNPACK #-} !Word64 {-# UNPACK #-} !Word64+  deriving stock (Eq)+++instance Show Decimal128 where+  showsPrec d = showsPrec d . decimal128ToInteger+++instance Storable Decimal128 where+  sizeOf _ = 16+  alignment _ = 8+  peek p = Decimal128 <$> peekByteOff p 0 <*> peekByteOff p 8+  poke p (Decimal128 lo hi) = pokeByteOff p 0 lo >> pokeByteOff p 8 hi+  {-# INLINE peek #-}+  {-# INLINE poke #-}+++-- | Two's complement order.+instance Ord Decimal128 where+  compare a b = compare (decimal128ToInteger a) (decimal128ToInteger b)+++decimal128ToInteger :: Decimal128 -> Integer+decimal128ToInteger (Decimal128 lo hi) =+  toInteger (fromIntegral hi :: Int64) `shiftL` 64 + toInteger lo+++-- | Wraps modulo 2^128.+decimal128FromInteger :: Integer -> Decimal128+decimal128FromInteger x = Decimal128 (fromInteger x) (fromInteger (x `shiftR` 64))+++-- | 256-bit little-endian two's complement decimal payload, least significant word first.+data Decimal256 = Decimal256 {-# UNPACK #-} !Word64 {-# UNPACK #-} !Word64 {-# UNPACK #-} !Word64 {-# UNPACK #-} !Word64+  deriving stock (Eq)+++instance Show Decimal256 where+  showsPrec d = showsPrec d . decimal256ToInteger+++instance Storable Decimal256 where+  sizeOf _ = 32+  alignment _ = 8+  peek p = Decimal256 <$> peekByteOff p 0 <*> peekByteOff p 8 <*> peekByteOff p 16 <*> peekByteOff p 24+  poke p (Decimal256 a b c d) = pokeByteOff p 0 a >> pokeByteOff p 8 b >> pokeByteOff p 16 c >> pokeByteOff p 24 d+  {-# INLINE peek #-}+  {-# INLINE poke #-}+++instance Ord Decimal256 where+  compare a b = compare (decimal256ToInteger a) (decimal256ToInteger b)+++decimal256ToInteger :: Decimal256 -> Integer+decimal256ToInteger (Decimal256 w0 w1 w2 w3) =+  toInteger (fromIntegral w3 :: Int64) `shiftL` 192+    + toInteger w2 `shiftL` 128+    + toInteger w1 `shiftL` 64+    + toInteger w0+++-- | Wraps modulo 2^256.+decimal256FromInteger :: Integer -> Decimal256+decimal256FromInteger x =+  Decimal256 (fromInteger x) (fromInteger (x `shiftR` 64)) (fromInteger (x `shiftR` 128)) (fromInteger (x `shiftR` 192))+++{- | Copy @n@ elements of @w@ bytes, reversing the byte order of each+group of the pattern within every element. This converts either+byte-order direction; callers supply complete, checked elements.+-}+swapCopy :: [Int] -> Int -> Int -> ByteString -> ByteString+swapCopy pat w n src = createAligned (n * w) $ \dst -> withBytesPtr src $ \s ->+  let elemLoop !e+        | e >= n = pure ()+        | otherwise = groups (e * w) pat >> elemLoop (e + 1)+      groups !_ [] = pure ()+      groups !base (g : gs) = swapGroup (s `plusPtr` base) (dst `plusPtr` base) g >> groups (base + g) gs+  in elemLoop 0+++swapGroup :: Ptr Word8 -> Ptr Word8 -> Int -> IO ()+swapGroup s d g = go 0+  where+    go !k+      | k >= g = pure ()+      | otherwise = do+          b <- peekByteOff s (g - 1 - k) :: IO Word8+          pokeByteOff d k b+          go (k + 1)+++{- | Convert a checked 16-byte view descriptor buffer between little and big+endian. @sourceBig@ describes the input. Inline string bytes and the+four-byte long-string prefix stay byte-for-byte unchanged; only length,+buffer index and offset are numeric. The caller supplies complete views.+-}+swapViewCopy :: Bool -> ByteString -> ByteString+swapViewCopy sourceBig src = createAligned (BS.length src) $ \dst -> withBytesPtr src $ \s -> do+  copyBytes dst s (BS.length src)+  let go !base+        | base >= BS.length src = pure ()+        | otherwise = do+            a <- peekByteOff s base :: IO Word8+            b <- peekByteOff s (base + 1) :: IO Word8+            c <- peekByteOff s (base + 2) :: IO Word8+            d <- peekByteOff s (base + 3) :: IO Word8+            let !len = if sourceBig then word d c b a else word a b c d+            swapGroup (s `plusPtr` base) (dst `plusPtr` base) 4+            if len > 12+              then do+                swapGroup (s `plusPtr` (base + 8)) (dst `plusPtr` (base + 8)) 4+                swapGroup (s `plusPtr` (base + 12)) (dst `plusPtr` (base + 12)) 4+              else pure ()+            go (base + 16)+      word a b c d = fromIntegral a .|. (fromIntegral b `shiftL` 8) .|. (fromIntegral c `shiftL` 16) .|. (fromIntegral d `shiftL` 24) :: Word64+  go 0
+ src/Arrow/Column/Builder.hs view
@@ -0,0 +1,456 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}++{- | Streaming column builders in 'ST' or 'IO' ('PrimMonad').++Each builder appends rows into growable, pinned, 64-byte aligned+buffers (amortised doubling) and freezes into a 'ColumnArray' whose+buffers are those allocations: no per-row heap objects, no final+copy. The validity bitmap is only allocated once the first null is+appended, so a builder that never sees a null freezes with no+validity.++Values are valid by construction ('appendText' copies the bytes of a+'Text', which are UTF-8), so freezing runs no validation.++'freezeBuilder' hands the buffers to the column and resets the+builder to empty; appending afterwards starts a new column in fresh+memory and never touches the frozen one.++Appending more than 2^31 - 1 bytes to a 32-bit offset builder+('newUtf8Builder', 'newBinaryBuilder') is an error; use the large+variants for such data.+-}+module Arrow.Column.Builder (+  -- * Common interface+  ColumnBuilder (..),++  -- * Fixed width+  PrimBuilder,+  newPrimBuilder,+  appendPrim,+  appendPrimMaybe,++  -- * Booleans+  BoolBuilder,+  newBoolBuilder,+  appendBool,+  appendBoolMaybe,++  -- * Strings and bytes+  Utf8Builder,+  newUtf8Builder,+  newLargeUtf8Builder,+  appendText,+  appendTextMaybe,+  BinaryBuilder,+  newBinaryBuilder,+  newLargeBinaryBuilder,+  appendBytes,+  appendBytesMaybe,+) where++import Arrow.Column.Internal+import Control.Monad.Primitive (PrimMonad, PrimState, unsafeIOToPrim)+import Data.Bits (complement, unsafeShiftL, unsafeShiftR, (.&.), (.|.))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Internal qualified as BSI+import Data.ByteString.Unsafe qualified as BSU+import Data.Int (Int32, Int64)+import Data.Primitive.MutVar (MutVar, newMutVar, readMutVar, writeMutVar)+import Data.Primitive.PrimArray (MutablePrimArray, newPrimArray, readPrimArray, setPrimArray, writePrimArray)+import Data.Text (Text)+import Data.Text.Foreign qualified as TF+import Data.Vector.Storable qualified as VS+import Data.Word (Word8)+import Foreign.ForeignPtr (ForeignPtr, castForeignPtr)+import GHC.ForeignPtr (unsafeWithForeignPtr)+import Foreign.Marshal.Utils (copyBytes, fillBytes)+import Foreign.Ptr (Ptr, castPtr, plusPtr)+import Foreign.Storable (Storable (..))+++-- | Operations every builder supports.+class ColumnBuilder b where+  -- | Append a null row.+  appendNull :: PrimMonad m => b (PrimState m) -> m ()++  -- | Rows appended so far.+  builderLength :: PrimMonad m => b (PrimState m) -> m Int++  -- | Freeze into a column and reset the builder to empty.+  freezeBuilder :: PrimMonad m => b (PrimState m) -> m ColumnArray+++-- ============================================================+-- Growable byte buffers+-- ============================================================++-- | A growable pinned byte buffer: the memory and its capacity in bytes.+data Grow s = Grow !(MutVar s (ForeignPtr Word8)) !(MutablePrimArray s Int)+++{-# INLINABLE newGrow #-}+newGrow :: PrimMonad m => Int -> m (Grow (PrimState m))+newGrow cap0 = do+  let !cap = max 0 cap0+  fp <- unsafeIOToPrim $ do+    fp <- mallocAligned cap+    unsafeWithForeignPtr fp $ \p -> fillBytes p 0 cap+    pure fp+  ref <- newMutVar fp+  meta <- newPrimArray 1+  writePrimArray meta 0 cap+  pure (Grow ref meta)+++-- | Make room for @need@ bytes, keeping the first @used@ bytes and zeroing the rest.+ensureGrow :: PrimMonad m => Grow (PrimState m) -> Int -> Int -> m ()+ensureGrow (Grow ref meta) !used !need = do+  cap <- readPrimArray meta 0+  if need <= cap+    then pure ()+    else do+      let !cap' = max need (max 64 (cap * 2))+      old <- readMutVar ref+      new <- unsafeIOToPrim $ do+        fp <- mallocAligned cap'+        unsafeWithForeignPtr fp $ \dst -> unsafeWithForeignPtr old $ \src -> do+          copyBytes dst src used+          fillBytes (dst `plusPtr` used) 0 (cap' - used)+        pure fp+      writeMutVar ref new+      writePrimArray meta 0 cap'+{-# INLINE ensureGrow #-}+++-- | Hand the memory out and replace it with an empty allocation.+{-# INLINABLE takeGrow #-}+takeGrow :: PrimMonad m => Grow (PrimState m) -> m (ForeignPtr Word8)+takeGrow (Grow ref meta) = do+  fp <- readMutVar ref+  fresh <- unsafeIOToPrim (mallocAligned 0)+  writeMutVar ref fresh+  writePrimArray meta 0 0+  pure fp+++-- | Validity under construction: the bits, plus [null count, materialised flag].+data ValidBuf s = ValidBuf !(Grow s) !(MutablePrimArray s Int)+++{-# INLINABLE newValidBuf #-}+newValidBuf :: PrimMonad m => m (ValidBuf (PrimState m))+newValidBuf = do+  g <- newGrow 0+  meta <- newPrimArray 2+  setPrimArray meta 0 2 0+  pure (ValidBuf g meta)+++writeBit :: PrimMonad m => Grow (PrimState m) -> Int -> Bool -> m ()+writeBit g@(Grow ref _) !i !b = do+  let !byte = i `unsafeShiftR` 3+  ensureGrow g byte (byte + 1)+  fp <- readMutVar ref+  unsafeIOToPrim $ unsafeWithForeignPtr fp $ \p -> do+    x <- peekByteOff p byte :: IO Word8+    let !m = 1 `unsafeShiftL` (i .&. 7)+    pokeByteOff p byte (if b then x .|. m else x .&. complement m)+{-# INLINE writeBit #-}+++-- | Row @i@ is valid.+validAt :: PrimMonad m => ValidBuf (PrimState m) -> Int -> m ()+validAt (ValidBuf g meta) !i = do+  materialised <- readPrimArray meta 1+  if materialised == 0 then pure () else writeBit g i True+{-# INLINE validAt #-}+++-- | Row @i@ is null; rows before it that were appended valid get their bits now.+{-# INLINE nullAt #-}+nullAt :: PrimMonad m => ValidBuf (PrimState m) -> Int -> m ()+nullAt (ValidBuf g@(Grow ref _) meta) !i = do+  materialised <- readPrimArray meta 1+  if materialised /= 0+    then pure ()+    else do+      let !full = i `unsafeShiftR` 3+      ensureGrow g 0 (full + 1)+      fp <- readMutVar ref+      unsafeIOToPrim $ unsafeWithForeignPtr fp $ \p -> do+        fillBytes p 0xFF full+        pokeByteOff p full ((1 `unsafeShiftL` (i .&. 7)) - 1 :: Word8)+      writePrimArray meta 1 1+  writeBit g i False+  n <- readPrimArray meta 0+  writePrimArray meta 0 (n + 1)+++{-# INLINABLE freezeValid #-}+freezeValid :: PrimMonad m => ValidBuf (PrimState m) -> Int -> m (Maybe Validity)+freezeValid (ValidBuf g meta) !rows = do+  nulls <- readPrimArray meta 0+  fp <- takeGrow g+  setPrimArray meta 0 2 0+  pure $+    if nulls == 0+      then Nothing+      else Just (Validity (Bitmap (BSI.BS fp ((rows + 7) `unsafeShiftR` 3)) 0 rows) nulls)+++-- ============================================================+-- Fixed width+-- ============================================================++-- | Builder for a fixed-width column of tag @PrimType a@.+data PrimBuilder a s = PrimBuilder !(PrimType a) !(Grow s) !(ValidBuf s) !(MutablePrimArray s Int)+++-- | A builder with room for the given number of rows before it first grows.+{-# INLINABLE newPrimBuilder #-}+newPrimBuilder :: PrimMonad m => PrimType a -> Int -> m (PrimBuilder a (PrimState m))+newPrimBuilder t hint = do+  g <- newGrow (max 0 hint * primWidth t)+  vb <- newValidBuf+  len <- newPrimArray 1+  writePrimArray len 0 0+  pure (PrimBuilder t g vb len)+++appendPrim :: (PrimMonad m, Storable a) => PrimBuilder a (PrimState m) -> a -> m ()+appendPrim (PrimBuilder _ g@(Grow ref _) vb lenRef) x = do+  n <- readPrimArray lenRef 0+  let !w = sizeOf x+  ensureGrow g (n * w) ((n + 1) * w)+  fp <- readMutVar ref+  unsafeIOToPrim $ unsafeWithForeignPtr fp $ \p -> pokeElemOff (castPtr p) n x+  validAt vb n+  writePrimArray lenRef 0 (n + 1)+{-# INLINE appendPrim #-}+++appendPrimMaybe :: (PrimMonad m, Storable a) => PrimBuilder a (PrimState m) -> Maybe a -> m ()+appendPrimMaybe b = maybe (appendNull b) (appendPrim b)+{-# INLINE appendPrimMaybe #-}+++instance ColumnBuilder (PrimBuilder a) where+  appendNull (PrimBuilder t g vb lenRef) = do+    n <- readPrimArray lenRef 0+    let !w = primWidth t+    ensureGrow g (n * w) ((n + 1) * w) -- new bytes are zeroed+    nullAt vb n+    writePrimArray lenRef 0 (n + 1)+  builderLength (PrimBuilder _ _ _ lenRef) = readPrimArray lenRef 0+  {-# INLINE appendNull #-}+  {-# INLINE builderLength #-}+  {-# INLINABLE freezeBuilder #-}+  freezeBuilder (PrimBuilder t g vb lenRef) = withPrim t $ do+    n <- readPrimArray lenRef 0+    fp <- takeGrow g+    v <- freezeValid vb n+    writePrimArray lenRef 0 0+    pure (ColPrim t v (VS.unsafeFromForeignPtr0 (castForeignPtr fp) n))+++-- ============================================================+-- Booleans+-- ============================================================++data BoolBuilder s = BoolBuilder !(Grow s) !(ValidBuf s) !(MutablePrimArray s Int)+++{-# INLINABLE newBoolBuilder #-}+newBoolBuilder :: PrimMonad m => Int -> m (BoolBuilder (PrimState m))+newBoolBuilder hint = do+  g <- newGrow ((max 0 hint + 7) `unsafeShiftR` 3)+  vb <- newValidBuf+  len <- newPrimArray 1+  writePrimArray len 0 0+  pure (BoolBuilder g vb len)+++appendBool :: PrimMonad m => BoolBuilder (PrimState m) -> Bool -> m ()+appendBool (BoolBuilder g vb lenRef) x = do+  n <- readPrimArray lenRef 0+  writeBit g n x+  validAt vb n+  writePrimArray lenRef 0 (n + 1)+{-# INLINE appendBool #-}+++appendBoolMaybe :: PrimMonad m => BoolBuilder (PrimState m) -> Maybe Bool -> m ()+appendBoolMaybe b = maybe (appendNull b) (appendBool b)+{-# INLINE appendBoolMaybe #-}+++instance ColumnBuilder BoolBuilder where+  appendNull (BoolBuilder g vb lenRef) = do+    n <- readPrimArray lenRef 0+    writeBit g n False+    nullAt vb n+    writePrimArray lenRef 0 (n + 1)+  builderLength (BoolBuilder _ _ lenRef) = readPrimArray lenRef 0+  {-# INLINE appendNull #-}+  {-# INLINE builderLength #-}+  {-# INLINABLE freezeBuilder #-}+  freezeBuilder (BoolBuilder g vb lenRef) = do+    n <- readPrimArray lenRef 0+    -- A builder with no rows may never have allocated a byte.+    let !nbytes = (n + 7) `unsafeShiftR` 3+    ensureGrow g nbytes nbytes+    fp <- takeGrow g+    v <- freezeValid vb n+    writePrimArray lenRef 0 0+    pure (ColBool v (Bitmap (BSI.BS fp nbytes) 0 n))+++-- ============================================================+-- Var-length+-- ============================================================++{- | Offsets, data, validity, [rows, data length], and the constructor+for the frozen column.+-}+data VarBuilder o s+  = VarBuilder+      !(Grow s)+      !(Grow s)+      !(ValidBuf s)+      !(MutablePrimArray s Int)+      !(Maybe Validity -> VS.Vector o -> ByteString -> ColumnArray)+++{-# INLINABLE newVarBuilder #-}+newVarBuilder+  :: forall o m+   . (PrimMonad m, Offset o)+  => (Maybe Validity -> VS.Vector o -> ByteString -> ColumnArray)+  -> Int+  -> m (VarBuilder o (PrimState m))+newVarBuilder mk hint = do+  let !w = sizeOf (0 :: o)+  offs <- newGrow ((max 0 hint + 1) * w)+  ensureGrow offs 0 w -- offset 0, zeroed+  dat <- newGrow (max 0 hint * 8)+  vb <- newValidBuf+  meta <- newPrimArray 2+  setPrimArray meta 0 2 0+  pure (VarBuilder offs dat vb meta mk)+++-- | Append one row of @len@ bytes copied by the action.+appendVar :: forall o m. (PrimMonad m, Offset o) => VarBuilder o (PrimState m) -> Int -> (Ptr Word8 -> IO ()) -> m ()+appendVar (VarBuilder offs@(Grow offRef _) dat@(Grow datRef _) vb meta _) !len copy = do+  n <- readPrimArray meta 0+  used <- readPrimArray meta 1+  let !end = used + len+      !w = sizeOf (0 :: o)+  if toInteger end > toInteger (maxBound :: o)+    then errorWithoutStackTrace "Arrow.Column.Builder: var-length data exceeds the offset width (use a large builder)"+    else pure ()+  ensureGrow dat used end+  fp <- readMutVar datRef+  unsafeIOToPrim $ unsafeWithForeignPtr fp $ \p -> copy (p `plusPtr` used)+  ensureGrow offs ((n + 1) * w) ((n + 2) * w)+  ofp <- readMutVar offRef+  unsafeIOToPrim $ unsafeWithForeignPtr ofp $ \p -> pokeElemOff (castPtr p) (n + 1) (fromIntegral end :: o)+  validAt vb n+  writePrimArray meta 0 (n + 1)+  writePrimArray meta 1 end+{-# INLINE appendVar #-}+++{-# INLINE appendVarNull #-}+appendVarNull :: forall o m. (PrimMonad m, Offset o) => VarBuilder o (PrimState m) -> m ()+appendVarNull (VarBuilder offs@(Grow offRef _) _ vb meta _) = do+  n <- readPrimArray meta 0+  used <- readPrimArray meta 1+  let !w = sizeOf (0 :: o)+  ensureGrow offs ((n + 1) * w) ((n + 2) * w)+  ofp <- readMutVar offRef+  unsafeIOToPrim $ unsafeWithForeignPtr ofp $ \p -> pokeElemOff (castPtr p) (n + 1) (fromIntegral used :: o)+  nullAt vb n+  writePrimArray meta 0 (n + 1)+++{-# INLINABLE freezeVar #-}+freezeVar :: forall o m. (PrimMonad m, Offset o) => VarBuilder o (PrimState m) -> m ColumnArray+freezeVar (VarBuilder offs dat vb meta mk) = do+  n <- readPrimArray meta 0+  used <- readPrimArray meta 1+  ofp <- takeGrow offs+  dfp <- takeGrow dat+  v <- freezeValid vb n+  setPrimArray meta 0 2 0+  -- Restart with offset 0 in place for the next column.+  ensureGrow offs 0 (sizeOf (0 :: o))+  pure (mk v (VS.unsafeFromForeignPtr0 (castForeignPtr ofp) (n + 1)) (BSI.BS dfp used))+++-- | Builder for a utf8 column (32-bit offsets) or large utf8 column (64-bit).+newtype Utf8Builder o s = Utf8Builder (VarBuilder o s)+++-- | Builder for a binary column (32-bit offsets) or large binary column (64-bit).+newtype BinaryBuilder o s = BinaryBuilder (VarBuilder o s)+++newUtf8Builder :: PrimMonad m => Int -> m (Utf8Builder Int32 (PrimState m))+newUtf8Builder hint = Utf8Builder <$> newVarBuilder ColUtf8 hint+++newLargeUtf8Builder :: PrimMonad m => Int -> m (Utf8Builder Int64 (PrimState m))+newLargeUtf8Builder hint = Utf8Builder <$> newVarBuilder ColLargeUtf8 hint+++newBinaryBuilder :: PrimMonad m => Int -> m (BinaryBuilder Int32 (PrimState m))+newBinaryBuilder hint = BinaryBuilder <$> newVarBuilder ColBinary hint+++newLargeBinaryBuilder :: PrimMonad m => Int -> m (BinaryBuilder Int64 (PrimState m))+newLargeBinaryBuilder hint = BinaryBuilder <$> newVarBuilder ColLargeBinary hint+++-- | Append the UTF-8 bytes of a 'Text' (one copy, no validation needed).+appendText :: (PrimMonad m, Offset o) => Utf8Builder o (PrimState m) -> Text -> m ()+appendText (Utf8Builder vb) t = appendVar vb (TF.lengthWord8 t) (TF.unsafeCopyToPtr t)+{-# INLINE appendText #-}+++appendTextMaybe :: (PrimMonad m, Offset o) => Utf8Builder o (PrimState m) -> Maybe Text -> m ()+appendTextMaybe b = maybe (appendNull b) (appendText b)+{-# INLINE appendTextMaybe #-}+++appendBytes :: (PrimMonad m, Offset o) => BinaryBuilder o (PrimState m) -> ByteString -> m ()+appendBytes (BinaryBuilder vb) bs =+  appendVar vb (BS.length bs) $ \dst -> BSU.unsafeUseAsCStringLen bs $ \(src, l) -> copyBytes dst (castPtr src) l+{-# INLINE appendBytes #-}+++appendBytesMaybe :: (PrimMonad m, Offset o) => BinaryBuilder o (PrimState m) -> Maybe ByteString -> m ()+appendBytesMaybe b = maybe (appendNull b) (appendBytes b)+{-# INLINE appendBytesMaybe #-}+++instance Offset o => ColumnBuilder (Utf8Builder o) where+  appendNull (Utf8Builder vb) = appendVarNull vb+  {-# INLINE appendNull #-}+  {-# INLINE builderLength #-}+  {-# INLINE freezeBuilder #-}+  builderLength (Utf8Builder (VarBuilder _ _ _ meta _)) = readPrimArray meta 0+  freezeBuilder (Utf8Builder vb) = freezeVar vb+++instance Offset o => ColumnBuilder (BinaryBuilder o) where+  appendNull (BinaryBuilder vb) = appendVarNull vb+  {-# INLINE appendNull #-}+  {-# INLINE builderLength #-}+  {-# INLINE freezeBuilder #-}+  builderLength (BinaryBuilder (VarBuilder _ _ _ meta _)) = readPrimArray meta 0+  freezeBuilder (BinaryBuilder vb) = freezeVar vb
+ src/Arrow/Column/Internal.hs view
@@ -0,0 +1,641 @@+{-# 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)))
+ src/Arrow/Derive.hs view
@@ -0,0 +1,293 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TemplateHaskell #-}++{- | Annotation-driven Template Haskell deriver for Apache Arrow+'Arrow.Record.Table' values.++Arrow is a /columnar/ format: records become rows in a columnar+batch rather than recursive trees. The deriver therefore+produces row-level instances:++* For a record type, it emits a 'HasTable' instance so that+  @'encodeTable' 'hasTable'@ and @'decodeTable' 'hasTable'@+  round-trip a @'Data.Vector.Vector'@ of the record through an+  Arrow 'Schema' + @'Data.Vector.Vector' 'ColumnArray'@.+* For a @newtype@, it emits 'HasEncoder' and 'HasDecoder'+  instances that pass through to the inner type. This lets the+  newtype be used as a record field by the record deriver.++Sum types and enums fail at splice time because Arrow has no+first-class column-level union representation.++Modifiers honoured (resolved against 'backendArrow'):++* 'rename', 'renameStyle', 'renameWith' — name of the column in+  the emitted schema.+* 'skip' — drop the column entirely from both encode and decode+  sides; if combined with 'defaults', the decoder fills the+  field from the named default function.+* 'defaults' — the default value used for skipped fields when+  decoding.+* 'coerced' — the field is encoded/decoded via the named target+  type's 'HasEncoder' / 'HasDecoder'; @Data.Coerce.coerce@ wraps+  the boundary in both directions.+-}+module Arrow.Derive (+  -- * Class+  HasTable (..),++  -- * Derivers+  deriveArrow,+  deriveArrowTable,+  deriveArrowColumn,+) where++import Arrow.Record (+  Decoder,+  Encoder,+  Table,+  columnD,+  fieldE,+  table,+ )+import Arrow.Record.Generic (HasDecoder (..), HasEncoder (..))+import Data.Coerce (coerce)+import Data.Foldable (foldlM)+import Data.Functor.Contravariant (contramap)+import Data.Maybe (catMaybes)+import Data.Text qualified as T+import Language.Haskell.TH+import Wireform.Derive.Backend (backendArrow)+import Wireform.Derive.ModifierInfo+import Wireform.Derive.TypeInfo+++-- ---------------------------------------------------------------------------+-- HasTable+-- ---------------------------------------------------------------------------++{- | A record type with a derived Arrow 'Table'.++Analogous to 'HasEncoder' / 'HasDecoder' from+"Arrow.Record.Generic", but operates at the row level rather+than the column level: a 'Table' bundles a 'RowEncoder' and a+'RowDecoder' for the same record type.+-}+class HasTable a where+  hasTable :: Table a+++-- ---------------------------------------------------------------------------+-- Public entry points+-- ---------------------------------------------------------------------------++{- | Derive Arrow instances for the named type.++Dispatches on the data shape:++* Records get a 'HasTable' instance.+* Newtypes get 'HasEncoder' + 'HasDecoder' instances that+  pass through to the inner type.+* Sum types and enums fail at splice time.+-}+deriveArrow :: Name -> Q [Dec]+deriveArrow nm = do+  ti <- reifyTypeInfo nm+  case typeInfoShape ti of+    TypeShapeRecord c -> deriveRecordTable ti c+    TypeShapeNewtype c -> deriveNewtypeColumn ti c+    TypeShapeEnum _ -> fail (rejectEnumOrSum nm "enum")+    TypeShapeSum _ -> fail (rejectEnumOrSum nm "sum")+++{- | Synonym for 'deriveArrow', used when calling sites want to+emphasise the record code path.+-}+deriveArrowTable :: Name -> Q [Dec]+deriveArrowTable = deriveArrow+++{- | Synonym for 'deriveArrow', used when calling sites want to+emphasise the newtype (column) code path.+-}+deriveArrowColumn :: Name -> Q [Dec]+deriveArrowColumn = deriveArrow+++rejectEnumOrSum :: Name -> String -> String+rejectEnumOrSum nm shape =+  "Arrow.Derive: cannot derive Arrow instances for "+    ++ shape+    ++ " type "+    ++ show nm+    ++ ". Arrow has no first-class column-level union representation; "+    ++ "model alternatives as a discriminator + payload columns by hand."+++-- ---------------------------------------------------------------------------+-- Records → HasTable+-- ---------------------------------------------------------------------------++deriveRecordTable :: TypeInfo -> ConInfo -> Q [Dec]+deriveRecordTable ti c = do+  let typ = applyTypeArgs (ConT (typeInfoName ti)) (typeInfoVarTypes ti)+  encExp <- buildRowEncoder c+  decExp <- buildRowDecoder c+  body <- [|table $(pure encExp) $(pure decExp)|]+  let decl =+        InstanceD+          Nothing+          []+          (AppT (ConT ''HasTable) typ)+          [FunD 'hasTable [Clause [] (NormalB body) []]]+  pure [decl]+++{- | Build a 'RowEncoder' expression. Skipped fields drop out+entirely; if every field is skipped the encoder collapses to+'mempty' (a no-column batch).+-}+buildRowEncoder :: ConInfo -> Q Exp+buildRowEncoder c = do+  pieces <- mapM mkFieldEncoder (conInfoFields c)+  case catMaybes pieces of+    [] -> [|mempty|]+    (e : es) -> foldlM (\acc x -> [|$(pure acc) <> $(pure x)|]) e es+++mkFieldEncoder :: FieldInfo -> Q (Maybe Exp)+mkFieldEncoder fi@(FieldInfo mSel _) = do+  selName <- requireSelector mSel+  mi <- reifyModifierInfoFor backendArrow selName+  if miSkip mi+    then pure Nothing+    else do+      let selBase = T.pack (nameBase selName)+      keyExp <- renderWireKey mi selBase+      enc <- case miCoerce mi of+        Nothing ->+          [|fieldE $(pure keyExp) $(varE selName) hasEncoder|]+        Just tgtName ->+          let tgtTy = ConT tgtName+              srcTy = fieldInfoType fi+          in [|+               fieldE+                 $(pure keyExp)+                 ( (coerce :: $(pure srcTy) -> $(pure tgtTy))+                     . $(varE selName)+                 )+                 (hasEncoder :: Encoder $(pure tgtTy))+               |]+      pure (Just enc)+++{- | Build a 'RowDecoder' expression. Every field of the+constructor needs a slot in the applicative chain; skipped+fields use @pure <default>@ to fill in.+-}+buildRowDecoder :: ConInfo -> Q Exp+buildRowDecoder c = do+  let conName = conInfoName c+      fields = conInfoFields c+  case fields of+    [] -> [|pure $(conE conName)|]+    (f0 : fs) -> do+      e0 <- mkFieldDecoder f0+      hd <- [|$(conE conName) <$> $(pure e0)|]+      foldlM+        ( \acc f -> do+            ef <- mkFieldDecoder f+            [|$(pure acc) <*> $(pure ef)|]+        )+        hd+        fs+++mkFieldDecoder :: FieldInfo -> Q Exp+mkFieldDecoder fi@(FieldInfo mSel _) = do+  selName <- requireSelector mSel+  mi <- reifyModifierInfoFor backendArrow selName+  if miSkip mi+    then case miDefaults mi of+      Just defNm ->+        [|pure $(varE defNm)|]+      Nothing ->+        fail $+          "Arrow.Derive: skipped field "+            ++ show selName+            ++ " requires a 'defaults' modifier so the decoder can "+            ++ "fill the missing value."+    else do+      let selBase = T.pack (nameBase selName)+      keyExp <- renderWireKey mi selBase+      case miCoerce mi of+        Nothing ->+          [|columnD $(pure keyExp) hasDecoder|]+        Just tgtName ->+          let tgtTy = ConT tgtName+              srcTy = fieldInfoType fi+          in [|+               (coerce :: $(pure tgtTy) -> $(pure srcTy))+                 <$> columnD+                   $(pure keyExp)+                   (hasDecoder :: Decoder $(pure tgtTy))+               |]+++-- ---------------------------------------------------------------------------+-- Newtypes → HasEncoder / HasDecoder pass-through+-- ---------------------------------------------------------------------------++deriveNewtypeColumn :: TypeInfo -> ConInfo -> Q [Dec]+deriveNewtypeColumn ti c = do+  let typ = applyTypeArgs (ConT (typeInfoName ti)) (typeInfoVarTypes ti)+  innerTy <- case conInfoFields c of+    [FieldInfo _ ty] -> pure ty+    _ ->+      fail $+        "Arrow.Derive: newtype "+          ++ show (typeInfoName ti)+          ++ " must have exactly one field"+  encBody <-+    [|+      contramap+        (coerce :: $(pure typ) -> $(pure innerTy))+        (hasEncoder :: Encoder $(pure innerTy))+      |]+  decBody <-+    [|+      (coerce :: $(pure innerTy) -> $(pure typ))+        <$> (hasDecoder :: Decoder $(pure innerTy))+      |]+  let encInst =+        InstanceD+          Nothing+          []+          (AppT (ConT ''HasEncoder) typ)+          [FunD 'hasEncoder [Clause [] (NormalB encBody) []]]+      decInst =+        InstanceD+          Nothing+          []+          (AppT (ConT ''HasDecoder) typ)+          [FunD 'hasDecoder [Clause [] (NormalB decBody) []]]+  pure [encInst, decInst]+++-- ---------------------------------------------------------------------------+-- Helpers+-- ---------------------------------------------------------------------------++requireSelector :: Maybe Name -> Q Name+requireSelector (Just n) = pure n+requireSelector Nothing =+  fail+    "Arrow.Derive: cannot derive Arrow instances for a positional \+    \(non-record) constructor; Arrow columns are addressed by name."+++applyTypeArgs :: Type -> [Type] -> Type+applyTypeArgs = foldl AppT
+ src/Arrow/File.hs view
@@ -0,0 +1,90 @@+{-# LANGUAGE BangPatterns #-}++{- | Apache Arrow IPC file and stream readers for the standard Arrow+format (FlatBuffers metadata; files carry the @ARROW1@ magic and a+FlatBuffers footer), as written by pyarrow, arrow-cpp, arrow-rs,+"Arrow.Stream" and "Arrow.Write".++'readArrowFileColumns' is the convenience entry point: it returns the+schema and every record batch materialized, with body compression+undone and dictionary columns resolved. 'readArrowFile' and+'readArrowStream' expose the raw @(RecordBatchDef, body)@ pairs in wire+layout; decode them with 'Arrow.Read.Columns.decodeRecordBatch' (which+also undoes body compression).+-}+module Arrow.File (+  ArrowFile (..),+  ArrowStream (..),+  readArrowFile,+  readArrowFileColumns,+  readArrowStream,+  readIPCMessage,+) where++import Arrow.Column (ColumnArray)+import Arrow.FlatBufferIPC (decodeMessageFrame, readArrowFileFBWithDicts, readArrowStreamFBWithDicts, readFrameHeader)+import Arrow.Stream (ReadOptions, decodeArrowFile)+import Arrow.Types+import Data.Bits (complement, (.&.))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.Vector qualified as V+++-- | Schema and record batches of an Arrow IPC file, in wire layout.+data ArrowFile = ArrowFile+  { afSchema :: !Schema+  , afBatches :: !(V.Vector (RecordBatchDef, ByteString))+  }+  deriving stock (Show, Eq)+++-- | Schema and record batches of an Arrow IPC stream, in wire layout.+data ArrowStream = ArrowStream+  { asSchema :: !Schema+  , asBatches :: !(V.Vector (RecordBatchDef, ByteString))+  }+  deriving stock (Show, Eq)+++{- | Read the encapsulated message starting at the given byte offset.+Returns @(message, body, offset of the next message)@. The+end-of-stream marker is reported as 'Left'.+-}+readIPCMessage :: ByteString -> Int -> Either String (Message, ByteString, Int)+readIPCMessage bs !off+  | off < 0 || off > BS.length bs = Left "Arrow.File: message offset out of range"+  | otherwise = do+      let !input = BS.drop off bs+      (msg, bodyLen) <- decodeMessageFrame input+      (_, _, rest) <- readFrameHeader input+      let !metaEnd = off + (BS.length input - BS.length rest)+          !n = fromIntegral bodyLen :: Int+      if bodyLen < 0 || n > BS.length rest+        then Left "Arrow.File: message body runs past the end of the input"+        else+          let !padded = min (BS.length rest) ((n + 7) .&. complement 7)+          in Right (msg, BS.take n rest, metaEnd + padded)+++-- | Read an Arrow IPC stream (schema, dictionary and record batches up to the end-of-stream marker).+readArrowStream :: ByteString -> Either String ArrowStream+readArrowStream bs = do+  (schema, _, batches) <- readArrowStreamFBWithDicts bs+  Right ArrowStream {asSchema = schema, asBatches = V.fromList batches}+++-- | Read an Arrow IPC file.+readArrowFile :: ByteString -> Either String ArrowFile+readArrowFile bs = do+  (schema, _, batches) <- readArrowFileFBWithDicts bs+  Right ArrowFile {afSchema = schema, afBatches = V.fromList batches}+++{- | Read an Arrow IPC file and materialize every record batch+(decompressed, dictionaries resolved).+-}+readArrowFileColumns :: ReadOptions -> ByteString -> Either String (Schema, V.Vector (V.Vector ColumnArray))+readArrowFileColumns opts bs = do+  (schema, batches) <- decodeArrowFile opts bs+  Right (schema, V.fromList batches)
+ src/Arrow/FlatBufferIPC.hs view
@@ -0,0 +1,94 @@+-- | Real Apache Arrow IPC framing — binary-compatible with the+-- reference implementations (arrow-cpp, arrow-rs, pyarrow).+--+-- "Arrow.IPC" uses a simplified flatbuffer-shaped encoding that+-- only self-round-trips; this module constructs Arrow's metadata+-- tables as actual FlatBuffers (per @format/Schema.fbs@ and+-- @format/Message.fbs@) and emits the encapsulated-message framing+-- so pyarrow / arrow-rs / arrow-cpp can consume the output.+--+-- The /generic/ flatbuffer primitives (back-to-front 'Builder',+-- vtable dedup, soffset chains, reader peek + table resolution)+-- live in "FlatBuffers.Builder" and "FlatBuffers.Reader" inside+-- @wireform-flatbuffers@. This module only owns the+-- Arrow-specific layout — the @Schema@, @Field@, @Type@,+-- @RecordBatch@, @Message@, @Tensor@, @SparseTensor@ tables, plus+-- the encapsulated-frame / file-format glue. That split lets+-- "FlatBuffers.Encode" / "FlatBuffers.Decode" stay focused on+-- value-shaped use cases while the spec-precise encoder is+-- shared rather than reimplemented per call site.+--+-- The encoder is standards-compliant:+--+--   * Buffer is built back-to-front.+--   * Tables carry a signed int32 soffset to their vtable at offset 0.+--   * Vtables share when structurally identical (via a deduplication map).+--   * Scalars are aligned to their width; vectors/strings/tables+--     are 4-aligned.+--   * The root offset at byte 0 is an unsigned uoffset_t pointing to+--     the root table.+--+-- The implementation is split across "Arrow.FlatBufferIPC.Write"+-- (encoder), "Arrow.FlatBufferIPC.Read" (decoder) and+-- "Arrow.FlatBufferIPC.Common" (types both sides share); this+-- module re-exports their public API.+module Arrow.FlatBufferIPC+  ( -- * Top-level builders+    buildSchemaMessage+  , buildRecordBatchMessage+    -- * Encapsulated-message framing+  , encapsulateMessage+    -- * Stream / file writers+  , writeArrowStreamFB+  , writeArrowFileFB+  , writeArrowFileFBWithDicts+    -- * Column-based convenience writer+  , buildRecordBatchBytes+  , buildRecordBatchBytesWith+  , writeArrowStreamFBFromColumns+    -- * Body compression helpers+  , compressBody+  , compressBufferEither+  , bodyCompressionAvailable+  , decompressBody+    -- * Reader (parses pyarrow / arrow-cpp output)+  , readArrowStreamFB+  , readArrowFileFB+  , readArrowFileFBWithDicts+  , decodeSchemaMessage+  , decodeRecordBatchMessage+  , decodeDictionaryBatchMessage+  , decodeRecordBatch+  , decodeDictionaryBatch+    -- * Dictionary support+  , DictBatch (..)+  , readArrowStreamFBWithDicts+  , readArrowStreamFBInterleaved+  , StreamFrame (..)+  , buildDictionaryBatchMessage+    -- * File footer index (sections)+  , ArrowBlock (..)+  , FileIndex (..)+  , readFileIndex+  , parseFileIndex+    -- * Tensor / SparseTensor+  , Tensor (..)+  , TensorDim (..)+  , buildTensorMessage+  , decodeTensorMessage+  , encodeTensorFrame+  , decodeTensorFrame+  , SparseTensor (..)+  , buildSparseTensorMessageCOO+  , encodeSparseTensorFrame+  , decodeSparseTensorFrame+  , writeArrowStreamFBWithDicts+    -- * Single-message frames+  , encodeMessageFrame+  , decodeMessageFrame+  , readFrameHeader+  ) where++import Arrow.FlatBufferIPC.Common+import Arrow.FlatBufferIPC.Read+import Arrow.FlatBufferIPC.Write
+ src/Arrow/FlatBufferIPC/Common.hs view
@@ -0,0 +1,122 @@+{-# LANGUAGE StrictData #-}+-- | Data types shared by the Arrow IPC encoder+-- ("Arrow.FlatBufferIPC.Write") and decoder+-- ("Arrow.FlatBufferIPC.Read"): dictionary batches and the+-- tensor / sparse-tensor message payloads.+module Arrow.FlatBufferIPC.Common+  ( DictBatch (..)+  , ArrowBlock (..)+  , Tensor (..)+  , TensorDim (..)+  , SparseTensor (..)+  ) where++import Data.ByteString (ByteString)+import Data.Int (Int32, Int64)+import qualified Data.Text as T+import qualified Data.Vector as V++import Arrow.Types (ArrowType, RecordBatchDef)++-- ============================================================+-- Footer blocks+-- ============================================================++{- | A @Block@ struct as emitted in the @Footer.dictionaries@ /+@Footer.recordBatches@ vectors (@format/File.fbs@). Inline+fixed-size struct (24 bytes total): @offset: i64@,+@metaDataLength: i32@ + 4 bytes padding, @bodyLength: i64@. Each+block points at its message's continuation marker, with the whole+header length (marker, length, metadata, padding) and the padded+body length.+-}+data ArrowBlock = ArrowBlock+  { abOffset  :: !Int64+  , abMetaLen :: !Int32+  , abBodyLen :: !Int64+  } deriving (Show, Eq)++-- ============================================================+-- Tensor / SparseTensor messages+-- ============================================================++-- | Dense tensor metadata. Mirrors @table Tensor@ from Arrow's+-- @format/Tensor.fbs@. @tensorBody@ is the raw buffer carrying+-- @product(shape) * bytesPerElement(tensorType)@ little-endian+-- elements in row-major order (unless 'tensorStrides' is set).+data Tensor = Tensor+  { tensorType    :: !ArrowType+  , tensorShape   :: !(V.Vector TensorDim)+  , tensorStrides :: !(V.Vector Int64)+  , tensorBody    :: !ByteString+  } deriving (Show, Eq)++-- | One entry of a tensor's shape vector. Name is optional+-- (empty string by default).+data TensorDim = TensorDim+  { tdSize :: !Int64+  , tdName :: !T.Text+  } deriving (Show, Eq)++-- ============================================================+-- SparseTensor+-- ============================================================+--+-- Arrow's SparseTensor union covers four index formats (COO,+-- CSR, CSC, CSF). We model the common case — SparseTensorIndexCOO+-- — explicitly; callers with other formats can drop down to the+-- raw flatbuffer builder. A SparseTensor message (header_type=5)+-- wraps:+--+-- @+-- table SparseTensor {+--   type_type: Type;           // slot 0+--   type:      Type;           // slot 1+--   shape:    [TensorDim];     // slot 2+--   non_zero_length: long;     // slot 3+--   sparseIndex_type: SparseTensorIndex;  // slot 4 (union tag)+--   sparseIndex: SparseTensorIndex;       // slot 5+--   data: Buffer;              // slot 6+-- }+-- @+--+-- For COO:+--+-- @+-- table SparseTensorIndexCOO {+--   indicesType: Int;       // slot 0+--   indicesStrides: [long]; // slot 1+--   indicesBuffer: Buffer;  // slot 2+--   isCanonical: bool;      // slot 3+-- }+-- @++-- | A sparse tensor in coordinate (COO) layout. @indicesType@+-- is the integer width of each coordinate; @indicesBuffer@+-- carries @non_zero_length * ndim@ integers; @tensorBody@+-- carries @non_zero_length@ values of @tensorType@.+data SparseTensor = SparseTensor+  { sparseTensorType    :: !ArrowType+  , sparseTensorShape   :: !(V.Vector TensorDim)+  , sparseNonZeroLength :: !Int64+  , sparseIndicesType   :: !ArrowType+    -- ^ typically @AInt 64 True@ or @AInt 32 True@+  , sparseIndicesBody   :: !ByteString+  , sparseIndicesCanonical :: !Bool+  , sparseTensorBody    :: !ByteString+  } deriving (Show, Eq)++-- | One decoded dictionary batch — the raw payload that defines a+-- dictionary's index → value mapping. The @data@ field is the+-- inner @RecordBatch@ (a single column whose values are the+-- dictionary values, in index order).+data DictBatch = DictBatch+  { dbId      :: !Int64+    -- ^ Dictionary id; matches @DictionaryEncoding.id@ in the+    -- schema.+  , dbIsDelta :: !Bool+    -- ^ When @True@, the values append to the existing dictionary+    -- with this id; otherwise they replace it.+  , dbData    :: !RecordBatchDef+  , dbBody    :: !ByteString+  } deriving stock (Show, Eq)
+ src/Arrow/FlatBufferIPC/Read.hs view
@@ -0,0 +1,1022 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE StrictData #-}+{-# LANGUAGE TypeApplications #-}+-- | Arrow IPC decoder: parses the @Schema@, @Field@, @Type@,+-- @RecordBatch@, @DictionaryBatch@, @Tensor@ and @SparseTensor@+-- flatbuffer tables emitted by pyarrow / arrow-rs / arrow-cpp,+-- walks the encapsulated-message framing of streams and files,+-- and undoes body compression. See "Arrow.FlatBufferIPC" for the+-- format notes.+module Arrow.FlatBufferIPC.Read+  ( -- * Body compression+    decompressBody+    -- * Reader (parses pyarrow / arrow-cpp output)+  , readArrowStreamFB+  , readArrowFileFB+  , readArrowFileFBWithDicts+  , decodeSchemaMessage+  , decodeRecordBatchMessage+  , decodeDictionaryBatchMessage+  , decodeRecordBatch+  , decodeDictionaryBatch+    -- * Dictionary support+  , readArrowStreamFBWithDicts+  , readArrowStreamFBInterleaved+  , StreamFrame (..)+    -- * File footer index+  , FileIndex (..)+  , readFileIndex+  , parseFileIndex+    -- * Tensor / SparseTensor+  , decodeTensorMessage+  , decodeTensorFrame+  , decodeSparseTensorFrame+    -- * Single-message frames+  , decodeMessageFrame+  , readFrameHeader+  ) where++import Data.ByteString (ByteString)+import qualified Data.ByteString as BS+import Data.Int (Int64)+import qualified Data.Text as T+import qualified Data.Text.Encoding as TE+import qualified Data.Vector as V+import qualified Data.Vector.Storable as VS+import Foreign.Storable (Storable)++import Arrow.Column.Internal (bytesToStorable)+import Arrow.FlatBufferIPC.Common+import Arrow.Read.Columns (decodeDictionaryBatch, decodeRecordBatch, decompressBody)+import Arrow.Types+import FlatBuffers.Builder (alignUp)+import FlatBuffers.Reader+  ( Pos+  , followUOffset+  , peekI16+  , peekI32+  , peekI64+  , peekU8+  , peekU32+  , readString+  , readStringSlice+  , readVectorInt64+  , readVectorOfOffsets+  , readVectorOfStructs+  , resolveTable+  , vectorLength+  )++-- | Mini-helper: run @e@ when the condition holds, else 'Right' '()'.+-- (This is shaped like 'Control.Monad.when' but specialised to+-- 'Either String' and lifted to a top-level binding so all the+-- decoders below can share one definition.)+when' :: Bool -> Either String () -> Either String ()+when' True  e = e+when' False _ = Right ()++-- | Slice @len@ bytes at @off@ out of @bs@, where both come off the+-- wire: negative or out-of-range values are a 'Left', never a+-- silently clamped slice.+sliceWire :: String -> Int64 -> Int64 -> ByteString -> Either String ByteString+sliceWire what off len bs+  | off < 0 || len < 0 || len > n || off > n - len =+      Left+        ( "Arrow.FlatBufferIPC: "+            ++ what+            ++ " [offset "+            ++ show off+            ++ ", length "+            ++ show len+            ++ "] lies outside the "+            ++ show n+            ++ " available bytes"+        )+  | otherwise = Right $! BS.take (fromIntegral len) (BS.drop (fromIntegral off) bs)+  where+    n = fromIntegral (BS.length bs) :: Int64++-- | Parse a @Tensor@ message header (flatbuffer metadata only;+-- the body bytes are the caller's concern — they sit in the+-- encapsulated frame's body section).+decodeTensorMessage+  :: ByteString -> Either String (Tensor, Int64)+decodeTensorMessage meta = do+  msgPos <- fromIntegral <$> peekU32 meta 0+  mSlot  <- resolveTable meta msgPos+  ht <- case mSlot 1 of+    Nothing -> Right (0 :: Int)+    Just b  -> fromIntegral <$> peekU8 meta b+  when' (ht /= 4) $+    Left ("Arrow.FlatBufferIPC.decodeTensorMessage: expected header_type=4, got "+          ++ show ht)+  headerPos <- case mSlot 2 of+    Nothing  -> Left "Tensor message missing header slot"+    Just off -> Right off+  tensorPos <- followUOffset meta headerPos+  tSlot     <- resolveTable meta tensorPos+  tyTag <- case tSlot 0 of+    Nothing -> Left "Tensor missing type_type"+    Just b  -> fromIntegral <$> peekU8 meta b+  tyFieldPos <- case tSlot 1 of+    Nothing  -> Left "Tensor missing type"+    Just off -> Right off+  tyPos <- followUOffset meta tyFieldPos+  arrowTy <- readType meta tyTag (Just tyPos)+  shape <- case tSlot 2 of+    Nothing  -> Right V.empty+    Just off -> do+      shapePos <- followUOffset meta off+      dimOffs  <- readVectorOfOffsets meta shapePos+      V.mapM (\dimPos -> readTensorDim meta dimPos) dimOffs+  strides <- case tSlot 3 of+    Nothing  -> Right V.empty+    Just off -> do+      sPos <- followUOffset meta off+      V.fromList <$> readVectorInt64 meta sPos+  bodyLen <- case mSlot 3 of+    Nothing -> Right 0+    Just b  -> peekI64 meta b+  Right (Tensor arrowTy shape strides BS.empty, bodyLen)++-- | Parse a SparseTensor message (header_type=5) carrying a COO+-- index. Returns the decoded sparse tensor with both its+-- @sparseIndicesBody@ and @sparseTensorBody@ sliced out of the+-- frame's body bytes, and the remainder of the input.+decodeSparseTensorFrame+  :: ByteString -> Either String (SparseTensor, ByteString)+decodeSparseTensorFrame bs = do+  (mlen, meta, rest1) <- readFrameHeader bs+  when' (mlen <= 0) $ Left "decodeSparseTensorFrame: unexpected EOS"+  msgPos <- fromIntegral <$> peekU32 meta 0+  mSlot  <- resolveTable meta msgPos+  ht <- case mSlot 1 of+    Nothing -> Right (0 :: Int)+    Just b  -> fromIntegral <$> peekU8 meta b+  when' (ht /= 5) $+    Left ("decodeSparseTensorFrame: expected header_type=5, got " ++ show ht)+  hdrPos <- case mSlot 2 of+    Nothing -> Left "SparseTensor missing header"+    Just p  -> Right p+  stPos  <- followUOffset meta hdrPos+  sSlot  <- resolveTable meta stPos+  tyTag  <- case sSlot 0 of+    Nothing -> Left "SparseTensor missing type_type"+    Just b  -> fromIntegral <$> peekU8 meta b+  tyFieldPos <- case sSlot 1 of+    Nothing -> Left "SparseTensor missing type"+    Just p  -> Right p+  tyPos <- followUOffset meta tyFieldPos+  arrowTy <- readType meta tyTag (Just tyPos)+  shape <- case sSlot 2 of+    Nothing  -> Right V.empty+    Just off -> do+      shapePos <- followUOffset meta off+      dimOffs  <- readVectorOfOffsets meta shapePos+      V.mapM (\dimPos -> readTensorDim meta dimPos) dimOffs+  nnz <- case sSlot 3 of+    Nothing -> Right 0+    Just p  -> peekI64 meta p+  idxTag <- case sSlot 4 of+    Nothing -> Left "SparseTensor missing sparseIndex_type"+    Just p  -> fromIntegral <$> peekU8 meta p+  when' (idxTag /= 1) $+    Left ("SparseTensor: only COO index format supported by this decoder, got tag=" ++ show (idxTag :: Int))+  idxFieldPos <- case sSlot 5 of+    Nothing -> Left "SparseTensor missing sparseIndex"+    Just p  -> Right p+  cooPos <- followUOffset meta idxFieldPos+  cSlot  <- resolveTable meta cooPos+  iTyFieldPos <- case cSlot 0 of+    Nothing -> Left "SparseTensorIndexCOO missing indicesType"+    Just p  -> Right p+  iTyPos <- followUOffset meta iTyFieldPos+  idxArrowTy <- readType meta 2 (Just iTyPos)  -- always an Int table+  (idxOffset, idxLen) <- case cSlot 2 of+    Nothing -> Left "SparseTensorIndexCOO missing indicesBuffer"+    Just p  -> do+      lo  <- peekI64 meta p+      ln' <- peekI64 meta (p + 8)+      Right (lo, ln')+  canonical <- case cSlot 3 of+    Nothing -> Right False+    Just p  -> (/= 0) <$> peekU8 meta p+  (valOffset, valLen) <- case sSlot 6 of+    Nothing -> Left "SparseTensor missing data buffer"+    Just p  -> do+      lo <- peekI64 meta p+      ln <- peekI64 meta (p + 8)+      Right (lo, ln)+  iSlice <- sliceWire "sparse tensor indices buffer" idxOffset idxLen rest1+  vSlice <- sliceWire "sparse tensor data buffer" valOffset valLen rest1+  -- Both slices were range-checked above, so this sum cannot overflow.+  let !bodyLen = fromIntegral valOffset + fromIntegral valLen :: Int+      rest2   = BS.drop (alignUp bodyLen 8) rest1+  Right ( SparseTensor+            { sparseTensorType       = arrowTy+            , sparseTensorShape      = shape+            , sparseNonZeroLength    = nnz+            , sparseIndicesType      = idxArrowTy+            , sparseIndicesBody      = iSlice+            , sparseIndicesCanonical = canonical+            , sparseTensorBody       = vSlice+            }+        , rest2+        )++-- | Parse an Arrow IPC frame that carries a 'Tensor'. Returns+-- the decoded tensor (with 'tensorBody' filled in from the+-- frame's body bytes) and the remaining bytes. Fails if the+-- frame's header_type isn't 4 (Tensor).+decodeTensorFrame+  :: ByteString -> Either String (Tensor, ByteString)+decodeTensorFrame bs = do+  (mlen, meta, rest1) <- readFrameHeader bs+  when' (mlen <= 0) $ Left "decodeTensorFrame: unexpected EOS"+  (t, bodyLen) <- decodeTensorMessage meta+  (body, rest2) <- takeFrameBody "tensor body" bodyLen rest1+  Right (t { tensorBody = body }, rest2)++-- | Split a frame's @bodyLen@-byte body (a wire value) off the bytes+-- that follow its metadata, also dropping the body's padding to the+-- next 8-byte boundary.+takeFrameBody :: String -> Int64 -> ByteString -> Either String (ByteString, ByteString)+takeFrameBody what bodyLen rest = do+  body <- sliceWire what 0 bodyLen rest+  Right (body, BS.drop (alignUp (BS.length body) 8) rest)++-- | Parse a @TensorDim@ table at the given position.+readTensorDim :: ByteString -> Int -> Either String TensorDim+readTensorDim meta pos = do+  slot <- resolveTable meta pos+  size <- case slot 0 of+    Nothing -> Right 0+    Just b  -> peekI64 meta b+  name <- case slot 1 of+    Nothing  -> Right T.empty+    Just off -> do+      sPos <- followUOffset meta off+      readString meta sPos+  Right (TensorDim size name)++-- | Decode a complete @Schema@ table at the given position.+readSchemaTable :: ByteString -> Pos -> Either String Schema+readSchemaTable bs schPos = do+  slot <- resolveTable bs schPos+  endian <- case slot 0 of+    Nothing -> Right Little+    Just p  -> do+      v <- peekI16 bs p+      case v of+        0 -> Right Little+        1 -> Right Big+        _ -> Left ("Arrow.FlatBufferIPC: unknown endianness " ++ show v)+  (fieldsVec, budget0) <- case slot 1 of+    Nothing -> Right (V.empty, schemaBudget bs)+    Just p  -> readChargedOffsets bs (schemaBudget bs) p+  (fields, budget1) <- readFields bs 0 budget0 fieldsVec+  (customMd, _) <- readKeyValues bs budget1 (slot 2)+  features <- case slot 3 of+    Nothing -> Right V.empty+    Just p -> do+      vecPos <- followUOffset bs p+      (_, elems) <- readVectorOfStructs bs vecPos 8+      V.mapM (\off -> peekI64 bs off >>= readFeature) elems+  Right Schema+    { arrowFields     = fields+    , arrowEndianness = endian+    , arrowMetadata   = customMd+    , arrowFeatures   = features+    }++-- | Unknown features may change stream semantics, so they cannot be silently+-- discarded as if this implementation supported them. Values are Int64,+-- as specified by Schema.fbs, rather than the usual byte-sized enums.+readFeature :: Int64 -> Either String Feature+readFeature value = case value of+  0 -> Right FeatureUnused+  1 -> Right FeatureDictionaryReplacement+  2 -> Right FeatureCompressedBody+  _ -> Left ("Arrow.FlatBufferIPC: unsupported schema feature " ++ show value)++-- | Nesting limit for decoded schemas (arrow-cpp's IPC reader uses+-- the same bound).+maxSchemaDepth :: Int+maxSchemaDepth = 64++-- | FlatBuffers offsets may alias, so a few hundred bytes can+-- describe an exponentially large DAG of fields or reuse one long+-- string for every name. Every decoded field, string and type id is+-- charged against this budget; an honest (unaliased) buffer spends+-- at most its own length.+schemaBudget :: ByteString -> Int+schemaBudget bs = 4 * BS.length bs + 4096++-- | Charge @cost@ against the remaining schema budget.+charge :: Int -> Int -> Either String Int+charge cost budget+  | cost > budget =+      Left "Arrow.FlatBufferIPC: schema decodes to more fields and strings than its metadata holds (aliased offsets)"+  | otherwise = Right (budget - cost)++-- | Decode the string at the uoffset in @fieldPos@, charging its length first.+readChargedString :: ByteString -> Int -> Pos -> Either String (T.Text, Int)+readChargedString bs budget fieldPos = do+  strPos <- followUOffset bs fieldPos+  raw <- readStringSlice bs strPos+  budget' <- charge (4 + BS.length raw) budget+  case TE.decodeUtf8' raw of+    Left _ -> Left "Arrow.FlatBufferIPC: invalid UTF-8 in schema string"+    Right t -> Right (t, budget')++-- | Read the @[offset]@ vector at the uoffset in @fieldPos@, charging+-- four units per element before the position vector is allocated.+readChargedOffsets :: ByteString -> Int -> Pos -> Either String (V.Vector Pos, Int)+readChargedOffsets bs budget fieldPos = do+  vecPos <- followUOffset bs fieldPos+  n <- vectorLength bs vecPos+  budget' <- charge (4 * n) budget+  ps <- readVectorOfOffsets bs vecPos+  Right (ps, budget')++-- | 'readType' with its variable-size parts (union type ids, the+-- timestamp zone string) charged before they are decoded.+readTypeCharged :: ByteString -> Int -> Int -> Maybe Pos -> Either String (ArrowType, Int)+readTypeCharged bs budget tag mpos = do+  cost <- case (tag, mpos) of+    (10, Just p) -> slotCost p (\vp -> (\s -> 4 + BS.length s) <$> readStringSlice bs vp)+    (14, Just p) -> slotCost p (\vp -> (* 4) <$> vectorLength bs vp)+    _ -> Right 0+  budget' <- charge cost budget+  ty <- readType bs tag mpos+  Right (ty, budget')+  where+    -- Both variable parts live behind slot 1 of their type table.+    slotCost p measure = do+      s <- resolveTable bs p+      case s 1 of+        Nothing -> Right 0+        Just fp -> followUOffset bs fp >>= measure++readFields :: ByteString -> Int -> Int -> V.Vector Pos -> Either String (V.Vector Field, Int)+readFields bs depth budget0 ps = go 0 budget0 []+  where+    !n = V.length ps+    go !i !budget acc+      | i >= n = Right (V.fromListN n (reverse acc), budget)+      | otherwise = do+          (f, budget') <- readField bs depth budget (V.unsafeIndex ps i)+          go (i + 1) budget' (f : acc)++-- | Decode one @Field@ table (and its children), returning the+-- remaining schema budget.+readField :: ByteString -> Int -> Int -> Pos -> Either String (Field, Int)+readField bs depth budget0 fldPos = do+  when' (depth >= maxSchemaDepth) $+    Left ("Arrow.FlatBufferIPC: schema nests deeper than " ++ show maxSchemaDepth ++ " levels")+  budget1 <- charge 8 budget0+  slot <- resolveTable bs fldPos+  (name, budget2) <- case slot 0 of+    Nothing -> Right ("", budget1)+    Just p  -> readChargedString bs budget1 p+  nullable <- case slot 1 of+    Nothing -> Right False+    Just p  -> do+      v <- peekU8 bs p+      Right (v /= 0)+  tyTag <- case slot 2 of+    Nothing -> Right 0+    Just p  -> peekU8 bs p+  (ty, budget3) <- case slot 3 of+    Nothing  -> readTypeCharged bs budget2 (fromIntegral tyTag) Nothing+    Just p   -> do+      tyPos <- followUOffset bs p+      readTypeCharged bs budget2 (fromIntegral tyTag) (Just tyPos)+  dictionary <- case slot 4 of+    Nothing -> Right Nothing+    Just p  -> do+      dePos <- followUOffset bs p+      Just <$> readDictionaryEncodingTable bs dePos+  (customMd, budget4) <- readKeyValues bs budget3 (slot 6)+  (children, budget5) <- case slot 5 of+    Nothing -> Right (V.empty, budget4)+    Just p  -> do+      (childPositions, budget4') <- readChargedOffsets bs budget4 p+      readFields bs (depth + 1) budget4' childPositions+  Right+    ( Field+        { fieldName     = name+        , fieldNullable = nullable+        , fieldType     = ty+        , fieldChildren = children+        , fieldDictionary = dictionary+        , fieldMetadata = customMd+        }+    , budget5+    )++-- | Decode an optional @[KeyValue]@ vector slot, charging each entry.+readKeyValues :: ByteString -> Int -> Maybe Pos -> Either String (V.Vector (T.Text, T.Text), Int)+readKeyValues _ budget Nothing = Right (V.empty, budget)+readKeyValues bs budget00 (Just p) = do+  (kvPositions, budget0) <- readChargedOffsets bs budget00 p+  let !n = V.length kvPositions+      go !i !budget acc+        | i >= n = Right (V.fromListN n (reverse acc), budget)+        | otherwise = do+            (kv, budget') <- readKeyValue bs budget (V.unsafeIndex kvPositions i)+            go (i + 1) budget' (kv : acc)+  go 0 budget0 []++-- | Decode one @KeyValue@ table (per @format/Schema.fbs@):+--+-- @+-- table KeyValue {+--   key   : string;   // 0+--   value : string;   // 1+-- }+-- @+readKeyValue :: ByteString -> Int -> Pos -> Either String ((T.Text, T.Text), Int)+readKeyValue bs budget0 kvPos = do+  budget1 <- charge 8 budget0+  slot <- resolveTable bs kvPos+  (k, budget2) <- case slot 0 of+    Nothing -> Right ("", budget1)+    Just p  -> readChargedString bs budget1 p+  (v, budget3) <- case slot 1 of+    Nothing -> Right ("", budget2)+    Just p  -> readChargedString bs budget2 p+  Right ((k, v), budget3)++-- | Decode a 'DictionaryEncoding' table:+--+-- @+-- table DictionaryEncoding {+--   id: long;+--   indexType: Int;+--   isOrdered: bool;+--   dictionaryKind: DictionaryKind;+-- }+-- @+readDictionaryEncodingTable :: ByteString -> Pos -> Either String DictionaryEncoding+readDictionaryEncodingTable bs dePos = do+  s <- resolveTable bs dePos+  did <- case s 0 of+    Nothing -> Right 0+    Just b  -> peekI64 bs b+  idxTy <- case s 1 of+    Nothing -> Right (AInt 32 True)   -- spec default+    Just b  -> do+      tyPos <- followUOffset bs b+      readType bs 2 (Just tyPos)+  ordered <- case s 2 of+    Nothing -> Right False+    Just b  -> do+      v <- peekU8 bs b+      Right (v /= 0)+  Right (DictionaryEncoding did idxTy ordered)++-- | Decode a @Type@ union variant. The discriminator (@type_type@)+-- selects which sub-table layout to read at @typePos@.+readType :: ByteString -> Int -> Maybe Pos -> Either String ArrowType+readType _  0 _ = Right ANull   -- "None" / Null+readType _  1 _ = Right ANull+readType bs 2 (Just p) = do+  -- Int { bitWidth: i32, is_signed: bool }. Schema.fbs gives+  -- is_signed no explicit default, so an absent slot is the bool+  -- default (false); pyarrow omits it for unsigned columns.+  s <- resolveTable bs p+  bits <- case s 0 of+    Nothing -> Right 32+    Just b  -> peekI32 bs b+  signed <- case s 1 of+    Nothing -> Right False+    Just b  -> do+      v <- peekU8 bs b+      Right (v /= 0)+  Right (AInt (fromIntegral bits) signed)+readType bs 3 (Just p) = do+  -- FloatingPoint { precision: Precision }; an absent slot is the+  -- enum's first value, HALF (pyarrow omits it for float16).+  s <- resolveTable bs p+  prec <- case s 0 of+    Nothing -> Right 0+    Just b  -> peekI16 bs b+  case prec of+    0 -> Right (AFloatingPoint Half)+    1 -> Right (AFloatingPoint Single)+    2 -> Right (AFloatingPoint DoublePrecision)+    n -> Left $ "Arrow.FlatBufferIPC: unknown precision " ++ show n+readType _  4 _ = Right ABinary+readType _  5 _ = Right AUtf8+readType _  6 _ = Right ABool+readType bs 7 (Just p) = do+  s <- resolveTable bs p+  prec  <- case s 0 of { Nothing -> Right 0; Just b -> peekI32 bs b }+  scale <- case s 1 of { Nothing -> Right 0; Just b -> peekI32 bs b }+  bw    <- case s 2 of { Nothing -> Right 128; Just b -> peekI32 bs b }+  decimal <- case bw of+      32 -> Right (ADecimal32 (fromIntegral prec) (fromIntegral scale))+      64 -> Right (ADecimal64 (fromIntegral prec) (fromIntegral scale))+      128 -> Right (ADecimal (fromIntegral prec) (fromIntegral scale))+      256 -> Right (ADecimal256 (fromIntegral prec) (fromIntegral scale))+      n   -> Left $ "Arrow.FlatBufferIPC: unsupported decimal bitWidth " ++ show n+  validateDecimalType decimal+  Right decimal+readType bs 8 (Just p) = do+  s <- resolveTable bs p+  u <- case s 0 of { Nothing -> Right 1; Just b -> peekI16 bs b }+  case u of+    0 -> Right (ADate DateDay)+    1 -> Right (ADate DateMillisecond)+    n -> Left $ "Arrow.FlatBufferIPC: unknown date unit " ++ show n+readType bs 9 (Just p) = do+  s <- resolveTable bs p+  u  <- case s 0 of { Nothing -> Right 1;  Just b -> peekI16 bs b }+  bw <- case s 1 of { Nothing -> Right 32; Just b -> peekI32 bs b }+  unit <- timeUnitFromTag (fromIntegral u)+  Right (ATime unit (fromIntegral bw))+readType bs 10 (Just p) = do+  s <- resolveTable bs p+  u  <- case s 0 of { Nothing -> Right 0; Just b -> peekI16 bs b }+  tz <- case s 1 of+    Nothing -> Right Nothing+    Just b  -> do+      strPos <- followUOffset bs b+      Just <$> readString bs strPos+  unit <- timeUnitFromTag (fromIntegral u)+  Right (ATimestamp unit tz)+readType bs 11 (Just p) = do+  s <- resolveTable bs p+  u <- case s 0 of { Nothing -> Right 0; Just b -> peekI16 bs b }+  iu <- case u of+    0 -> Right YearMonth+    1 -> Right DayTime+    2 -> Right MonthDayNano+    n -> Left $ "Arrow.FlatBufferIPC: unknown interval unit " ++ show n+  Right (AInterval iu)+readType _  12 _ = Right AList+readType _  13 _ = Right AStruct+readType bs 14 (Just p) = do+  s <- resolveTable bs p+  m <- case s 0 of { Nothing -> Right 0; Just b -> peekI16 bs b }+  mode <- case m of+    0 -> Right Sparse+    1 -> Right Dense+    n -> Left $ "Arrow.FlatBufferIPC: unknown union mode " ++ show n+  ids <- case s 1 of+    Nothing -> Right V.empty+    Just b  -> do+      vecPos <- followUOffset bs b+      (_, elems) <- readVectorOfStructs bs vecPos 4+      V.mapM (peekI32 bs) elems+  Right (AUnion mode ids)+readType bs 15 (Just p) = do+  s <- resolveTable bs p+  bw <- case s 0 of { Nothing -> Right 0; Just b -> peekI32 bs b }+  Right (AFixedSizeBinary (fromIntegral bw))+readType bs 16 (Just p) = do+  s <- resolveTable bs p+  ls <- case s 0 of { Nothing -> Right 0; Just b -> peekI32 bs b }+  Right (AFixedSizeList (fromIntegral ls))+readType bs 17 (Just p) = do+  s <- resolveTable bs p+  sorted <- case s 0 of+    Nothing -> Right False+    Just b  -> do+      v <- peekU8 bs b+      Right (v /= 0)+  Right (AMap sorted)+readType bs 18 (Just p) = do+  s <- resolveTable bs p+  u <- case s 0 of { Nothing -> Right 1; Just b -> peekI16 bs b }+  unit <- timeUnitFromTag (fromIntegral u)+  Right (ADuration unit)+readType _  19 _ = Right ALargeBinary+readType _  20 _ = Right ALargeUtf8+readType _  21 _ = Right ALargeList+readType _  22 _ = Right ARunEndEncoded+readType _  23 _ = Right ABinaryView+readType _  24 _ = Right AUtf8View+readType _  25 _ = Right AListView+readType _  26 _ = Right ALargeListView+readType _  n  _ = Left $ "Arrow.FlatBufferIPC: unsupported Type discriminator " ++ show n++timeUnitFromTag :: Int -> Either String TimeUnit+timeUnitFromTag 0 = Right Second+timeUnitFromTag 1 = Right Millisecond+timeUnitFromTag 2 = Right Microsecond+timeUnitFromTag 3 = Right Nanosecond+timeUnitFromTag n = Left $ "Arrow.FlatBufferIPC: unknown time unit " ++ show n++-- | The vector of 16-byte structs (@FieldNode@ or @Buffer@) at the+-- uoffset in @slotPos@: one length check, then an O(1) alias of the+-- payload (copied once only when it is not 8-aligned in memory).+structVector :: Storable a => ByteString -> Pos -> Either String (VS.Vector a)+structVector bs slotPos = do+  vecPos <- followUOffset bs slotPos+  n <- vectorLength bs vecPos+  -- vectorLength succeeded, so 0 <= vecPos <= length - 4.+  if n > (BS.length bs - vecPos - 4) `quot` 16+    then Left ("Arrow.FlatBufferIPC: record batch vector of " ++ show n ++ " structs overruns the message")+    else Right $! bytesToStorable (BS.take (n * 16) (BS.drop (vecPos + 4) bs))++-- | Decode a @RecordBatch@ table.+readRecordBatchTable :: ByteString -> Pos -> Either String RecordBatchDef+readRecordBatchTable bs rbPos = do+  s <- resolveTable bs rbPos+  len <- case s 0 of+    Nothing -> Right 0+    Just b  -> peekI64 bs b+  nodes <- maybe (Right VS.empty) (structVector bs) (s 1)+  bufs <- maybe (Right VS.empty) (structVector bs) (s 2)+  variadic <- case s 4 of+    Nothing -> Right V.empty+    Just b  -> do+      vecPos <- followUOffset bs b+      V.fromList <$> readVectorInt64 bs vecPos+  -- Slot 3 is BodyCompression (a table). When present we read+  -- the codec discriminator and translate to our enum.+  bodyComp <- case s 3 of+    Nothing -> Right Nothing+    Just b  -> do+      bcPos <- followUOffset bs b+      bcSlot <- resolveTable bs bcPos+      codec <- case bcSlot 0 of+        Nothing -> Right (0 :: Int)+        Just p  -> do+          v <- peekU8 bs p+          Right (fromIntegral v)+      case codec of+        0 -> Right (Just LZ4Frame)+        1 -> Right (Just BodyZstd)+        n -> Left $ "Arrow.FlatBufferIPC: unknown BodyCompression codec " ++ show n+  Right RecordBatchDef+    { rbLength  = len+    , rbNodes   = nodes+    , rbBuffers = bufs+    , rbVariadicBufferCounts = variadic+    , rbBodyCompression = bodyComp+    }++-- | Decode a Schema-typed @Message@ flatbuffer (just the metadata+-- bytes; the caller has already stripped the encapsulated framing).+decodeSchemaMessage :: ByteString -> Either String Schema+decodeSchemaMessage meta = do+  msgPos <- fromIntegral <$> peekU32 meta 0+  s <- resolveTable meta msgPos+  ht <- case s 1 of+    Nothing -> Left "Arrow.FlatBufferIPC: Message header_type missing"+    Just b  -> peekU8 meta b+  when' (ht /= 1) $+    Left ("Arrow.FlatBufferIPC: expected Schema header (1), got " ++ show ht)+  case s 2 of+    Nothing -> Left "Arrow.FlatBufferIPC: Message header (Schema) missing"+    Just b  -> do+      schPos <- followUOffset meta b+      readSchemaTable meta schPos++-- | Decode a RecordBatch-typed @Message@ flatbuffer to+-- @(RecordBatchDef, bodyLength)@.+decodeRecordBatchMessage :: ByteString -> Either String (RecordBatchDef, Int64)+decodeRecordBatchMessage meta = do+  msgPos <- fromIntegral <$> peekU32 meta 0+  s <- resolveTable meta msgPos+  ht <- case s 1 of+    Nothing -> Left "Arrow.FlatBufferIPC: Message header_type missing"+    Just b  -> peekU8 meta b+  when' (ht /= 3) $+    Left ("Arrow.FlatBufferIPC: expected RecordBatch header (3), got " ++ show ht)+  rb <- case s 2 of+    Nothing -> Left "Arrow.FlatBufferIPC: Message header (RecordBatch) missing"+    Just b  -> do+      rbPos <- followUOffset meta b+      readRecordBatchTable meta rbPos+  bodyLen <- case s 3 of+    Nothing -> Right 0+    Just b  -> peekI64 meta b+  Right (rb, bodyLen)++-- | Decode a DictionaryBatch-typed @Message@ flatbuffer to+-- @(id, isDelta, RecordBatchDef, bodyLength)@.+decodeDictionaryBatchMessage+  :: ByteString -> Either String (Int64, Bool, RecordBatchDef, Int64)+decodeDictionaryBatchMessage meta = do+  msgPos <- fromIntegral <$> peekU32 meta 0+  s <- resolveTable meta msgPos+  ht <- case s 1 of+    Nothing -> Left "Arrow.FlatBufferIPC: Message header_type missing"+    Just b  -> peekU8 meta b+  when' (ht /= 2) $+    Left ("Arrow.FlatBufferIPC: expected DictionaryBatch header (2), got " ++ show ht)+  dbTblPos <- case s 2 of+    Nothing -> Left "Arrow.FlatBufferIPC: Message header (DictionaryBatch) missing"+    Just b  -> followUOffset meta b+  ds <- resolveTable meta dbTblPos+  did <- case ds 0 of+    Nothing -> Right 0+    Just b  -> peekI64 meta b+  rb <- case ds 1 of+    Nothing -> Left "Arrow.FlatBufferIPC: DictionaryBatch.data missing"+    Just b  -> do+      rbPos <- followUOffset meta b+      readRecordBatchTable meta rbPos+  isDelta <- case ds 2 of+    Nothing -> Right False+    Just b  -> do+      v <- peekU8 meta b+      Right (v /= 0)+  bodyLen <- case s 3 of+    Nothing -> Right 0+    Just b  -> peekI64 meta b+  Right (did, isDelta, rb, bodyLen)++-- | Parse an Arrow IPC stream produced by any spec-compliant+-- writer (pyarrow / arrow-cpp / arrow-rs) into wireform's+-- 'Schema' + a list of @(RecordBatchDef, body bytes)@ pairs.+--+-- Recognises both the post-0.15.0 framing (continuation marker ++-- length) and the legacy framing (positive length first, no+-- continuation), per the @ConsumeInitial@ logic in arrow-cpp's+-- @message.cc@.+readArrowStreamFB+  :: ByteString+  -> Either String (Schema, [(RecordBatchDef, ByteString)])+readArrowStreamFB bs0 = do+  (sch, _, batches) <- readArrowStreamFBWithDicts bs0+  Right (sch, batches)++-- | Frame variant for 'readArrowStreamFBInterleaved'. Preserves+-- the stream-order sequence of dict + record messages so a+-- downstream reader can honour replacement / delta dictionary+-- semantics.+data StreamFrame+  = SFDict  !DictBatch+  | SFBatch !RecordBatchDef !ByteString+  deriving (Show, Eq)++-- | Like 'readArrowStreamFBWithDicts' but returns the dict /+-- record batches /interleaved/ in stream order. Use this when+-- a writer may emit @isDelta=false@ replacement dict batches+-- between record batches: the flat list gives each record batch+-- the opportunity to resolve against the most-recently-seen dict+-- for that id.+readArrowStreamFBInterleaved+  :: ByteString+  -> Either String (Schema, [StreamFrame])+readArrowStreamFBInterleaved bs0 = do+  (schema, after) <- consumeSchema bs0+  frames <- goFrames after []+  Right (schema, frames)+  where+    consumeSchema bs = do+      (mlen, meta, rest) <- readFrameHeader bs+      when' (mlen <= 0) $+        Left "Arrow.FlatBufferIPC: unexpected EOS while reading schema"+      sch <- decodeSchemaMessage meta+      Right (sch, rest)++    goFrames bs acc+      | BS.length bs < 4 = Right (reverse acc)+      | otherwise = do+          (mlen, meta, rest1) <- readFrameHeader bs+          if mlen == 0+            then Right (reverse acc)+            else do+              ht <- peekHeaderType meta+              case ht of+                3 -> do+                  (rb, bodyLen) <- decodeRecordBatchMessage meta+                  (body, rest2) <- takeFrameBody "record batch body" bodyLen rest1+                  goFrames rest2 (SFBatch rb body : acc)+                2 -> do+                  (did, isDelta, rb, bodyLen) <-+                    decodeDictionaryBatchMessage meta+                  (body, rest2) <- takeFrameBody "dictionary batch body" bodyLen rest1+                  let !db = DictBatch { dbId = did, dbIsDelta = isDelta+                                      , dbData = rb, dbBody = body }+                  goFrames rest2 (SFDict db : acc)+                _ ->+                  Left ("Arrow.FlatBufferIPC: unsupported message header_type "+                        ++ show ht)++-- | Like 'readArrowStreamFB' but also returns any 'DictBatch'+-- frames encountered (in stream order). Most pyarrow / arrow-cpp+-- streams emit dictionary batches before the first record batch+-- whose schema references their @id@.+readArrowStreamFBWithDicts+  :: ByteString+  -> Either String (Schema, [DictBatch], [(RecordBatchDef, ByteString)])+readArrowStreamFBWithDicts bs0 = do+  (schema, frames) <- readArrowStreamFBInterleaved bs0+  let split (SFDict db) (ds, bs) = (db : ds, bs)+      split (SFBatch rb body) (ds, bs) = (ds, (rb, body) : bs)+      (dicts, batches) = foldr split ([], []) frames+  Right (schema, dicts, batches)++-- | Look up the @header_type@ ubyte from a Message flatbuffer.+peekHeaderType :: ByteString -> Either String Int+peekHeaderType meta = do+  msgPos <- fromIntegral <$> peekU32 meta 0+  s <- resolveTable meta msgPos+  case s 1 of+    Nothing -> Right 0+    Just b  -> fromIntegral <$> peekU8 meta b++-- | Parse one encapsulated frame (with or without the continuation+-- marker) into its 'Message' and the @bodyLength@ it declares. The+-- end-of-stream marker and unsupported header types are 'Left'.+decodeMessageFrame :: ByteString -> Either String (Message, Int64)+decodeMessageFrame bs = do+  (mlen, meta, _) <- readFrameHeader bs+  when' (mlen == 0) $+    Left "Arrow.FlatBufferIPC: end-of-stream marker, not a message"+  when' (BS.length meta < mlen) $+    Left "Arrow.FlatBufferIPC: truncated message metadata"+  ht <- peekHeaderType meta+  case ht of+    1 -> (\s -> (SchemaMessage s, 0)) <$> decodeSchemaMessage meta+    2 -> (\(did, isDelta, rb, n) -> (DictionaryBatch did isDelta rb, n)) <$> decodeDictionaryBatchMessage meta+    3 -> (\(rb, n) -> (RecordBatch rb, n)) <$> decodeRecordBatchMessage meta+    _ -> Left ("Arrow.FlatBufferIPC: unsupported message header_type " ++ show ht)++-- | Parse an Arrow IPC /file/ (per @format/File.fbs@) using its footer+-- index. File padding may align messages to 64 bytes rather than 8, and+-- a file must follow the footer's block lists rather than stream positions.+readArrowFileFB+  :: ByteString+  -> Either String (Schema, [(RecordBatchDef, ByteString)])+readArrowFileFB bs = do+  (sch, _, batches) <- readArrowFileFBWithDicts bs+  Right (sch, batches)++-- | Like 'readArrowFileFB' but also returns any 'DictBatch'+-- frames the file contains.+readArrowFileFBWithDicts+  :: ByteString+  -> Either String (Schema, [DictBatch], [(RecordBatchDef, ByteString)])+readArrowFileFBWithDicts bs = do+  index <- readFileIndex bs+  dicts <- V.mapM (block dictionary) (fiDictBlocks index)+  batches <- V.mapM (block batch) (fiRecordBatches index)+  Right (fiSchema index, V.toList dicts, V.toList batches)+  where+    dictionary message body = case message of+      DictionaryBatch ident delta definition -> Right (DictBatch ident delta definition body)+      _ -> Left "Arrow.FlatBufferIPC: file dictionary block has the wrong message type"+    batch message body = case message of+      RecordBatch definition -> Right (definition, body)+      _ -> Left "Arrow.FlatBufferIPC: file record block has the wrong message type"+    block consume descriptor = do+      let offset = fromIntegral (abOffset descriptor)+          metadataLength = fromIntegral (abMetaLen descriptor)+          metadata = BS.take metadataLength (BS.drop offset bs)+      (message, declaredBody) <- decodeMessageFrame metadata+      when' (declaredBody < 0 || declaredBody > abBodyLen descriptor) $+        Left "Arrow.FlatBufferIPC: file block body length contradicts its message"+      let body = BS.take (fromIntegral declaredBody) (BS.drop (offset + metadataLength) bs)+      consume message body++-- ============================================================+-- File footer index+-- ============================================================++{- | The @Footer@ of an Arrow IPC file (@format/File.fbs@): the+schema plus the byte ranges of every dictionary and record batch.+Each 'ArrowBlock' names one self-contained section of the file; a+reader can seek to it, decode just that message and its body, and+move on without parsing anything in between.+-}+data FileIndex = FileIndex+  { fiSchema :: !Schema+  , fiDictBlocks :: !(V.Vector ArrowBlock)+  -- ^ Dictionary batches, in file order.+  , fiRecordBatches :: !(V.Vector ArrowBlock)+  -- ^ Record batches, in file order.+  }+  deriving (Show, Eq)+++{- | Read the 'FileIndex' out of whole file bytes: check the+@ARROW1@ magics, peel the trailing @<i32 footer length> ARROW1@+trailer, and parse the footer flatbuffer. Every block's offset and+length is validated against the file size.+-}+readFileIndex :: ByteString -> Either String FileIndex+readFileIndex bs = do+  when' (BS.length bs < 14) $+    Left "Arrow.FlatBufferIPC: input too small to be an Arrow file"+  when' (BS.take 6 bs /= "ARROW1") $+    Left "Arrow.FlatBufferIPC: missing leading ARROW1 magic"+  when' (BS.takeEnd 6 bs /= "ARROW1") $+    Left "Arrow.FlatBufferIPC: missing trailing ARROW1 magic"+  let !n = BS.length bs+      !trailer = BS.drop (n - 10) bs+  footerLen <- peekI32 trailer 0+  when' (footerLen <= 0 || fromIntegral footerLen > n - 16) $+    Left "Arrow.FlatBufferIPC: invalid footer length in file trailer"+  let !footer = BS.take (fromIntegral footerLen) (BS.drop (n - 10 - fromIntegral footerLen) bs)+  parseFileIndex (fromIntegral n) footer+++{- | Parse a @Footer@ flatbuffer from its bytes alone, for readers+that fetched only the footer region (seek-and-read streaming).+@fileSize@ is the total file size the block ranges are validated+against.+-}+parseFileIndex :: Int64 -> ByteString -> Either String FileIndex+parseFileIndex fileSize footer = do+  when' (BS.length footer < 8) $+    Left "Arrow.FlatBufferIPC: footer too small"+  when' (fileSize < 18 || fromIntegral (BS.length footer) > fileSize - 18) $+    Left "Arrow.FlatBufferIPC: footer runs before the file header"+  let dataEnd = fileSize - 10 - fromIntegral (BS.length footer)+  rootPos <- fromIntegral <$> peekU32 footer 0+  s <- resolveTable footer rootPos+  schField <- case s 1 of+    Nothing -> Left "Arrow.FlatBufferIPC: footer has no schema"+    Just b -> Right b+  sch <- followUOffset footer schField >>= readSchemaTable footer+  dictBlocks <- readFooterBlocks dataEnd footer (s 2)+  rbBlocks <- readFooterBlocks dataEnd footer (s 3)+  Right+    FileIndex+      { fiSchema = sch+      , fiDictBlocks = dictBlocks+      , fiRecordBatches = rbBlocks+      }+++-- | The @[Block]@ struct vector at a footer slot, range-validated.+readFooterBlocks :: Int64 -> ByteString -> Maybe Pos -> Either String (V.Vector ArrowBlock)+readFooterBlocks _ _ Nothing = Right V.empty+readFooterBlocks fileSize footer (Just slotPos) = do+  vecPos <- followUOffset footer slotPos+  (n, elems) <- readVectorOfStructs footer vecPos 24+  V.generateM n $ \i -> do+    let !p = V.unsafeIndex elems i+    off <- peekI64 footer p+    metaLen <- peekI32 footer (p + 8)+    bodyLen <- peekI64 footer (p + 16)+    if off < 8 || metaLen <= 0 || bodyLen < 0+      then Left "Arrow.FlatBufferIPC: footer block has an invalid offset or length"+      else+        if off > fileSize - fromIntegral metaLen+          then Left "Arrow.FlatBufferIPC: footer block header runs past the end of the file"+          else+            if bodyLen > fileSize - off - fromIntegral metaLen+              then Left "Arrow.FlatBufferIPC: footer block body runs past the end of the file"+              else+                Right+                  ArrowBlock+                    { abOffset = off+                    , abMetaLen = metaLen+                    , abBodyLen = bodyLen+                    }++-- | Strip one encapsulated-message frame:+--+--   * 4 bytes continuation (0xFFFFFFFF) — optional in legacy mode+--   * 4 bytes metadata_length (i32 LE)+--   * @metadata_length@ metadata bytes (already padded)+--+-- Returns @(mlen, metadata bytes, rest of stream after metadata)@.+-- @mlen == 0@ signals the EOS marker.+readFrameHeader+  :: ByteString+  -> Either String (Int, ByteString, ByteString)+readFrameHeader bs = do+  when' (BS.length bs < 4) $+    Left "Arrow.FlatBufferIPC: truncated frame header"+  first4 <- peekU32 bs 0+  if first4 == 0xFFFFFFFF+    then do+      when' (BS.length bs < 8) $+        Left "Arrow.FlatBufferIPC: truncated frame after continuation"+      mlen <- peekI32 bs 4+      metadataAt 8 (fromIntegral mlen)+    else+      -- Legacy: first 4 bytes are the metadata length itself.+      if first4 == 0+        then Right (0, BS.empty, BS.drop 4 bs)+        else metadataAt 4 (fromIntegral first4)+  where+    metadataAt :: Int -> Int -> Either String (Int, ByteString, ByteString)+    metadataAt start mlenI+      | mlenI < 0 = Left "Arrow.FlatBufferIPC: negative metadata length"+      | mlenI > BS.length bs - start =+          Left+            ( "Arrow.FlatBufferIPC: frame claims "+                ++ show mlenI+                ++ " metadata bytes, "+                ++ show (BS.length bs - start)+                ++ " remain"+            )+      | otherwise =+          Right ( mlenI+                , BS.take mlenI (BS.drop start bs)+                , BS.drop (start + mlenI) bs+                )
+ src/Arrow/FlatBufferIPC/Write.hs view
@@ -0,0 +1,1126 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE StrictData #-}+{-# LANGUAGE TypeApplications #-}+-- | Arrow IPC encoder: builds the @Schema@, @Field@, @Type@,+-- @RecordBatch@, @DictionaryBatch@, @Tensor@ and @SparseTensor@+-- flatbuffer tables, the encapsulated-message framing, the stream+-- and file layouts, and body compression. See "Arrow.FlatBufferIPC"+-- for the format notes.+--+-- Record batch and dictionary batch headers are written front to+-- back straight into their final bytes (no FlatBuffers builder): the+-- field nodes and buffers are 'Storable' vectors whose memory layout+-- is the wire layout, so each is one memcpy. Streams and files are+-- laid out from 'Frame's: 'renderStream' / 'renderFile' make one+-- allocation and copy every body byte once, 'renderStreamLazy' /+-- 'renderFileLazy' alias the body pieces as lazy chunks.+module Arrow.FlatBufferIPC.Write+  ( -- * Top-level builders+    buildSchemaMessage+  , buildRecordBatchMessage+    -- * Encapsulated-message framing+  , encapsulateMessage+    -- * Frames and layout+  , Frame (..)+  , EncodedBatch (..)+  , encodeBatch+  , encodeBatchWithEndianness+  , schemaFrame+  , batchFrame+  , dictionaryFrame+  , renderStream+  , renderStreamLazy+  , renderFile+  , renderFileLazy+    -- * Stream / file writers+  , writeArrowStreamFB+  , writeArrowFileFB+  , writeArrowFileFBWithDicts+    -- * Column-based convenience writer+  , buildRecordBatchBytes+  , buildRecordBatchBytesWith+  , writeArrowStreamFBFromColumns+    -- * Body compression helpers+  , compressBody+  , compressBufferEither+  , bodyCompressionAvailable+    -- * Dictionary support+  , buildDictionaryBatchMessage+  , writeArrowStreamFBWithDicts+    -- * Tensor / SparseTensor+  , buildTensorMessage+  , encodeTensorFrame+  , buildSparseTensorMessageCOO+  , encodeSparseTensorFrame+    -- * Single-message frames+  , encodeMessageFrame+  ) where++#include "MachDeps.h"+#if defined(WORDS_BIGENDIAN)+#error "Arrow.FlatBufferIPC.Write copies Storable FieldNode/Buffer vectors as little-endian wire structs"+#endif++import Control.Monad (foldM, when)+import Data.Bits ((.&.), complement)+import Data.ByteString (ByteString)+import qualified Data.ByteString as BS+import qualified Data.ByteString.Internal as BSI+import qualified Data.ByteString.Lazy as BL+import qualified Data.ByteString.Unsafe as BSU+import Data.Int (Int16, Int32, Int64)+import Data.Maybe (isJust)+import qualified Data.Text as T+import qualified Data.Vector as V+import qualified Data.Vector.Storable as VS+import Data.Word (Word16, Word32, Word8)+import Foreign.Marshal.Utils (copyBytes, fillBytes)+import Foreign.Ptr (Ptr, castPtr, plusPtr)+import Foreign.Storable (pokeByteOff)+import System.IO.Unsafe (unsafePerformIO)++import Arrow.Column (ColumnArray)+import Arrow.FlatBufferIPC.Common+import Arrow.Types+import Arrow.Write.Columns (BatchPlan (..), planBatchWithEndianness)+import FlatBuffers.Builder+  ( Builder+  , alignUp+  , finish+  , newBuilder+  , prependBS+  , prependI16+  , prependI32+  , prependI64+  , prependU8+  , scalar+  , struct+  , voff+  , writeString+  , writeTable+  , writeVectorInt32+  , writeVectorInt64+  , writeVectorOfOffsets+  , writeVectorOfStructs+  )++#ifdef HAVE_ZSTD+import qualified Codec.Compression.Zstd as Zstd+#endif++#ifdef HAVE_LZ4+import qualified Codec.Lz4 as Lz4+#endif+-- ============================================================+-- Arrow-specific: Type tables+-- ============================================================++-- | Returns (union_tag, UOffset of the type table).+writeType :: Builder -> ArrowType -> IO (Word8, Int)+writeType b ty = case ty of+  ANull             -> emptyT 1+  AInt bits signed  -> do+    -- Arrow's Int.fbs declares is_signed with no default, but+    -- arrow-cpp's generated reader defaults absent slots to+    -- @true@. The writer used to omit the slot when+    -- @signed = False@, which silently coerced unsigned columns+    -- back to signed on round-trip. Emit the slot explicitly+    -- whenever @signed = False@ so both paths survive.+    u <- writeTable b+           [ Just (scalar 4 (\bb -> prependI32 bb (fromIntegral bits)))+           , Just (scalar 1 (\bb -> prependU8 bb (if signed then 1 else 0)))+           ]+    pure (2, u)+  AFloatingPoint p  -> do+    u <- writeTable b [ Just (scalar 2 (\bb -> prependI16 bb (fromIntegral (precisionTag p)))) ]+    pure (3, u)+  ABinary           -> emptyT 4+  AUtf8             -> emptyT 5+  ABool             -> emptyT 6+  ADecimal32 p s   -> decimal 32 p s+  ADecimal64 p s   -> decimal 64 p s+  ADecimal p s      -> do+    u <- writeTable b+           [ Just (scalar 4 (\bb -> prependI32 bb (fromIntegral p)))+           , Just (scalar 4 (\bb -> prependI32 bb (fromIntegral s)))+           , Nothing   -- bitWidth default 128+           ]+    pure (7, u)+  ADecimal256 p s   -> do+    u <- writeTable b+           [ Just (scalar 4 (\bb -> prependI32 bb (fromIntegral p)))+           , Just (scalar 4 (\bb -> prependI32 bb (fromIntegral s)))+           , Just (scalar 4 (\bb -> prependI32 bb 256))+           ]+    pure (7, u)+  ADate u' -> do+    u <- writeTable b [ Just (scalar 2 (\bb -> prependI16 bb (fromIntegral (dateUnitTag u')))) ]+    pure (8, u)+  ATime u' bits -> do+    u <- writeTable b+           [ Just (scalar 2 (\bb -> prependI16 bb (fromIntegral (timeUnitTag u'))))+           , Just (scalar 4 (\bb -> prependI32 bb (fromIntegral bits)))+           ]+    pure (9, u)+  ATimestamp u' tz -> do+    tzOff <- case tz of+      Nothing -> pure Nothing+      Just t  -> Just <$> writeString b t+    u <- writeTable b+           [ Just (scalar 2 (\bb -> prependI16 bb (fromIntegral (timeUnitTag u'))))+           , case tzOff of+               Nothing  -> Nothing+               Just uo  -> Just (voff uo)+           ]+    pure (10, u)+  AInterval u' -> do+    u <- writeTable b [ Just (scalar 2 (\bb -> prependI16 bb (fromIntegral (intervalUnitTag u')))) ]+    pure (11, u)+  AList  -> emptyT 12+  AStruct -> emptyT 13+  AUnion mode typeIds -> do+    idsOff <- if V.null typeIds+                then pure Nothing+                else Just <$> writeVectorInt32 b (V.toList typeIds)+    u <- writeTable b+           [ Just (scalar 2 (\bb -> prependI16 bb (fromIntegral (unionModeTag mode))))+           , case idsOff of { Nothing -> Nothing; Just uo -> Just (voff uo) }+           ]+    pure (14, u)+  AFixedSizeBinary n -> do+    u <- writeTable b [ Just (scalar 4 (\bb -> prependI32 bb (fromIntegral n))) ]+    pure (15, u)+  AFixedSizeList n -> do+    u <- writeTable b [ Just (scalar 4 (\bb -> prependI32 bb (fromIntegral n))) ]+    pure (16, u)+  AMap sorted -> do+    u <- writeTable b+           [ if sorted then Just (scalar 1 (\bb -> prependU8 bb 1)) else Nothing ]+    pure (17, u)+  ADuration u' -> do+    u <- writeTable b [ Just (scalar 2 (\bb -> prependI16 bb (fromIntegral (timeUnitTag u')))) ]+    pure (18, u)+  ALargeBinary    -> emptyT 19+  ALargeUtf8      -> emptyT 20+  ALargeList      -> emptyT 21+  ARunEndEncoded  -> emptyT 22+  ABinaryView     -> emptyT 23+  AUtf8View       -> emptyT 24+  AListView       -> emptyT 25+  ALargeListView  -> emptyT 26+  where+    decimal width precision scale = do+      u <- writeTable b+             [ Just (scalar 4 (\bb -> prependI32 bb (fromIntegral precision)))+             , Just (scalar 4 (\bb -> prependI32 bb (fromIntegral scale)))+             , Just (scalar 4 (\bb -> prependI32 bb width))+             ]+      pure (7, u)+    emptyT !tag = do+      u <- writeTable b []+      pure (tag, u)++precisionTag :: Precision -> Int+precisionTag Half            = 0+precisionTag Single          = 1+precisionTag DoublePrecision = 2++dateUnitTag :: DateUnit -> Int+dateUnitTag DateDay         = 0+dateUnitTag DateMillisecond = 1++timeUnitTag :: TimeUnit -> Int+timeUnitTag Second      = 0+timeUnitTag Millisecond = 1+timeUnitTag Microsecond = 2+timeUnitTag Nanosecond  = 3++intervalUnitTag :: IntervalUnit -> Int+intervalUnitTag YearMonth    = 0+intervalUnitTag DayTime      = 1+intervalUnitTag MonthDayNano = 2++unionModeTag :: UnionMode -> Int+unionModeTag Sparse = 0+unionModeTag Dense  = 1++-- ============================================================+-- Field + Schema tables+-- ============================================================++-- | @+-- table Field {+--   name            : string;           // 0+--   nullable        : bool;             // 1+--   type_type       : ubyte;            // 2+--   type            : Type;             // 3+--   dictionary      : DictionaryEncoding; // 4+--   children        : [Field];          // 5+--   custom_metadata : [KeyValue];       // 6+-- }+-- @+writeField :: Builder -> Field -> IO Int+writeField b fld = do+  childrenVec <- if V.null (fieldChildren fld)+                   then pure Nothing+                   else do+                     childUOffs <- mapM (writeField b) (V.toList (fieldChildren fld))+                     Just <$> writeVectorOfOffsets b childUOffs+  (tyTag, tyUOff) <- writeType b (fieldType fld)+  dictOff <- case fieldDictionary fld of+    Nothing -> pure Nothing+    Just de -> Just <$> writeDictionaryEncoding b de+  nameOff <- if T.null (fieldName fld)+               then pure Nothing+               else Just <$> writeString b (fieldName fld)+  customMd <- writeKeyValueVector b (fieldMetadata fld)+  writeTable b+    [ case nameOff of { Nothing -> Nothing; Just uo -> Just (voff uo) }+    , if fieldNullable fld then Just (scalar 1 (\bb -> prependU8 bb 1)) else Nothing+    , Just (scalar 1 (\bb -> prependU8 bb tyTag))+    , Just (voff tyUOff)+    , case dictOff of { Nothing -> Nothing; Just uo -> Just (voff uo) }+    , case childrenVec of { Nothing -> Nothing; Just uo -> Just (voff uo) }+    , case customMd  of { Nothing -> Nothing; Just uo -> Just (voff uo) }+    ]++-- | Build a @[KeyValue]@ vector for @custom_metadata@ slots+-- (used by both 'Field' and 'Schema'). Returns 'Nothing' for an+-- empty metadata vector so the caller can omit the slot+-- entirely (FlatBuffers' 'Nothing' = default).+writeKeyValueVector :: Builder -> V.Vector (T.Text, T.Text) -> IO (Maybe Int)+writeKeyValueVector b kvs+  | V.null kvs = pure Nothing+  | otherwise = do+      kvUOffs <- mapM (writeKeyValueTable b) (V.toList kvs)+      Just <$> writeVectorOfOffsets b kvUOffs++writeKeyValueTable :: Builder -> (T.Text, T.Text) -> IO Int+writeKeyValueTable b (k, v) = do+  kOff <- writeString b k+  vOff <- writeString b v+  writeTable b [ Just (voff kOff), Just (voff vOff) ]++-- | Build a 'DictionaryEncoding' table:+--+-- @+-- table DictionaryEncoding {+--   id: long;+--   indexType: Int;+--   isOrdered: bool;+--   dictionaryKind: DictionaryKind;+-- }+-- @+writeDictionaryEncoding :: Builder -> DictionaryEncoding -> IO Int+writeDictionaryEncoding b (DictionaryEncoding did indexTy ordered) = do+  -- The indexType is always an Int table; build via 'writeType' to+  -- reuse the layout, but we must always emit the table even if+  -- indexTy is the default Int32-signed.+  (_, intUOff) <- writeType b indexTy+  writeTable b+    [ Just (scalar 8 (\bb -> prependI64 bb did))+    , Just (voff intUOff)+    , if ordered then Just (scalar 1 (\bb -> prependU8 bb 1)) else Nothing+    -- dictionaryKind defaults to DenseArray (0); omit.+    ]++-- | @+-- table Schema {+--   endianness     : Endianness = Little;+--   fields         : [Field];+--   custom_metadata: [KeyValue];+--   features       : [long];+-- }+-- @+writeSchema :: Builder -> Schema -> IO Int+writeSchema b sch = do+  fieldUOffs <- mapM (writeField b) (V.toList (arrowFields sch))+  fieldsVec  <- writeVectorOfOffsets b fieldUOffs+  customMd   <- writeKeyValueVector b (arrowMetadata sch)+  features   <- if V.null (arrowFeatures sch)+                  then pure Nothing+                  else Just <$> writeVectorInt64 b (map featureValue (V.toList (arrowFeatures sch)))+  writeTable b+    [ case arrowEndianness sch of+        Little -> Nothing+        Big    -> Just (scalar 2 (\bb -> prependI16 bb 1))+    , Just (voff fieldsVec)+    , case customMd of { Nothing -> Nothing; Just uo -> Just (voff uo) }+    , fmap voff features+    ]++-- | Keep the wire values explicit: future features use powers of two,+-- so their values need not match Haskell Enum indices.+featureValue :: Feature -> Int64+featureValue feature = case feature of+  FeatureUnused -> 0+  FeatureDictionaryReplacement -> 1+  FeatureCompressedBody -> 2++-- ============================================================+-- RecordBatch / DictionaryBatch messages (direct, front to back)+-- ============================================================++-- | The @DictionaryBatch@ wrapper fields: id and isDelta.+data DictHeader = DictHeader !Int64 !Bool+++align8 :: Int -> Int+align8 n = (n + 7) .&. complement 7+{-# INLINE align8 #-}+++{- | A record batch @Message@ (or, with a 'DictHeader', a dictionary+batch @Message@ wrapping it), written front to back into one+allocation. With @encapsulated@ the result is the whole frame header:+continuation marker, metadata length, the metadata, and zero padding+to an 8-byte multiple; otherwise just the metadata.++Layout (offsets relative to the metadata start, which the framing+keeps 8-aligned): root offset; @Message@ vtable at 4 and table at 16+(version, header type, header, bodyLength at 32); for a dictionary+batch its vtable at 40 and table at 56 (id, data, isDelta); then the+@RecordBatch@ vtable and table (length, nodes, buffers, compression,+variadicBufferCounts), the nodes and buffers vectors (16-byte structs,+one memcpy each from the 'Storable' vectors, 8-aligned data), the+@BodyCompression@ vtable and table, and the variadic counts vector.+Every offset points forward and every vtable precedes its table.+-}+batchMessage :: Bool -> Maybe DictHeader -> RecordBatchDef -> Int64 -> ByteString+batchMessage encapsulated mDict rb bodyLen = BSI.unsafeCreate total $ \p0 -> do+  fillBytes p0 0 total+  when encapsulated $ do+    pokeByteOff p0 0 (0xFFFFFFFF :: Word32)+    pokeByteOff p0 4 (fromIntegral metaPadded :: Int32)+  let !p = p0 `plusPtr` pre+      w8 :: Int -> Int -> IO ()+      w8 o x = pokeByteOff p o (fromIntegral x :: Word8)+      w16 :: Int -> Int -> IO ()+      w16 o x = pokeByteOff p o (fromIntegral x :: Word16)+      w32 :: Int -> Int -> IO ()+      w32 o x = pokeByteOff p o (fromIntegral x :: Word32)+      w64 :: Int -> Int64 -> IO ()+      w64 o x = pokeByteOff p o x+      uoff at target = w32 at (target - at)+  -- Message+  w32 0 16+  w16 4 12 >> w16 6 24 >> w16 8 4 >> w16 10 6 >> w16 12 8 >> w16 14 16+  w32 16 12+  pokeByteOff p 20 metadataVersionV5+  case mDict of+    Nothing -> do+      w8 22 3+      uoff 24 rbTable+    Just (DictHeader did isDelta) -> do+      w8 22 2+      uoff 24 56+      w16 40 10 >> w16 42 24 >> w16 44 8 >> w16 46 16 >> w16 48 (if isDelta then 20 else 0)+      w32 56 16+      w64 64 did+      uoff 72 rbTable+      when isDelta (w8 76 1)+  w64 32 bodyLen+  -- RecordBatch+  w16 rbVt rbVtSize+  w16 (rbVt + 2) rbTableSize+  w16 (rbVt + 4) 8 >> w16 (rbVt + 6) 16 >> w16 (rbVt + 8) 20+  when (rbFields >= 4) (w16 (rbVt + 10) (if hasComp then 24 else 0))+  when (rbFields >= 5) (w16 (rbVt + 12) 28)+  w32 rbTable (rbTable - rbVt)+  w64 (rbTable + 8) (rbLength rb)+  uoff (rbTable + 16) nodesCount+  uoff (rbTable + 20) bufsCount+  w32 nodesCount nN+  VS.unsafeWith (rbNodes rb) $ \src -> copyBytes (p `plusPtr` (nodesCount + 4)) (castPtr src) (16 * nN)+  w32 bufsCount nB+  VS.unsafeWith (rbBuffers rb) $ \src -> copyBytes (p `plusPtr` (bufsCount + 4)) (castPtr src) (16 * nB)+  case rbBodyCompression rb of+    Nothing -> pure ()+    Just codec -> do+      uoff (rbTable + 24) compTable+      w16 afterBufs 6 >> w16 (afterBufs + 2) 8 >> w16 (afterBufs + 4) 4+      w32 compTable (compTable - afterBufs)+      w8 (compTable + 4) (codecTag codec)+  when hasVar $ do+    uoff (rbTable + 28) varCount+    w32 varCount nV+    V.imapM_ (\i x -> w64 (varCount + 4 + 8 * i) x) (rbVariadicBufferCounts rb)+  where+    !pre = if encapsulated then 8 else 0+    !nN = VS.length (rbNodes rb)+    !nB = VS.length (rbBuffers rb)+    !nV = V.length (rbVariadicBufferCounts rb)+    !hasComp = isJust (rbBodyCompression rb)+    !hasVar = nV > 0+    !rbFields = if hasVar then 5 else if hasComp then 4 else 3 :: Int+    !rbVt = case mDict of+      Nothing -> 40+      Just _ -> 80+    !rbVtSize = 4 + 2 * rbFields+    !rbTable = align8 (rbVt + rbVtSize)+    !rbTableSize = if rbFields > 3 then 32 else 24+    !nodesCount = align8 (rbTable + rbTableSize + 4) - 4+    !bufsCount = align8 (nodesCount + 4 + 16 * nN + 4) - 4+    !afterBufs = bufsCount + 4 + 16 * nB+    !compTable = afterBufs + 8+    !afterComp = if hasComp then compTable + 8 else afterBufs+    !varCount = align8 (afterComp + 4) - 4+    !metaSize = if hasVar then varCount + 4 + 8 * nV else afterComp+    !metaPadded = align8 metaSize+    !total = if encapsulated then 8 + metaPadded else metaSize+++codecTag :: BodyCompressionCodec -> Int+codecTag = \case+  LZ4Frame -> 0+  BodyZstd -> 1++-- ============================================================+-- Message envelope+-- ============================================================++-- | @+-- table Message {+--   version        : MetadataVersion;  // short, V5 = 4+--   header_type    : MessageHeader;    // ubyte (1=Schema, 3=RecordBatch)+--   header         : MessageHeader;    // union payload+--   bodyLength     : long;+--   custom_metadata: [KeyValue];+-- }+-- @+--+-- File-identifier is /not/ emitted for Arrow Messages; the+-- encapsulating stream framing distinguishes message boundaries.+-- This low-level assembler assumes valid schema parameters. Checked+-- column writers in Arrow.Stream reject invalid decimal precision/scale.+buildSchemaMessage :: Schema -> ByteString+buildSchemaMessage sch = unsafePerformIO $ do+  b <- newBuilder+  schUOff <- writeSchema b sch+  msgUOff <- writeTable b+    [ Just (scalar 2 (\bb -> prependI16 bb metadataVersionV5))+    , Just (scalar 1 (\bb -> prependU8 bb 1))   -- header type: Schema+    , Just (voff schUOff)+    , Just (scalar 8 (\bb -> prependI64 bb 0))  -- bodyLength+    , Nothing+    ]+  finish b msgUOff+{-# NOINLINE buildSchemaMessage #-}++-- | The @Message@ flatbuffer of a record batch whose body is @bodyLen@ bytes.+buildRecordBatchMessage :: RecordBatchDef -> Int64 -> ByteString+buildRecordBatchMessage = batchMessage False Nothing++-- | Build a @Message@ flatbuffer wrapping a @DictionaryBatch@:+--+-- @+-- table DictionaryBatch {+--   id      : long;          // 0+--   data    : RecordBatch;   // 1+--   isDelta : bool = false;  // 2+-- }+-- @+--+-- The @bodyLength@ is the batch body's length padded to 8 bytes.+buildDictionaryBatchMessage :: DictBatch -> ByteString+buildDictionaryBatchMessage (DictBatch did isDelta rb body) =+  batchMessage False (Just (DictHeader did isDelta)) rb (fromIntegral (align8 (BS.length body)))++metadataVersionV5 :: Int16+metadataVersionV5 = 4++-- | Build a @Message@ flatbuffer wrapping a @Tensor@:+--+-- @+-- table Tensor {+--   type_type: Type;       // slot 0 (union tag)+--   type:      Type;       // slot 1+--   shape:    [TensorDim]; // slot 2+--   strides:  [long];      // slot 3+--   data:      Buffer;     // slot 4 (struct)+-- }+-- @+buildTensorMessage :: Tensor -> ByteString+buildTensorMessage t = unsafePerformIO $ do+  b <- newBuilder+  (tyTag, tyUOff) <- writeType b (tensorType t)+  -- Each TensorDim is a table (not a struct) because @name@ is a+  -- variable-length string.+  dimUOffs <- mapM (writeTensorDim b) (V.toList (tensorShape t))+  shapeVec <- writeVectorOfOffsets b dimUOffs+  stridesVec <- if V.null (tensorStrides t)+                  then pure Nothing+                  else Just <$> writeVectorInt64 b (V.toList (tensorStrides t))+  -- Data Buffer struct: i64 offset (always 0, we emit body+  -- right after the metadata) + i64 length.+  let !bodyLen = fromIntegral (BS.length (tensorBody t)) :: Int64+  tensorUOff <- writeTable b+    [ Just (scalar 1 (\bb -> prependU8 bb tyTag))+    , Just (voff tyUOff)+    , Just (voff shapeVec)+    , fmap voff stridesVec+    , Just (struct 16 8 (\bb -> do+        prependI64 bb bodyLen+        prependI64 bb 0))+    ]+  -- Now wrap in a Message table: header_type = 4 (Tensor).+  msgUOff <- writeTable b+    [ Just (scalar 2 (\bb -> prependI16 bb metadataVersionV5))+    , Just (scalar 1 (\bb -> prependU8 bb 4))+    , Just (voff tensorUOff)+    , Just (scalar 8 (\bb -> prependI64 bb bodyLen))+    , Nothing+    ]+  finish b msgUOff+{-# NOINLINE buildTensorMessage #-}++-- | Encode one @TensorDim@ as a flatbuffer table.+writeTensorDim :: Builder -> TensorDim -> IO Int+writeTensorDim b (TensorDim size name) = do+  nameOff <- if T.null name+               then pure Nothing+               else Just <$> writeString b name+  writeTable b+    [ Just (scalar 8 (\bb -> prependI64 bb size))+    , fmap voff nameOff+    ]++-- | Build a SparseTensor @Message@ flatbuffer (COO index format).+-- The body of the encapsulated frame carries indices followed by+-- values, concatenated; callers must handle the per-buffer+-- offsets on the encoding side.+buildSparseTensorMessageCOO :: SparseTensor -> ByteString+buildSparseTensorMessageCOO st = unsafePerformIO $ do+  b <- newBuilder+  (tyTag, tyUOff) <- writeType b (sparseTensorType st)+  dimUOffs <- mapM (writeTensorDim b) (V.toList (sparseTensorShape st))+  shapeVec <- writeVectorOfOffsets b dimUOffs+  -- Indices buffer struct (offset = 0, length)+  let !indicesLen = fromIntegral (BS.length (sparseIndicesBody st)) :: Int64+      !valuesLen  = fromIntegral (BS.length (sparseTensorBody  st)) :: Int64+      !indicesOffset = 0 :: Int64+      !valuesOffset  = alignUp (BS.length (sparseIndicesBody st)) 8+  (_, idxTypeUOff) <- writeType b (sparseIndicesType st)+  cooUOff <- writeTable b+    [ Just (voff idxTypeUOff)+    , Nothing                   -- indicesStrides (empty)+    , Just (struct 16 8 (\bb -> do+        prependI64 bb indicesLen+        prependI64 bb indicesOffset))+    , if sparseIndicesCanonical st+        then Just (scalar 1 (\bb -> prependU8 bb 1)) else Nothing+    ]+  -- SparseTensor table+  stUOff <- writeTable b+    [ Just (scalar 1 (\bb -> prependU8 bb tyTag))+    , Just (voff tyUOff)+    , Just (voff shapeVec)+    , Just (scalar 8 (\bb -> prependI64 bb (sparseNonZeroLength st)))+    , Just (scalar 1 (\bb -> prependU8 bb 1))   -- SparseTensorIndex = COO+    , Just (voff cooUOff)+    , Just (struct 16 8 (\bb -> do+        prependI64 bb valuesLen+        prependI64 bb (fromIntegral valuesOffset)))+    ]+  msgUOff <- writeTable b+    [ Just (scalar 2 (\bb -> prependI16 bb metadataVersionV5))+    , Just (scalar 1 (\bb -> prependU8 bb 5))   -- header type: SparseTensor+    , Just (voff stUOff)+    , Just (scalar 8 (\bb -> prependI64 bb (indicesLen + fromIntegral valuesOffset - indicesLen + valuesLen)))+    , Nothing+    ]+  finish b msgUOff+{-# NOINLINE buildSparseTensorMessageCOO #-}++-- | Encode a sparse tensor as a standalone Arrow IPC frame.+-- The encapsulated body is @indicesBody <pad8> tensorBody@, so+-- the body offsets recorded in the message match the slice+-- layout.+encodeSparseTensorFrame :: SparseTensor -> ByteString+encodeSparseTensorFrame st =+  let !iLen   = BS.length (sparseIndicesBody st)+      !iPad   = alignUp iLen 8 - iLen+      !body   = BS.concat+                  [ sparseIndicesBody st+                  , BS.replicate iPad 0+                  , sparseTensorBody  st+                  ]+  in  encapsulateMessage (buildSparseTensorMessageCOO st) body++-- | Encode a 'Tensor' as a standalone Arrow IPC frame:+-- continuation + metadata length + padded flatbuffer + body+-- + body padding. Suitable as a file payload or as one element+-- of an application-level container; Arrow's stream framing+-- itself accepts Tensor messages interleaved with any other+-- Message.+encodeTensorFrame :: Tensor -> ByteString+encodeTensorFrame t = encapsulateMessage (buildTensorMessage t) (tensorBody t)++-- ============================================================+-- Encapsulated framing+-- ============================================================++-- | Wrap a raw flatbuffer @Message@ in the encapsulated IPC frame+-- (one allocation):+--+-- @+-- <continuation 0xFFFFFFFF : u32-LE>+-- <metadata_length : i32-LE, padded so body starts aligned to 8>+-- <flatbuffer bytes, padded>+-- <body bytes, padded to 8-byte alignment>+-- @+encapsulateMessage :: ByteString -> ByteString -> ByteString+encapsulateMessage meta body =+  let !metaLen = BS.length meta+      !padded = align8 metaLen+      !bodyLen = BS.length body+      !total = 8 + padded + align8 bodyLen+  in BSI.unsafeCreate total $ \p -> do+       pokeByteOff p 0 (0xFFFFFFFF :: Word32)+       pokeByteOff p 4 (fromIntegral padded :: Int32)+       copyInto (p `plusPtr` 8) meta+       fillBytes (p `plusPtr` (8 + metaLen)) 0 (padded - metaLen)+       copyInto (p `plusPtr` (8 + padded)) body+       fillBytes (p `plusPtr` (8 + padded + bodyLen)) 0 (total - 8 - padded - bodyLen)+++copyInto :: Ptr Word8 -> ByteString -> IO ()+copyInto dst bs = BSU.unsafeUseAsCStringLen bs $ \(src, n) -> copyBytes dst (castPtr src) n+{-# INLINE copyInto #-}+++-- ============================================================+-- Frames and layout+-- ============================================================++{- | One encapsulated message ready to be laid out: its header+(continuation marker, metadata length, metadata, zero padding; a+multiple of 8 bytes) and its body pieces. On output every piece is+followed by zero padding to a multiple of 8, and 'frameBodyLength' is+the padded total (the @bodyLength@ the header declares).+-}+data Frame = Frame+  { frameHeader :: !ByteString+  , framePieces :: !(V.Vector ByteString)+  , frameBodyLength :: !Int+  }+++frameSize :: Frame -> Int+frameSize fr = BS.length (frameHeader fr) + frameBodyLength fr+++-- | The schema message frame (no body).+schemaFrame :: Schema -> Frame+schemaFrame sch = Frame (encapsulateMessage (buildSchemaMessage sch) BS.empty) V.empty 0+++{- | A record batch laid out for the wire: the header definition (buffer+offsets relative to the body start, 8-byte aligned) and the body pieces+(one per buffer, compressed when the definition says so).+-}+data EncodedBatch = EncodedBatch+  { ebDef :: !RecordBatchDef+  , ebPieces :: !(V.Vector ByteString)+  , ebBodyLength :: !Int+  }+++{- | Plan a batch (see 'Arrow.Write.Columns.planBatch') and compress its+pieces when a codec is requested. Body words use little-endian order;+use 'encodeBatchWithEndianness' for a big-endian schema.+A codec that was not compiled in+(see 'bodyCompressionAvailable') is ignored, so the batch is always a+valid uncompressed one.+-}+encodeBatch :: Maybe BodyCompressionCodec -> V.Vector Field -> V.Vector ColumnArray -> Either String EncodedBatch+encodeBatch = encodeBatchWithEndianness Little+++-- | Encode a column batch in the declared schema byte order.+encodeBatchWithEndianness :: Endianness -> Maybe BodyCompressionCodec -> V.Vector Field -> V.Vector ColumnArray -> Either String EncodedBatch+encodeBatchWithEndianness endian mCodec0 fields cols = do+  plan <- planBatchWithEndianness endian fields cols+  let !mCodec = case mCodec0 of+        Just codec | bodyCompressionAvailable codec -> Just codec+        _ -> Nothing+  pieces <- case mCodec of+    Nothing -> Right (bpPieces plan)+    Just codec -> V.mapM (compressPiece codec) (bpPieces plan)+  let !(bufs, bodyLen) = layoutBuffers pieces+  Right+    EncodedBatch+      { ebDef =+          RecordBatchDef+            { rbLength = bpLength plan+            , rbNodes = bpNodes plan+            , rbBuffers = bufs+            , rbVariadicBufferCounts = bpVariadic plan+            , rbBodyCompression = mCodec+            }+      , ebPieces = pieces+      , ebBodyLength = bodyLen+      }+++-- | Buffer descriptors for pieces laid out back to back, each padded to 8 bytes, and the padded total.+layoutBuffers :: V.Vector ByteString -> (VS.Vector Buffer, Int)+layoutBuffers pieces =+  let !lens = VS.generate (V.length pieces) (BS.length . V.unsafeIndex pieces)+      !offs = VS.prescanl' (\o l -> o + align8 l) 0 lens+      !bufs = VS.zipWith (\o l -> Buffer (fromIntegral o) (fromIntegral l)) offs lens+  in (bufs, VS.foldl' (\t l -> t + align8 l) 0 lens)+++-- | The frame of an encoded record batch.+batchFrame :: EncodedBatch -> Frame+batchFrame eb =+  Frame (batchMessage True Nothing (ebDef eb) (fromIntegral (ebBodyLength eb))) (ebPieces eb) (ebBodyLength eb)+++-- | The frame of an encoded dictionary batch (id, isDelta, the values batch).+dictionaryFrame :: Int64 -> Bool -> EncodedBatch -> Frame+dictionaryFrame did isDelta eb =+  Frame+    (batchMessage True (Just (DictHeader did isDelta)) (ebDef eb) (fromIntegral (ebBodyLength eb)))+    (ebPieces eb)+    (ebBodyLength eb)+++-- | A frame for a definition and a pre-built body (one piece).+rawBatchFrame :: Maybe DictHeader -> RecordBatchDef -> ByteString -> Frame+rawBatchFrame mDict rb body =+  let !bodyLen = align8 (BS.length body)+  in Frame (batchMessage True mDict rb (fromIntegral bodyLen)) (V.singleton body) bodyLen+++-- | Write a frame at @off@; returns the offset after it.+pokeFrame :: Ptr Word8 -> Int -> Frame -> IO Int+pokeFrame p off fr = do+  let !h = frameHeader fr+  copyInto (p `plusPtr` off) h+  V.foldM'+    ( \o pc -> do+        let !l = BS.length pc+            !padded = align8 l+        copyInto (p `plusPtr` o) pc+        fillBytes (p `plusPtr` (o + l)) 0 (padded - l)+        pure (o + padded)+    )+    (off + BS.length h)+    (framePieces fr)+++-- | The end-of-stream marker: continuation, zero metadata length.+eosBytes :: ByteString+eosBytes = BS.pack [0xFF, 0xFF, 0xFF, 0xFF, 0, 0, 0, 0]+{-# NOINLINE eosBytes #-}+++-- | Zero bytes the lazy layouts slice their padding from.+zeroPad :: ByteString+zeroPad = BS.replicate 8 0+{-# NOINLINE zeroPad #-}+++{- | A stream (the frames in order, then the end-of-stream marker) in+one allocation; every body byte is copied once.+-}+renderStream :: [Frame] -> ByteString+renderStream frames =+  let !total = foldl (\t fr -> t + frameSize fr) (BS.length eosBytes) frames+  in BSI.unsafeCreate total $ \p -> do+       end <- foldM (pokeFrame p) 0 frames+       copyInto (p `plusPtr` end) eosBytes+++-- | Lazy chunks of one frame: the header, then each piece (aliased) and its padding.+frameChunks :: Frame -> [ByteString]+frameChunks fr = frameHeader fr : V.foldr (\pc rest -> pc : padChunk (BS.length pc) rest) [] (framePieces fr)+  where+    padChunk l rest =+      let !k = align8 l - l+      in if k == 0 then rest else BS.take k zeroPad : rest+++{- | 'renderStream' as a lazy 'BL.ByteString' whose chunks are the frame+headers, the body pieces themselves (no copy), and shared padding.+-}+renderStreamLazy :: [Frame] -> BL.ByteString+renderStreamLazy frames = BL.fromChunks (concatMap frameChunks frames ++ [eosBytes])+++-- | The leading and trailing file magic (@ARROW1@ plus two padding bytes).+fileMagic :: ByteString+fileMagic = "ARROW1\0\0"+{-# NOINLINE fileMagic #-}+++{- | The file layout around the stream: the footer (indexing the dictionary+and record batch frames) and its trailer (footer length, @ARROW1@).+-}+fileFooter :: Schema -> [Frame] -> [Frame] -> ByteString+fileFooter sch dictFrames batchFrames =+  let !schemaSize = frameSize (schemaFrame sch)+      blocks off0 frs = reverse (snd (foldl (\(!o, acc) fr -> (o + frameSize fr, block o fr : acc)) (off0, []) frs))+      block o fr =+        ArrowBlock+          { abOffset = fromIntegral o+          , abMetaLen = fromIntegral (BS.length (frameHeader fr))+          , abBodyLen = fromIntegral (frameBodyLength fr)+          }+      !dictStart = BS.length fileMagic + schemaSize+      !batchStart = foldl (\t fr -> t + frameSize fr) dictStart dictFrames+      !footer = buildFileFooter sch (blocks dictStart dictFrames) (blocks batchStart batchFrames)+      !trailer = BSI.unsafeCreate 10 $ \p -> do+        pokeByteOff p 0 (fromIntegral (BS.length footer) :: Int32)+        copyInto (p `plusPtr` 4) (BS.take 6 fileMagic)+  in footer <> trailer+++{- | An Arrow IPC /file/ (per @format/File.fbs@) in one allocation:++@+'ARROW1'\\0\\0+<schema message> <dictionary batches> <record batches> <EOS>+<Footer flatbuffer> <i32 footer length> 'ARROW1'+@++Each footer 'Block' points at the continuation marker of its message,+with the header length (marker, length, metadata, padding) and the+padded body length. The EOS marker is kept so the payload also parses+as a stream.+-}+renderFile :: Schema -> [Frame] -> [Frame] -> ByteString+renderFile sch dictFrames batchFrames =+  let !frames = schemaFrame sch : dictFrames ++ batchFrames+      !footer = fileFooter sch dictFrames batchFrames+      !total = foldl (\t fr -> t + frameSize fr) (BS.length fileMagic + BS.length eosBytes + BS.length footer) frames+  in BSI.unsafeCreate total $ \p -> do+       copyInto p fileMagic+       end <- foldM (pokeFrame p) (BS.length fileMagic) frames+       copyInto (p `plusPtr` end) eosBytes+       copyInto (p `plusPtr` (end + BS.length eosBytes)) footer+++-- | 'renderFile' as lazy chunks aliasing the body pieces (see 'renderStreamLazy').+renderFileLazy :: Schema -> [Frame] -> [Frame] -> BL.ByteString+renderFileLazy sch dictFrames batchFrames =+  BL.fromChunks+    ( fileMagic+        : concatMap frameChunks (schemaFrame sch : dictFrames ++ batchFrames)+        ++ [eosBytes, fileFooter sch dictFrames batchFrames]+    )+++-- ============================================================+-- Low-level writers over pre-built bodies+-- ============================================================++-- | Emit a complete Arrow IPC stream (schema + batches + EOS).+writeArrowStreamFB :: Schema -> [(RecordBatchDef, ByteString)] -> ByteString+writeArrowStreamFB sch = writeArrowStreamFBWithDicts sch []+++-- | Emit a stream with the dictionary batches (in the order given)+-- ahead of the record batches. Each batch carries an @id@ that must+-- match a @DictionaryEncoding.id@ of the schema.+writeArrowStreamFBWithDicts :: Schema -> [DictBatch] -> [(RecordBatchDef, ByteString)] -> ByteString+writeArrowStreamFBWithDicts sch dicts batches =+  renderStream (schemaFrame sch : map dictBatchFrame dicts ++ map (\(rb, body) -> rawBatchFrame Nothing rb body) batches)+++dictBatchFrame :: DictBatch -> Frame+dictBatchFrame (DictBatch did isDelta rb body) = rawBatchFrame (Just (DictHeader did isDelta)) rb body+++-- | Arrow IPC file with the given record batches (see 'renderFile').+writeArrowFileFB :: Schema -> [(RecordBatchDef, ByteString)] -> ByteString+writeArrowFileFB sch = writeArrowFileFBWithDicts sch []+++-- | Arrow IPC file with dictionary batches indexed in the footer.+writeArrowFileFBWithDicts :: Schema -> [DictBatch] -> [(RecordBatchDef, ByteString)] -> ByteString+writeArrowFileFBWithDicts sch dicts batches =+  renderFile sch (map dictBatchFrame dicts) (map (\(rb, body) -> rawBatchFrame Nothing rb body) batches)+++-- | The header and a contiguous body (one allocation) for a batch; see 'buildRecordBatchBytesWith'.+buildRecordBatchBytes :: Schema -> V.Vector ColumnArray -> Either String (RecordBatchDef, ByteString)+buildRecordBatchBytes = buildRecordBatchBytesWith Nothing+++{- | Lay out a batch (spec buffer layout, 8-byte aligned buffers) and+copy its pieces into one contiguous body. With a codec each buffer is+compressed independently (Arrow's @BodyCompression@ = @BUFFER@) and+'rbBodyCompression' is set; a codec that was not compiled in is ignored.+The columns are not checked against the schema beyond what the layout+needs; see 'Arrow.Write.Columns.validateColumns'.+-}+buildRecordBatchBytesWith+  :: Maybe BodyCompressionCodec+  -> Schema+  -> V.Vector ColumnArray+  -> Either String (RecordBatchDef, ByteString)+buildRecordBatchBytesWith mCodec sch cols = do+  eb <- encodeBatchWithEndianness (arrowEndianness sch) mCodec (arrowFields sch) cols+  Right (ebDef eb, concatPieces (ebBodyLength eb) (ebPieces eb))+++concatPieces :: Int -> V.Vector ByteString -> ByteString+concatPieces total pieces = BSI.unsafeCreate total $ \p -> do+  _ <- pokeFrame p 0 (Frame BS.empty pieces total)+  pure ()+++-- | A stream from column batches (no schema checks; see "Arrow.Stream" for the checked writer).+writeArrowStreamFBFromColumns :: Schema -> V.Vector (V.Vector ColumnArray) -> Either String ByteString+writeArrowStreamFBFromColumns sch batches = do+  ebs <- traverse (encodeBatchWithEndianness (arrowEndianness sch) Nothing (arrowFields sch)) (V.toList batches)+  Right (renderStream (schemaFrame sch : map batchFrame ebs))+++-- ============================================================+-- Body compression+-- ============================================================++{- | Apply Arrow's @BodyCompression = BUFFER@ scheme to a body: every+buffer slice is replaced by its envelope (see 'compressPiece') and the+buffers are laid out again, 8-aligned.+-}+compressBody+  :: BodyCompressionCodec+  -> VS.Vector Buffer+  -> ByteString+  -> Either String (VS.Vector Buffer, ByteString)+compressBody codec bufs body = do+  let slice (Buffer o l) = BS.take (fromIntegral l) (BS.drop (fromIntegral o) body)+  pieces <- V.mapM (compressPiece codec . slice) (V.convert bufs)+  let !(bufs', total) = layoutBuffers pieces+  Right (bufs', concatPieces total pieces)+++{- | One buffer per Arrow's @BUFFER@ method: an 8-byte little-endian+uncompressed length followed by the compressed bytes, or length @-1@+and the raw bytes when compression does not shrink them (the spec's+escape hatch). Empty buffers stay empty.+-}+compressPiece :: BodyCompressionCodec -> ByteString -> Either String ByteString+compressPiece codec raw+  | BS.null raw = Right raw+  | otherwise = do+      compressed <- compressBufferEither codec raw+      let !rawLen = BS.length raw+          !(prefix, payload)+            | BS.length compressed >= rawLen = (-1, raw)+            | otherwise = (fromIntegral rawLen, compressed)+      Right $ BSI.unsafeCreate (8 + BS.length payload) $ \p -> do+        pokeByteOff p 0 (prefix :: Int64)+        copyInto (p `plusPtr` 8) payload+++-- | Compress a single buffer's bytes. Routes to the right+-- codec backend; @-fzstd@ / @-flz4@ Cabal flags select+-- availability.+--+-- Returns 'Left' when the requested codec wasn't compiled in.+compressBufferEither+  :: BodyCompressionCodec -> ByteString+  -> Either String ByteString+compressBufferEither codec bs = case codec of+#ifdef HAVE_ZSTD+  BodyZstd ->+    Right (Zstd.compress 3 bs)   -- level 3 matches arrow-cpp's default+#else+  BodyZstd -> Left "Arrow.FlatBufferIPC: ZSTD body compression requires building wireform-arrow with -fzstd"+#endif+#ifdef HAVE_LZ4+  LZ4Frame ->+    -- lz4-hs's Codec.Lz4.compress produces the official+    -- LZ4_Frame format (magic 0x184D2204 + frame descriptor ++    -- one-or-more blocks), which is what arrow-cpp / pyarrow+    -- consume for BodyCompression codec=0 (LZ4_FRAME). The API+    -- works on lazy ByteStrings under the hood.+    Right (BL.toStrict (Lz4.compress (BL.fromStrict bs)))+#else+  LZ4Frame -> Left "Arrow.FlatBufferIPC: LZ4 body compression requires building wireform-arrow with -flz4"+#endif++-- | Predicate the caller can use to fail-fast before invoking+-- the writer on a codec that isn't compiled in.+bodyCompressionAvailable :: BodyCompressionCodec -> Bool+bodyCompressionAvailable = \case+#ifdef HAVE_ZSTD+  BodyZstd -> True+#else+  BodyZstd -> False+#endif+#ifdef HAVE_LZ4+  LZ4Frame -> True+#else+  LZ4Frame -> False+#endif++-- | Build the @Footer@ flatbuffer:+--+-- @+-- table Footer {+--   version       : MetadataVersion;   // i16+--   schema        : Schema;            // table uoffset+--   dictionaries  : [Block];           // vector of structs (struct size = 24 with padding)+--   recordBatches : [Block];+--   custom_metadata : [KeyValue];+-- }+-- @+--+-- @Block@ is a struct with layout+-- @offset: i64; metaDataLength: i32; bodyLength: i64;@. The+-- @metaDataLength@ field is 4 bytes wide but the struct is+-- 8-aligned, so each Block occupies 24 bytes (8 + 4 + 4 padding ++-- 8). FlatBuffers structs are compiler-generated, but we hand-roll+-- the same layout here.+buildFileFooter :: Schema -> [ArrowBlock] -> [ArrowBlock] -> ByteString+buildFileFooter sch dictBlocks rbBlocks = unsafePerformIO $ do+  b <- newBuilder+  schUOff <- writeSchema b sch+  rbVec   <- writeVectorOfStructs b 24 8 (map writeBlockStruct rbBlocks)+  dictVec <- if null dictBlocks+               then pure Nothing+               else Just <$> writeVectorOfStructs b 24 8+                                (map writeBlockStruct dictBlocks)+  msgUOff <- writeTable b+    [ Just (scalar 2 (\bb -> prependI16 bb metadataVersionV5))+    , Just (voff schUOff)+    , case dictVec of { Nothing -> Nothing; Just uo -> Just (voff uo) }+    , Just (voff rbVec)+    , Nothing                        -- custom_metadata+    ]+  finish b msgUOff+{-# NOINLINE buildFileFooter #-}++writeBlockStruct :: ArrowBlock -> Builder -> IO ()+writeBlockStruct (ArrowBlock o ml bl) bb = do+  -- Reverse order: bodyLength (i64), 4-byte pad, metaDataLength+  -- (i32), offset (i64).+  prependI64 bb bl+  prependBS  bb (BS.replicate 4 0)+  prependI32 bb ml+  prependI64 bb o+++-- | Encapsulated frame (continuation marker, length, FlatBuffers+-- @Message@ metadata, padding) for one 'Message'. The body is not+-- included: the metadata's @bodyLength@ is the 8-aligned extent of+-- the batch's buffers, and the caller appends that many body bytes.+encodeMessageFrame :: Message -> ByteString+encodeMessageFrame = \case+  SchemaMessage sch -> encapsulateMessage (buildSchemaMessage sch) BS.empty+  RecordBatch rb -> batchMessage True Nothing rb (bodyExtent rb)+  DictionaryBatch did isDelta rb -> batchMessage True (Just (DictHeader did isDelta)) rb (bodyExtent rb)+  where+    bodyExtent rb =+      let !end = VS.foldl' (\m buf -> max m (bufOffset buf + bufLength buf)) 0 (rbBuffers rb)+      in (end + 7) .&. complement 7
+ src/Arrow/IPC.hs view
@@ -0,0 +1,37 @@+{- | Apache Arrow IPC message framing.++An encapsulated Arrow IPC message is a continuation marker+(@0xFFFFFFFF@), a 4-byte little-endian metadata length, the+FlatBuffers @Message@ table (padded to 8 bytes) and then the body.+'encodeIPCMessage' / 'decodeIPCMessage' handle the metadata frame for+one 'Message' in exactly the format pyarrow, arrow-cpp and arrow-rs use+(they delegate to "Arrow.FlatBufferIPC"); bodies are appended / read by+the caller. "Arrow.Stream" and "Arrow.File" build whole streams and+files on top.+-}+module Arrow.IPC (+  encodeIPCMessage,+  decodeIPCMessage,+  validateRecordBatchBuffers,+) where++import Arrow.FlatBufferIPC (decodeMessageFrame, encodeMessageFrame)+import Arrow.Read.Columns (validateRecordBatchBuffers)+import Arrow.Types (Message)+import Data.ByteString (ByteString)+++{- | Encapsulated metadata frame for one message. For record and+dictionary batches the declared body length is the 8-aligned extent of+the batch's buffers; append that many body bytes to form the full+message.+-}+encodeIPCMessage :: Message -> ByteString+encodeIPCMessage = encodeMessageFrame+++{- | Decode the metadata frame at the start of the input (bytes after+the frame, such as the body, are ignored).+-}+decodeIPCMessage :: ByteString -> Either String Message+decodeIPCMessage bs = fst <$> decodeMessageFrame bs
+ src/Arrow/Read/Columns.hs view
@@ -0,0 +1,1024 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | Decode Arrow IPC record batch and dictionary batch bodies into+'ColumnArray's without copying.++Every buffer of the result aliases the message body (or, for a+compressed body, the freshly decompressed buffer). A buffer is copied+only when its address is not aligned for the element type (arrow-rs+@align_buffers@) or when the schema declares big-endian data, which is+byte-swapped into a fresh buffer.++Decoding validates everything arrow-rs validates by default, in O(1)+Haskell work per array plus C kernels: buffer bounds, buffer sizes+against array lengths, null counts against the validity bitmap+(popcount), offsets (monotonic, in range), UTF-8 and character+boundaries, list views, union type ids and dense offsets, run ends,+views. Dictionary keys are checked against their values when the+dictionary is resolved ('Arrow.Column.resolveDictionaryColumn'). A+corrupt or hostile batch is a 'Left', never an out-of-bounds read.++The body must use the IPC specification's buffer layout: a validity+slot for every array except null, union and run-end-encoded arrays+(an empty slot when the array has no nulls).+-}+module Arrow.Read.Columns (+  decodeRecordBatch,+  decodeDictionaryBatch,++  -- * Decode budgets+  ReadOptions (..),+  defaultReadOptions,++  -- * Body compression+  decompressBody,++  -- * Buffer bounds+  validateRecordBatchBuffers,+) where++import Arrow.Column (columnLength, placeholderColumn)+import Arrow.Column.Internal (+  ColumnArray,+  IntegralPrim (..),+  Offset,+  PrimType (..),+  SomePrimType (..),+  integralPrim,+  primTypeFor,+  withPrim,+ )+import Arrow.Column.Internal qualified as I+import Arrow.FlatBufferIPC.Common (DictBatch (..))+import Arrow.Types (+  ArrowType (..),+  BodyCompressionCodec (..),+  Buffer (..),+  DictionaryEncoding (..),+  Endianness (..),+  Field (..),+  FieldNode (..),+  RecordBatchDef (..),+  Schema (..),+  UnionMode (..),+ )+import Control.DeepSeq (NFData)+import Data.Bits (shiftL, (.|.))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Unsafe qualified as BSU+import Data.Int (Int32, Int64, Int8)+import Data.Text (Text)+import Data.Vector qualified as V+import Data.Vector.Storable qualified as VS+import Foreign.Marshal.Utils (copyBytes, fillBytes)+import Foreign.Ptr (castPtr, plusPtr)+import Foreign.Storable (Storable (..))+import GHC.Generics (Generic)+import System.IO.Unsafe (unsafeDupablePerformIO)+import Wireform.FFI (validateArrowBuffers)++#ifdef HAVE_ZSTD+import qualified Codec.Compression.Zstd as Zstd+#endif++#ifdef HAVE_LZ4+import qualified Codec.Lz4 as Lz4+import qualified Control.Exception as Exc+import qualified Data.ByteString.Lazy as BL+#endif+++-- ============================================================+-- Reader options+-- ============================================================++{- | Caller-supplied decode budgets. Every entry point that turns wire+bytes into 'ColumnArray's takes one, mirroring arrow-cpp's+@IpcReadOptions@: an application that reads untrusted input passes+tight limits and gets a decode 'Left' instead of an oversized+allocation or an absurd logical row count. A batch that fits its+declared @length@ in real body bytes is still bounded by the input+size (columns alias the body); the budgets cap what the wire+/says/:++  * 'roMaxRows' caps the record batch's declared row count.+  * 'roMaxValues' caps the total declared length across every field+    node of one batch, i.e. the logical values a consumer walking the+    columns would visit (a run-end-encoded or list column can name+    billions of rows from a few hundred bytes).+  * 'roMaxUnbackedRows' caps rows claimed with no body bytes backing+    them at all (nullable structs or fixed-size lists over null or+    run-end-encoded children, zero-width fixed-size binary).+  * 'roMaxDecompressedBytes' caps the total uncompressed bytes a+    compressed body may expand to.++'defaultReadOptions' is generous enough for any batch that can fit in+memory while keeping every budget finite.+-}+data ReadOptions = ReadOptions+  { roMaxRows :: !Int64+  -- ^ Largest row count ('rbLength') one record batch may declare.+  , roMaxValues :: !Int64+  -- ^ Largest total declared length across the field nodes of one batch.+  , roMaxUnbackedRows :: !Int64+  -- ^ Largest number of rows a batch may claim with no buffer backing.+  , roMaxDecompressedBytes :: !Int64+  -- ^ Largest total uncompressed size a compressed body may expand to.+  }+  deriving stock (Show, Eq, Generic)+  deriving anyclass (NFData)+++{- | The budgets used when a caller has no stricter policy: same row+cap as 'maxArrayLength', @2^34@ total declared values (a batch+claiming more than that cannot be backed by body bytes that fit in+memory), @2^20@ unbacked rows, and @2^34@ decompressed bytes.+-}+defaultReadOptions :: ReadOptions+defaultReadOptions =+  ReadOptions+    { roMaxRows = fromIntegral maxArrayLength+    , roMaxValues = 1 `shiftL` 34+    , roMaxUnbackedRows = 1 `shiftL` 20+    , roMaxDecompressedBytes = 1 `shiftL` 34+    }+++{- | Check the wire-claim budgets that apply before any buffer is+touched: the declared row count and the total declared value count+across all field nodes.+-}+checkBatchBudgets :: ReadOptions -> RecordBatchDef -> Either String ()+checkBatchBudgets opts rb = do+  let !rows = rbLength rb+  if rows < 0 || rows > roMaxRows opts+    then+      Left+        ( "Arrow.Read.Columns: record batch declares "+            ++ show rows+            ++ " rows (limit "+            ++ show (roMaxRows opts)+            ++ ")"+        )+    else Right ()+  let !values = nodeLengthSum (rbNodes rb)+  if values > roMaxValues opts+    then+      Left+        ( "Arrow.Read.Columns: record batch field nodes declare "+            ++ show values+            ++ " values in total (limit "+            ++ show (roMaxValues opts)+            ++ ")"+        )+    else Right ()+++-- | Sum of the field node lengths, saturating at @maxBound@.+nodeLengthSum :: VS.Vector FieldNode -> Int64+nodeLengthSum = VS.foldl' step 0+  where+    step !acc fn+      | acc >= maxBound = maxBound+      | otherwise =+          let !l = max 0 (fnLength fn)+          in if acc > maxBound - l then maxBound else acc + l+++-- ============================================================+-- Entry points+-- ============================================================++{- | Decode one record batch: the schema's top-level fields, in order,+from the batch header and its body. Body compression is undone first+(per buffer). Every top-level column has exactly 'rbLength' rows.++Dictionary-encoded columns come back as @ColDictionary id keys+placeholder@: the keys at their wire width with the row validity, and+an empty placeholder for the values; resolve them with+'Arrow.Column.resolveDictionaryColumn'.+-}+decodeRecordBatch :: ReadOptions -> Schema -> RecordBatchDef -> ByteString -> Either String (V.Vector ColumnArray)+decodeRecordBatch opts schema rb body = do+  checkBatchBudgets opts rb+  src <- bodySource opts rb body+  let !ctx = Ctx (arrowEndianness schema == Big) rb src opts+  Outs cols _ <- fieldsN ctx (arrowFields schema) (Cur 0 0 0 0)+  case V.findIndex (\c -> fromIntegral (columnLength c) /= rbLength rb) cols of+    Nothing -> Right cols+    Just k ->+      Left+        ( "Arrow.Read.Columns: column #"+            ++ show k+            ++ " has "+            ++ show (columnLength (V.unsafeIndex cols k))+            ++ " rows, the record batch declares "+            ++ show (rbLength rb)+        )+++{- | Decode the values of a dictionary batch. The values field is the+schema field (at any depth) whose dictionary encoding carries the+batch's id, with the encoding removed. Returns the id and the values+column; dictionaries nested inside the values stay unresolved+placeholders.+-}+decodeDictionaryBatch :: ReadOptions -> Schema -> DictBatch -> Either String (Int64, ColumnArray)+decodeDictionaryBatch opts sch db = do+  f <- case findDictField (dbId db) (arrowFields sch) of+    Just f -> Right f+    Nothing ->+      Left+        ( "Arrow.Read.Columns: dictionary batch with id "+            ++ show (dbId db)+            ++ " doesn't match any field in the schema"+        )+  let !valuesField = f {fieldDictionary = Nothing}+      !inner = sch {arrowFields = V.singleton valuesField, arrowMetadata = V.empty, arrowFeatures = V.empty}+  cols <- decodeRecordBatch opts inner (dbData db) (dbBody db)+  if V.length cols == 1+    then Right (dbId db, V.unsafeHead cols)+    else Left "Arrow.Read.Columns: dictionary batch must hold exactly one column"+++-- | The field whose dictionary encoding carries the id (depth first).+findDictField :: Int64 -> V.Vector Field -> Maybe Field+findDictField did = go+  where+    go fs = V.foldr (\f rest -> pick f rest) Nothing fs+    pick f rest = case fieldDictionary f of+      Just de | deId de == did -> Just f+      _ -> case go (fieldChildren f) of+        Just g -> Just g+        Nothing -> rest+++{- | Validate that all buffer offset/length pairs in a 'RecordBatchDef' are+non-negative, within the given body length, and non-overlapping.+SIMD-accelerated; reads the Storable buffer vector in place.+-}+validateRecordBatchBuffers :: RecordBatchDef -> Int64 -> Bool+validateRecordBatchBuffers rb bodyLen+  | n == 0 = True+  | otherwise = unsafeDupablePerformIO $ VS.unsafeWith bufs $ \p ->+      pure $! validateArrowBuffers (castPtr p) n bodyLen+  where+    !bufs = rbBuffers rb+    !n = VS.length bufs+++-- ============================================================+-- Body sources+-- ============================================================++-- | Where buffer @i@ lives.+data Source+  = -- | Uncompressed body; buffer descriptors already bounds-checked.+    Plain !ByteString+  | -- | One decompressed buffer per descriptor.+    Unpacked !(V.Vector ByteString)+++bodySource :: ReadOptions -> RecordBatchDef -> ByteString -> Either String Source+bodySource opts rb body = case rbBodyCompression rb of+  Nothing+    | validateRecordBatchBuffers rb (fromIntegral (BS.length body)) -> Right (Plain body)+    | otherwise -> Left "Arrow.Read.Columns: invalid buffer bounds in RecordBatchDef"+  Just codec -> Unpacked <$> decompressBuffers opts codec (rbBuffers rb) body+++-- ============================================================+-- Walk state+-- ============================================================++data Ctx = Ctx+  { ctxSwap :: !Bool+  -- ^ Big-endian body: multi-byte values are byte-swapped into fresh buffers.+  , ctxRb :: !RecordBatchDef+  , ctxSrc :: !Source+  , ctxOpts :: !ReadOptions+  }+++-- | Next field node, next buffer, next variadic-count entry, unbacked rows so far.+data Cur = Cur+  { curNode :: {-# UNPACK #-} !Int+  , curBuf :: {-# UNPACK #-} !Int+  , curVar :: {-# UNPACK #-} !Int+  , curUnbacked :: {-# UNPACK #-} !Int+  }+++data Out = Out !ColumnArray {-# UNPACK #-} !Cur+++data Outs = Outs !(V.Vector ColumnArray) {-# UNPACK #-} !Cur+++-- | Advance past @nodes@ field nodes and @bufs@ buffers.+adv :: Int -> Int -> Cur -> Cur+adv nodes bufs (Cur n b v u) = Cur (n + nodes) (b + bufs) v u+{-# INLINE adv #-}+++{- | Upper bound on any array length taken from a field node. Far above+any batch that fits in memory, and small enough that @len * 32@ and+@(len + 1) * 8@ cannot overflow 'Int'.+-}+maxArrayLength :: Int+maxArrayLength = 1 `shiftL` 40+++-- | Field node @i@ with its length checked to lie in @[0, 'maxArrayLength']@.+nodeAt :: Ctx -> Int -> Either String FieldNode+nodeAt ctx i = case rbNodes (ctxRb ctx) VS.!? i of+  Nothing ->+    Left+      ( "Arrow.Read.Columns: record batch has "+          ++ show (VS.length (rbNodes (ctxRb ctx)))+          ++ " field nodes, schema needs node #"+          ++ show i+      )+  Just fn+    | fnLength fn < 0 || fnLength fn > fromIntegral maxArrayLength ->+        Left ("Arrow.Read.Columns: invalid field node length " ++ show (fnLength fn))+    | otherwise -> Right fn+{-# INLINE nodeAt #-}+++-- | Bytes of buffer @i@.+bufferAt :: Ctx -> Int -> Either String ByteString+bufferAt ctx i+  | i < 0 || i >= VS.length bufs =+      Left+        ( "Arrow.Read.Columns: record batch has "+            ++ show (VS.length bufs)+            ++ " buffers, schema layout needs buffer #"+            ++ show i+        )+  | otherwise = case ctxSrc ctx of+      Unpacked v -> Right (V.unsafeIndex v i)+      Plain body ->+        let Buffer o l = VS.unsafeIndex bufs i+        in if o < 0 || l < 0 || o > fromIntegral (BS.length body) - l+             then Left "Arrow.Read.Columns: buffer outside the body"+             else Right $! BSU.unsafeTake (fromIntegral l) (BSU.unsafeDrop (fromIntegral o) body)+  where+    !bufs = rbBuffers (ctxRb ctx)+{-# INLINE bufferAt #-}+++-- ============================================================+-- Schema walk+-- ============================================================++fieldsN :: Ctx -> V.Vector Field -> Cur -> Either String Outs+fieldsN ctx fields cur0 = go 0 cur0 []+  where+    !n = V.length fields+    go !i !cur acc+      | i >= n = Right (Outs (V.fromListN n (reverse acc)) cur)+      | otherwise = do+          Out c cur' <- field ctx (V.unsafeIndex fields i) cur+          go (i + 1) cur' (c : acc)+++field :: Ctx -> Field -> Cur -> Either String Out+field ctx f cur = do+  fn <- nodeAt ctx (curNode cur)+  let !len = fromIntegral (fnLength fn) :: Int+      !lim = roMaxUnbackedRows (ctxOpts ctx)+      !ub = if allocatesUnbacked f then curUnbacked cur + len else curUnbacked cur+  if fromIntegral ub > lim+    then+      Left+        ( "Arrow.Read.Columns: record batch claims "+            ++ show ub+            ++ " rows with no buffer backing them (limit "+            ++ show lim+            ++ ")"+        )+    else fieldAt ctx f fn len cur {curUnbacked = ub}+++fieldAt :: Ctx -> Field -> FieldNode -> Int -> Cur -> Either String Out+fieldAt ctx f fn len cur = case fieldDictionary f of+  Just de -> dictKeys ctx f de fn len cur+  Nothing -> case fieldType f of+    ANull -> Right (Out (I.ColNull len) (adv 1 0 cur))+    ABool -> do+      v <- validityAt ctx "bool" len fn bi+      bs <- bufferAt ctx (bi + 1)+      bits <- case I.mkBitmap bs 0 len of+        Right b -> Right b+        Left _ -> Left "Arrow.Read.Columns: bool data buffer too small"+      Right (Out (I.ColBool v bits) (adv 1 2 cur))+    AUtf8 -> varLen @Int32 ctx "utf8" True I.ColUtf8 fn len cur+    ABinary -> varLen @Int32 ctx "binary" False I.ColBinary fn len cur+    ALargeUtf8 -> varLen @Int64 ctx "large utf8" True I.ColLargeUtf8 fn len cur+    ALargeBinary -> varLen @Int64 ctx "large binary" False I.ColLargeBinary fn len cur+    AFixedSizeBinary w -> fixedSizeBinary ctx w fn len cur+    AUtf8View -> views ctx True fn len cur+    ABinaryView -> views ctx False fn len cur+    AStruct -> struct ctx f fn len cur+    AList -> list @Int32 ctx "list" I.ColList f fn len cur+    ALargeList -> list @Int64 ctx "large list" I.ColLargeList f fn len cur+    AListView -> listView @Int32 ctx "list view" I.ColListView f fn len cur+    ALargeListView -> listView @Int64 ctx "large list view" I.ColLargeListView f fn len cur+    AFixedSizeList size -> fixedSizeList ctx size f fn len cur+    AMap _ -> mapCol ctx f fn len cur+    AUnion mode ids -> union ctx mode ids f len cur+    ARunEndEncoded -> runEndEncoded ctx f len cur+    ty -> case primTypeFor ty of+      Just (SomePrimType t) -> do+        c <- prim ctx t fn len bi+        Right (Out c (adv 1 2 cur))+      Nothing -> Left ("Arrow.Read.Columns: unsupported type: " ++ show ty)+  where+    !bi = curBuf cur+++-- | Does this array claim rows that no buffer backs?+allocatesUnbacked :: Field -> Bool+allocatesUnbacked f+  | Just _ <- fieldDictionary f = False+  | otherwise = case fieldType f of+      AStruct -> fieldNullable f && not (lengthBacked f)+      AFixedSizeList _ -> fieldNullable f && not (lengthBacked f)+      AFixedSizeBinary 0 -> True+      _ -> False+++{- | Is the length of an array of this field paid for by bytes in some+buffer of the array or its descendants? Null arrays have no buffers,+and a run-end-encoded array's logical length can be any value its last+run end names, so neither backs a parent.+-}+lengthBacked :: Field -> Bool+lengthBacked f+  | Just _ <- fieldDictionary f = True+  | otherwise = case fieldType f of+      ANull -> False+      ARunEndEncoded -> False+      AStruct -> V.any lengthBacked (fieldChildren f)+      AFixedSizeList n -> n > 0 && V.any lengthBacked (fieldChildren f)+      AFixedSizeBinary w -> w > 0+      _ -> True+++-- ============================================================+-- Buffers+-- ============================================================++{- | The validity of an array of @len@ rows from its slot. No nulls is+'Nothing' (whatever the slot holds, as in arrow-rs); otherwise the+bitmap must cover @len@ bits and its clear bits must equal the node's+null count (C popcount).+-}+validityAt :: Ctx -> String -> Int -> FieldNode -> Int -> Either String (Maybe I.Validity)+validityAt ctx what len fn bi+  | nc == 0 = bufferAt ctx bi >> Right Nothing+  | nc < 0 || nc > len = Left ("Arrow.Read.Columns: " ++ what ++ ": null count " ++ show nc ++ " for " ++ show len ++ " rows")+  | otherwise = do+      bs <- bufferAt ctx bi+      let !nbytes = (len + 7) `quot` 8+      if BS.length bs < nbytes+        then Left ("Arrow.Read.Columns: " ++ what ++ ": validity bitmap too small for " ++ show len ++ " rows")+        else do+          let !b = I.Bitmap (BSU.unsafeTake nbytes bs) 0 len+              !nulls = len - I.bitmapSetCount b+          if nulls /= nc+            then+              Left+                ( "Arrow.Read.Columns: "+                    ++ what+                    ++ ": null count "+                    ++ show nc+                    ++ " does not match the validity bitmap ("+                    ++ show nulls+                    ++ ")"+                )+            else Right (Just (I.Validity b nc))+  where+    !nc = fromIntegral (fnNullCount fn) :: Int+++-- | @n@ elements of @w@ bytes from buffer @i@, aligned and in native byte order.+storableAt :: forall a. Storable a => Ctx -> String -> [Int] -> Int -> Int -> Either String (VS.Vector a)+storableAt ctx what swapPat n i = do+  bs <- bufferAt ctx i+  let !w = sizeOf (undefined :: a)+      !need = n * w+  if BS.length bs < need+    then Left ("Arrow.Read.Columns: " ++ what ++ " buffer too small for " ++ show n ++ " values")+    else+      let !exact = BSU.unsafeTake need bs+      in Right $!+           if ctxSwap ctx && not (null swapPat)+             then I.unsafeBytesToStorable (I.swapCopy swapPat w n exact)+             else I.bytesToStorable exact+{-# INLINE storableAt #-}+++{- | @rows + 1@ offsets from buffer @i@. A zero-row array may omit the+offsets buffer (it then has the single offset 0).+-}+offsetsAt :: forall o. Offset o => Ctx -> String -> Int -> Int -> Either String (VS.Vector o)+offsetsAt ctx what rows i = do+  bs <- bufferAt ctx i+  if rows == 0 && BS.length bs < sizeOf (undefined :: o)+    then Right (VS.singleton 0)+    else storableAt ctx (what ++ " offsets") [sizeOf (undefined :: o)] (rows + 1) i+{-# INLINE offsetsAt #-}+++-- ============================================================+-- Leaf arrays+-- ============================================================++-- | A fixed-width array: validity at @bi@, values at @bi + 1@.+prim :: Ctx -> PrimType a -> FieldNode -> Int -> Int -> Either String ColumnArray+prim ctx t fn len bi = withPrim t $ do+  v <- validityAt ctx (I.primTypeName t) len fn bi+  xs <- storableAt ctx (I.primTypeName t) (I.swapPattern t) len (bi + 1)+  Right (I.ColPrim t v xs)+++varLen+  :: forall o+   . Offset o+  => Ctx+  -> String+  -> Bool+  -> (Maybe I.Validity -> VS.Vector o -> ByteString -> ColumnArray)+  -> FieldNode+  -> Int+  -> Cur+  -> Either String Out+varLen ctx what utf8 con fn len cur = do+  let !bi = curBuf cur+  v <- validityAt ctx what len fn bi+  offs <- offsetsAt @o ctx what len (bi + 1)+  dat <- bufferAt ctx (bi + 2)+  I.validateOffsets what (BS.length dat) offs+  if utf8 then I.validateUtf8 what offs dat else Right ()+  Right (Out (con v offs dat) (adv 1 3 cur))+{-# INLINE varLen #-}+++fixedSizeBinary :: Ctx -> Int -> FieldNode -> Int -> Cur -> Either String Out+fixedSizeBinary ctx w fn len cur+  | w < 0 = Left ("Arrow.Read.Columns: negative fixed-size binary width " ++ show w)+  | w > 0 && len > maxArrayLength `quot` w =+      Left ("Arrow.Read.Columns: fixed-size binary of width " ++ show w ++ " cannot have " ++ show len ++ " rows")+  | otherwise = do+      let !bi = curBuf cur+      v <- validityAt ctx "fixed-size binary" len fn bi+      dat <- bufferAt ctx (bi + 1)+      if BS.length dat < w * len+        then Left "Arrow.Read.Columns: fixed-size binary data buffer too small"+        else Right (Out (I.ColFixedSizeBinary w len v (BSU.unsafeTake (w * len) dat)) (adv 1 2 cur))+++{- | Utf8View / BinaryView: validity, the 16-byte views, then this+column's variadic data buffers (their count is the next+'rbVariadicBufferCounts' entry, in preorder over view columns).+-}+views :: Ctx -> Bool -> FieldNode -> Int -> Cur -> Either String Out+views ctx utf8 fn len cur = do+  let !rb = ctxRb ctx+      !bi = curBuf cur+      !what = if utf8 then "utf8 view" else "binary view"+  nVar <- case rbVariadicBufferCounts rb V.!? curVar cur of+    Nothing -> Left "Arrow.Read.Columns: view column has no variadicBufferCounts entry"+    Just c+      | c < 0 || c > fromIntegral (VS.length (rbBuffers rb)) ->+          Left ("Arrow.Read.Columns: invalid variadic buffer count " ++ show c)+      | otherwise -> Right (fromIntegral c :: Int)+  v <- validityAt ctx what len fn bi+  vs <- bufferAt ctx (bi + 1)+  if BS.length vs < len * 16+    then Left ("Arrow.Read.Columns: " ++ what ++ " buffer too small")+    else Right ()+  dataBufs <- V.generateM nVar (\k -> bufferAt ctx (bi + 2 + k))+  let !rawViews = BSU.unsafeTake (len * 16) vs+      !viewBytes = if ctxSwap ctx then I.swapViewCopy True rawViews else rawViews+  I.validateViews what utf8 len v viewBytes dataBufs+  let !col = if utf8 then I.ColUtf8View v viewBytes dataBufs else I.ColBinaryView v viewBytes dataBufs+      Cur n b x u = cur+  Right (Out col (Cur (n + 1) (b + 2 + nVar) (x + 1) u))+++-- ============================================================+-- Nested arrays+-- ============================================================++singleChild :: String -> Field -> Either String Field+singleChild what f+  | V.length (fieldChildren f) == 1 = Right (V.unsafeHead (fieldChildren f))+  | otherwise =+      Left ("Arrow.Read.Columns: " ++ what ++ " field must have exactly one child, has " ++ show (V.length (fieldChildren f)))+++-- | Every child of a struct-like parent must cover the parent's rows.+checkChildLengths :: String -> Int -> V.Vector ColumnArray -> Either String ()+checkChildLengths what len cols = case V.findIndex (\c -> columnLength c < len) cols of+  Nothing -> Right ()+  Just k ->+    Left+      ( "Arrow.Read.Columns: "+          ++ what+          ++ " child #"+          ++ show k+          ++ " has "+          ++ show (columnLength (V.unsafeIndex cols k))+          ++ " rows, parent has "+          ++ show len+      )+++struct :: Ctx -> Field -> FieldNode -> Int -> Cur -> Either String Out+struct ctx f fn len cur = do+  v <- validityAt ctx "struct" len fn (curBuf cur)+  Outs cols cur' <- fieldsN ctx (fieldChildren f) (adv 1 1 cur)+  checkChildLengths "struct" len cols+  let named = V.zipWith (\c col -> (fieldName c, col)) (fieldChildren f) cols :: V.Vector (Text, ColumnArray)+  Right (Out (I.ColStruct len v named) cur')+++list+  :: forall o+   . Offset o+  => Ctx+  -> String+  -> (Maybe I.Validity -> VS.Vector o -> ColumnArray -> ColumnArray)+  -> Field+  -> FieldNode+  -> Int+  -> Cur+  -> Either String Out+list ctx what con f fn len cur = do+  let !bi = curBuf cur+  childField <- singleChild what f+  v <- validityAt ctx what len fn bi+  offs <- offsetsAt @o ctx what len (bi + 1)+  Out child cur' <- field ctx childField (adv 1 2 cur)+  I.validateOffsets what (columnLength child) offs+  Right (Out (con v offs child) cur')+++listView+  :: forall o+   . Offset o+  => Ctx+  -> String+  -> (Maybe I.Validity -> VS.Vector o -> VS.Vector o -> ColumnArray -> ColumnArray)+  -> Field+  -> FieldNode+  -> Int+  -> Cur+  -> Either String Out+listView ctx what con f fn len cur = do+  let !bi = curBuf cur+      !w = sizeOf (undefined :: o)+  childField <- singleChild what f+  v <- validityAt ctx what len fn bi+  offs <- storableAt @o ctx (what ++ " offsets") [w] len (bi + 1)+  sizes <- storableAt @o ctx (what ++ " sizes") [w] len (bi + 2)+  Out child cur' <- field ctx childField (adv 1 3 cur)+  I.validateListView what v offs sizes (columnLength child)+  Right (Out (con v offs sizes child) cur')+++fixedSizeList :: Ctx -> Int -> Field -> FieldNode -> Int -> Cur -> Either String Out+fixedSizeList ctx size f fn len cur+  | size < 0 || (size > 0 && len > maxArrayLength `quot` size) =+      Left ("Arrow.Read.Columns: fixed-size list of size " ++ show size ++ " cannot have " ++ show len ++ " rows")+  | otherwise = do+      childField <- singleChild "fixed-size list" f+      v <- validityAt ctx "fixed-size list" len fn (curBuf cur)+      Out child cur' <- field ctx childField (adv 1 1 cur)+      if columnLength child < len * size+        then+          Left+            ( "Arrow.Read.Columns: fixed-size list child has "+                ++ show (columnLength child)+                ++ " rows, needs "+                ++ show (len * size)+            )+        else Right (Out (I.ColFixedSizeList size len v child) cur')+++{- | Map: validity, offsets, then the entries struct (key, value),+which may hold no nulls of its own (arrow-rs @MapArray@).+-}+mapCol :: Ctx -> Field -> FieldNode -> Int -> Cur -> Either String Out+mapCol ctx f fn len cur = do+  let !bi = curBuf cur+  entries <- singleChild "map" f+  v <- validityAt ctx "map" len fn bi+  offs <- offsetsAt @Int32 ctx "map" len (bi + 1)+  Out ent cur' <- field ctx entries (adv 1 2 cur)+  case ent of+    I.ColStruct n Nothing kv+      | V.length kv == 2 -> do+          I.validateOffsets "map" n offs+          Right (Out (I.ColMap v offs (snd (V.unsafeIndex kv 0)) (snd (V.unsafeIndex kv 1))) cur')+    _ -> Left "Arrow.Read.Columns: map entries must be a struct of exactly (key, value) without nulls"+++{- | Unions carry no validity. Wire type ids become child positions+through the declared @typeIds@ (aliased as is when they are+positional, the common case; otherwise one table-lookup pass).+-}+union :: Ctx -> UnionMode -> V.Vector Int32 -> Field -> Int -> Cur -> Either String Out+union ctx mode declared f len cur = do+  let !bi = curBuf cur+      !k = V.length (fieldChildren f)+      !what = "union"+  table <- unionChildTable k declared+  wire <- storableAt @Int8 ctx "union type ids" [] len bi+  let !types = case table of+        Nothing -> wire+        Just tbl -> VS.map (\t -> if t < 0 then -1 else VS.unsafeIndex tbl (fromIntegral t)) wire+  case mode of+    Dense -> do+      offs <- storableAt @Int32 ctx "dense union offsets" [4] len (bi + 1)+      Outs children cur' <- fieldsN ctx (fieldChildren f) (adv 1 2 cur)+      let !lens = VS.generate k (\j -> fromIntegral (columnLength (V.unsafeIndex children j))) :: VS.Vector Int64+      I.validateDenseUnion what types offs lens+      Right (Out (I.ColDenseUnion types offs children) cur')+    Sparse -> do+      Outs children cur' <- fieldsN ctx (fieldChildren f) (adv 1 1 cur)+      checkChildLengths "sparse union" len children+      I.validateSparseUnionTypes what types k+      Right (Out (I.ColSparseUnion types children) cur')+++{- | Lookup table from wire type id (0..127) to child position (@-1@ when+undeclared), or 'Nothing' when the ids are positional (child @k@ has+id @k@), so the wire buffer already holds child positions.+-}+unionChildTable :: Int -> V.Vector Int32 -> Either String (Maybe (VS.Vector Int8))+unionChildTable nChildren declared+  | nChildren > 128 = Left ("Arrow.Read.Columns: union has " ++ show nChildren ++ " children, at most 128 allowed")+  | V.null declared = Right Nothing+  | V.length declared /= nChildren =+      Left+        ( "Arrow.Read.Columns: union declares "+            ++ show (V.length declared)+            ++ " type ids for "+            ++ show nChildren+            ++ " children"+        )+  | V.any (\t -> t < 0 || t > 127) declared = Left "Arrow.Read.Columns: union type ids must be in 0..127"+  | V.and (V.imap (\j t -> fromIntegral t == j) declared) = Right Nothing+  | otherwise =+      let table = VS.generate 128 (\t -> maybe (-1) fromIntegral (V.elemIndex (fromIntegral t) declared)) :: VS.Vector Int8+      in if VS.length (VS.filter (>= 0) table) /= nChildren+           then Left "Arrow.Read.Columns: union declares a type id twice"+           else Right (Just table)+++{- | Run-end encoded: no buffers of its own, children run ends (non-null+int16/32/64, positive, strictly increasing, last >= length) and values+(at least one per run).+-}+runEndEncoded :: Ctx -> Field -> Int -> Cur -> Either String Out+runEndEncoded ctx f len cur+  | V.length (fieldChildren f) /= 2 =+      Left "Arrow.Read.Columns: RunEndEncoded must have exactly two children (run_ends, values)"+  | otherwise = do+      Out re cur1 <- field ctx (V.unsafeIndex (fieldChildren f) 0) (adv 1 0 cur)+      Out vals cur2 <- field ctx (V.unsafeIndex (fieldChildren f) 1) cur1+      I.validateRunEnds "run-end encoded" re len+      let !runs = columnLength re+      if columnLength vals < runs+        then+          Left+            ( "Arrow.Read.Columns: run-end encoded array has "+                ++ show runs+                ++ " runs but "+                ++ show (columnLength vals)+                ++ " values"+            )+        else Right (Out (I.ColRunEndEncoded 0 len re vals) cur2)+++{- | The wire payload of a dictionary-encoded field: the keys (an+integer array at the index type's width, carrying the row validity)+and an empty placeholder for the values.+-}+dictKeys :: Ctx -> Field -> DictionaryEncoding -> FieldNode -> Int -> Cur -> Either String Out+dictKeys ctx f de fn len cur = case primTypeFor (deIndexType de) of+  Just (SomePrimType t)+    | Just IntegralPrim <- integralPrim t -> do+        keys <- prim ctx t fn len (curBuf cur)+        placeholder <- placeholderColumn f+        Right (Out (I.ColDictionary (deId de) keys placeholder) (adv 1 2 cur))+  _ -> Left ("Arrow.Read.Columns: dictionary index type must be an 8/16/32/64-bit integer, got " ++ show (deIndexType de))+++-- ============================================================+-- Body compression+-- ============================================================++{- | Decompress every buffer of a compressed body (one envelope per+buffer descriptor). Each result is a fresh buffer.+-}+decompressBuffers :: ReadOptions -> BodyCompressionCodec -> VS.Vector Buffer -> ByteString -> Either String (V.Vector ByteString)+decompressBuffers opts codec bufs body = do+  checkDecompressedBudget opts bufs body+  V.generateM (VS.length bufs) $ \i -> do+    let Buffer o l = VS.unsafeIndex bufs i+    env <- sliceWire "compressed buffer" o l body+    decompressBufferEnvelope codec env+++{- | Sum the buffer envelopes' declared uncompressed sizes and reject+the body /before/ anything is decompressed when the expansion the+decoder would have to materialize exceeds 'roMaxDecompressedBytes'.+Envelopes marked stored (@-1@) do not expand and are not charged. Each+individual claim is still bounded by 'maxCompressionRatio' against its+payload; this bounds the total.+-}+checkDecompressedBudget :: ReadOptions -> VS.Vector Buffer -> ByteString -> Either String ()+checkDecompressedBudget opts bufs body = go 0 0+  where+    !n = VS.length bufs+    go !i !acc+      | acc > roMaxDecompressedBytes opts =+          Left+            ( "Arrow.Read.Columns: compressed body declares "+                ++ show acc+                ++ " decompressed bytes in total (limit "+                ++ show (roMaxDecompressedBytes opts)+                ++ ")"+            )+      | i >= n = Right ()+      | otherwise = do+          let Buffer o l = VS.unsafeIndex bufs i+          env <- sliceWire "compressed buffer" o l body+          let !claim+                | BS.length env < 8 = 0+                | otherwise = case leInt64 env of+                    r | r > 0 -> r+                    _ -> 0+          go (i + 1) (if acc > maxBound - claim then maxBound else acc + claim)+++{- | Decode a body written with body compression into an uncompressed+body plus buffer descriptors that point into it (each buffer padded to+8 bytes). 'decodeRecordBatch' does this itself, per buffer and without+the concatenation; this is for callers that need the uncompressed+message body.+-}+decompressBody :: ReadOptions -> BodyCompressionCodec -> VS.Vector Buffer -> ByteString -> Either String (VS.Vector Buffer, ByteString)+decompressBody opts codec bufs body = do+  payloads <- decompressBuffers opts codec bufs body+  let padded b = (BS.length b + 7) `quot` 8 * 8+      offsets = V.prescanl' (+) 0 (V.map padded payloads)+      !total = V.sum (V.map padded payloads)+      newBufs = VS.generate (V.length payloads) $ \i ->+        Buffer (fromIntegral (V.unsafeIndex offsets i)) (fromIntegral (BS.length (V.unsafeIndex payloads i)))+      out = I.createAligned total $ \dst -> do+        let fill !i+              | i >= V.length payloads = pure ()+              | otherwise = do+                  let !p = V.unsafeIndex payloads i+                      !off = V.unsafeIndex offsets i+                      !pl = padded p+                  I.withBytesPtr p $ \src -> copyBytes (dst `plusPtr` off) src (BS.length p)+                  fillBytes (dst `plusPtr` (off + BS.length p)) 0 (pl - BS.length p)+                  fill (i + 1)+        fill 0+  Right (newBufs, out)+++-- | Slice @len@ bytes at @off@ out of @bs@ (both wire values).+sliceWire :: String -> Int64 -> Int64 -> ByteString -> Either String ByteString+sliceWire what off len bs+  | off < 0 || len < 0 = Left ("Arrow.Read.Columns: " ++ what ++ ": negative offset or length")+  | off > n || len > n - off =+      Left+        ( "Arrow.Read.Columns: "+            ++ what+            ++ " ["+            ++ show off+            ++ ", +"+            ++ show len+            ++ ") runs past the "+            ++ show n+            ++ "-byte input"+        )+  | otherwise = Right $! BSU.unsafeTake (fromIntegral len) (BSU.unsafeDrop (fromIntegral off) bs)+  where+    n = fromIntegral (BS.length bs) :: Int64+++{- | Decompress one buffer envelope per the @BUFFER@ method.++The 8-byte uncompressed length is a wire claim. It must be @-1@+(stored uncompressed) or a size the payload could actually expand to+(see 'maxCompressionRatio'), and the decompressed bytes must match it+exactly; a lying header can never size an allocation.+-}+decompressBufferEnvelope :: BodyCompressionCodec -> ByteString -> Either String ByteString+decompressBufferEnvelope codec env+  | BS.null env = Right env+  | BS.length env < 8 = Left "Arrow.Read.Columns: buffer envelope shorter than 8 bytes"+  | rawLen == -1 = Right payload+  | rawLen < 0 = Left ("Arrow.Read.Columns: negative uncompressed buffer length " ++ show rawLen)+  | rawLen > fromIntegral (BS.length payload) * maxCompressionRatio =+      Left+        ( "Arrow.Read.Columns: buffer claims "+            ++ show rawLen+            ++ " uncompressed bytes from a "+            ++ show (BS.length payload)+            ++ "-byte payload"+        )+  | otherwise = do+      out <- decompressBuffer codec (fromIntegral rawLen) payload+      if BS.length out /= fromIntegral rawLen+        then+          Left+            ( "Arrow.Read.Columns: buffer decompressed to "+                ++ show (BS.length out)+                ++ " bytes, header says "+                ++ show rawLen+            )+        else Right out+  where+    -- Lazy on purpose: only forced once the guards above have checked+    -- that the envelope holds the 8-byte length.+    rawLen = leInt64 env+    payload = BSU.unsafeDrop 8 env+++{- | No LZ4 frame or ZSTD frame expands by more than this factor (a+ZSTD RLE block turns 4 bytes into at most 128 KiB; LZ4 tops out near+255x), so a larger claimed size is a lie, not data.+-}+maxCompressionRatio :: Int64+maxCompressionRatio = 32768+++-- | Little-endian 'Int64' in the first 8 bytes (the caller checked the length).+leInt64 :: ByteString -> Int64+leInt64 bs = go 7 0+  where+    go :: Int -> Int64 -> Int64+    go !k !acc+      | k < 0 = acc+      | otherwise = go (k - 1) ((acc `shiftL` 8) .|. fromIntegral (BSU.unsafeIndex bs k))+++{- | Decompress a payload that must expand to exactly @rawLen@ bytes+('decompressBufferEnvelope' has already bounded @rawLen@). Neither+backend may allocate more than @rawLen + 1@ output bytes.+-}+decompressBuffer :: BodyCompressionCodec -> Int -> ByteString -> Either String ByteString+#if defined(HAVE_ZSTD) || defined(HAVE_LZ4)+decompressBuffer codec rawLen comp = case codec of+#else+decompressBuffer codec _ _ = case codec of+#endif+#ifdef HAVE_ZSTD+  BodyZstd ->+    -- Zstd.decompress sizes its output from the frame header, so the+    -- header must agree with the (bounded) envelope length first.+    case Zstd.decompressedSize comp of+      Just n | n == rawLen ->+        case Zstd.decompress comp of+          Zstd.Decompress out -> Right out+          Zstd.Skip -> Left "Arrow.Read.Columns: ZSTD decompress: skipped frame"+          Zstd.Error msg -> Left ("Arrow.Read.Columns: ZSTD decompress: " ++ msg)+      other ->+        Left+          ( "Arrow.Read.Columns: ZSTD frame content size "+              ++ show other+              ++ " does not match the buffer's uncompressed length "+              ++ show rawLen+          )+#else+  BodyZstd -> Left "Arrow.Read.Columns: ZSTD body compression requires building wireform-arrow with -fzstd"+#endif+#ifdef HAVE_LZ4+  LZ4Frame ->+    -- lz4-hs's decompress is lazy and throws on malformed input; take at+    -- most one byte past the expected size (so an overlong frame is+    -- detected without being fully inflated) and catch the exception to+    -- keep the Either-shaped contract.+    case unsafeDupablePerformIO $+      Exc.try @Exc.SomeException+        (Exc.evaluate (BL.toStrict (BL.take (fromIntegral rawLen + 1) (Lz4.decompress (BL.fromStrict comp))))) of+      Right out -> Right out+      Left e -> Left ("Arrow.Read.Columns: LZ4_FRAME decompress: " ++ show e)+#else+  LZ4Frame -> Left "Arrow.Read.Columns: LZ4 body compression requires building wireform-arrow with -flz4"+#endif
+ src/Arrow/Record.hs view
@@ -0,0 +1,1523 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}++{- | @hasql@-shaped encoder / decoder combinators for Arrow's+columnar data model.++Five complementary abstractions:++* 'Encoder' @a@: column-level encoder, 'Contravariant'.+  Primitives 'int32E', 'utf8E', 'boolE', ... reshape with+  @'contramap' :: (a -> b) -> Encoder b -> Encoder a@ and+  @'nullable' :: Encoder a -> Encoder (Maybe a)@.++* 'Decoder' @a@: column-level decoder, 'Functor'. Mirror set+  of primitives ('int32D', ...) with @'nullableD'@.++* 'RowEncoder' @r@: record-level encoder. Combine+  'fieldE' calls via 'Semigroup' @<>@.++* 'RowDecoder' @r@: 'Applicative' row decoder. Build with+  @<$>@ + @<*>@ + 'columnD'.++* 'Table' @r@: pairs the two for round-trip use.++== Cost model++Encoding allocates one column buffer per field, sized for the row+count (fixed-width values and booleans are stored straight into their+slot; strings and bytes go through "Arrow.Column.Builder"), then+traverses the input records once, appending each record to every+column. Field selectors, 'contramap' and the 'nullable' unwrapping+are fused into that loop, so no intermediate vector and no per-row+heap object is allocated. Decoding binds each column once (type+check, dictionary handling) to a row reader, composes the readers+through the 'Applicative', and runs a single 'V.generate' over the+rows: the only per-row allocation is the record itself and its boxed+fields ('Text' fields are slices of one copy of their column's string+bytes, made when the column is bound; 'ByteString' fields alias the+input column).++== Example++@+data Trade = Trade { sym :: Text, qty :: Int32, note :: Maybe Text }++tradeTable :: 'Table' Trade+tradeTable = 'table' enc dec+  where+    enc = 'fieldE' "sym"  sym  'utf8E'+       <> 'fieldE' "qty"  qty  'int32E'+       <> 'fieldE' "note" note ('nullable' 'utf8E')+    dec = Trade+        \<$\> 'columnD' "sym"  'utf8D'+        \<*\> 'columnD' "qty"  'int32D'+        \<*\> 'columnD' "note" ('nullableD' 'utf8D')++encoded = 'encodeTable' tradeTable tradesVec+@+-}+module Arrow.Record (+  -- * Column-level encoder+  Encoder,+  encoderType,+  encoderNullable,+  runEncoder,+  contramapE,+  nullable,++  -- ** Primitive encoders+  int8E,+  int16E,+  int32E,+  int64E,+  word8E,+  word16E,+  word32E,+  word64E,+  floatE,+  doubleE,+  boolE,+  utf8E,+  binaryE,+  date32E,+  timestampE,++  -- * Column-level decoder+  Decoder,+  decoderType,+  runDecoder,+  nullableD,++  -- ** Primitive decoders+  int8D,+  int16D,+  int32D,+  int64D,+  word8D,+  word16D,+  word32D,+  word64D,+  floatD,+  doubleD,+  boolD,+  utf8D,+  binaryD,+  date32D,+  timestampD,++  -- * Row-level encoder+  RowEncoder,+  rowEncoderFields,+  runRowEncoder,+  fieldE,+  structE,+  structEMaybe,++  -- * Row-level decoder+  RowDecoder,+  runRowDecoder,+  rowDecoderRequiredColumns,+  columnD,+  columnDWithDefault,+  structD,+  structDMaybe,++  -- * Column-name strategies+  NameStrategy (..),+  applyNameStrategy,++  -- * Table (round-trip pair)+  Table (..),+  table,+  tableSchema,+  tableRequiredColumns,+  encodeTable,+  decodeTable,++  -- * Subset / projection+  subsetTable,+  projectTable,+) where++import Arrow.Column (+  BoolArray (..),+  BytesArray (..),+  ColumnArray,+  PrimArray (..),+  PrimType (..),+  Utf8Array (..),+  anyBytesAt,+  anyTextAt,+  asBinary,+  asBool,+  asLargeBinary,+  asLargeUtf8,+  asPrim,+  asUtf8,+  bitAt,+  boolArrayAt,+  columnLength,+  columnTag,+  expandDictionary,+  fillerColumn,+  isValidAt,+  mkBitmap,+  mkBool,+  mkPrim,+  mkStruct,+  mkValidity,+  nullCount,+  primColumn,+  unsafeBytesAt,+  unsafePrimAt,+  validity,+  validityGenerate,+  pattern ColBinaryView,+  pattern ColDictionary,+  pattern ColFixedSizeBinary,+  pattern ColPrim,+  pattern ColStruct,+  pattern ColUtf8View,+ )+import Arrow.Column qualified as AC+import Arrow.Column.Buffer (mallocAligned, unsafeIsValidAt)+import Arrow.Column.Internal qualified as I+import Arrow.Types (+  ArrowType (..),+  DateUnit (..),+  Endianness (..),+  Field (..),+  Precision (..),+  Schema (..),+  TimeUnit (..),+ )+import Control.Monad (when)+import Control.Monad.ST (ST, runST, stToIO)+import Control.Monad.ST.Unsafe (unsafeIOToST)+import Data.Bits (complement, unsafeShiftL, unsafeShiftR, (.&.), (.|.))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Internal qualified as BSI+import Data.ByteString.Unsafe qualified as BSU+import Data.Functor.Contravariant (Contravariant (..))+import Data.Int (Int16, Int32, Int64, Int8)+import Data.List (findIndex)+import Data.Maybe (fromMaybe, isJust)+import Data.Primitive.MutVar (newMutVar, readMutVar, writeMutVar)+import Data.Primitive.PrimArray (MutablePrimArray, newPrimArray, readPrimArray, writePrimArray)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Text.Array qualified as TA+import Data.Text.Internal qualified as TI+import Data.Text.Foreign qualified as TF+import Data.Vector qualified as V+import Data.Vector.Storable qualified as VS+import Data.Vector.Storable.Mutable qualified as VSM+import Data.Word (Word16, Word32, Word64, Word8)+import Foreign.ForeignPtr (ForeignPtr)+import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)+import Foreign.Marshal.Utils (copyBytes, fillBytes)+import Foreign.Ptr (Ptr, castPtr, minusPtr, nullPtr, plusPtr)+import Foreign.Storable (Storable, peekByteOff, pokeByteOff, pokeElemOff, sizeOf)+import GHC.Exts (Int (I#), Int#)+import GHC.ForeignPtr (unsafeWithForeignPtr)+import System.IO.Unsafe (unsafeDupablePerformIO)+++-- ============================================================+-- Encoder+-- ============================================================++{- | Serialises Haskell values as an Arrow column.++Pairs the Arrow type (used to populate the schema) with a column+sink: given the row count, it allocates the column's buffers and+returns an append action (indexed by row), a null-append action and a+freeze. 'contramap' composes the projection into the append action,+so a field selector runs inside the encoding loop and no intermediate+vector is built.+-}+data Encoder a = Encoder+  { encoderType :: ArrowType+  , encoderNullable :: Bool+  , encoderSink :: forall s. Int -> ST s (ColSink s a)+  }+++{- | A column under construction, addressed by row index (rows are+appended in order @0, 1, ..@ exactly once each): append a value,+append a null, freeze.+-}+data ColSink s a = ColSink (Int# -> a -> ST s ()) (Int# -> ST s ()) (ST s ColumnArray)+++-- | Encode a column with one row per input value.+runEncoder :: Encoder a -> V.Vector a -> ColumnArray+runEncoder e v = runST $ do+  ColSink push _ freeze <- encoderSink e (V.length v)+  V.imapM_ (\(I# i) x -> push i x) v+  freeze+{-# INLINE runEncoder #-}+++instance Contravariant Encoder where+  contramap f (Encoder ty nu mk) =+    Encoder+      { encoderType = ty+      , encoderNullable = nu+      , encoderSink = \n -> do+          ColSink push pushNull freeze <- mk n+          pure (ColSink (\i x -> push i (f x)) pushNull freeze)+      }+  {-# INLINE contramap #-}+++{- | Alias for 'contramap' that reads more naturally at call+sites.+-}+contramapE :: (a -> b) -> Encoder b -> Encoder a+contramapE = contramap+{-# INLINE contramapE #-}+++{- | Lift an encoder to build a nullable Arrow column: 'Nothing'+rows are null. Arrow has no nested-null representation, so running+an encoder wrapped in 'nullable' twice is a runtime error.+-}+nullable :: Encoder a -> Encoder (Maybe a)+nullable e =+  Encoder+    { encoderType = encoderType e+    , encoderNullable = True+    , encoderSink = \n ->+        if encoderNullable e+          then+            errorWithoutStackTrace+              "Arrow.Record.nullable: Arrow has no nested-null \+              \representation; don't wrap 'nullable' twice"+          else do+            ColSink push pushNull freeze <- encoderSink e n+            let pushMaybe i = \case+                  Nothing -> pushNull i+                  Just x -> push i x+            pure (ColSink pushMaybe pushNull freeze)+    }+{-# INLINE nullable #-}+++{- | Fixed-width encoder. The row count is known up front, so values+are stored straight into their slot of one pinned buffer (a single+store per row); nulls are recorded in a side mask and only turned+into a validity bitmap if one occurred.++Like every encoder here, the append actions write through the raw+addresses of their pinned buffers and capture nothing else: a row+encoder's per-row action holds every field's append state, and each+captured value costs a load and a stack slot per row. The buffers stay+alive because the freeze action, which runs after the last append,+holds them.+-}+primE :: Storable a => ArrowType -> PrimType a -> Encoder a+primE ty t =+  Encoder+    { encoderType = ty+    , encoderNullable = False+    , encoderSink = \n -> do+        vals <- VSM.unsafeNew n+        nulls <- newNullMask n+        let !p = mvectorPtr vals+            pushNull i = do+              zeroSlot p (I# i)+              markNull nulls (I# i)+            freeze = do+              v <- VS.unsafeFreeze vals+              valid <- freezeNullMask nulls n+              pure $ case valid of+                Nothing -> primColumn t v+                Just _ -> encoded "primE" (mkPrim t valid v)+        pure (ColSink (\i x -> unsafeIOToST (pokeElemOff p (I# i) x)) pushNull freeze)+    }+{-# INLINE primE #-}+++-- | The address of a pinned storable vector's first element.+mvectorPtr :: VSM.MVector s a -> Ptr a+mvectorPtr = unsafeForeignPtrToPtr . fst . VSM.unsafeToForeignPtr0+{-# INLINE mvectorPtr #-}+++-- | Zero a null row's value slot (deterministic bytes under nulls).+zeroSlot :: forall s a. Storable a => Ptr a -> Int -> ST s ()+zeroSlot p i = unsafeIOToST $ do+  let !w = sizeOf (undefined :: a)+  fillBytes (castPtr p `plusPtr` (i * w)) 0 w+++-- | Unwrap a column built from rows that satisfy its invariants by construction.+encoded :: String -> Either String ColumnArray -> ColumnArray+encoded who = \case+  Right c -> c+  Left e -> errorWithoutStackTrace ("Arrow.Record." ++ who ++ ": " ++ e)+++{- | Null rows seen so far: a zeroed bit per row (set = null), its+address, and the null count.+-}+data NullMask s = NullMask (Ptr Word8) (ForeignPtr Word8) (MutablePrimArray s Int)+++newNullMask :: Int -> ST s (NullMask s)+newNullMask n = do+  bits <- newBits n+  count <- newPrimArray 1+  writePrimArray count 0 0+  pure (NullMask (unsafeForeignPtrToPtr bits) bits count)+++markNull :: NullMask s -> Int -> ST s ()+markNull (NullMask p _ count) i = do+  setBit' p i+  c <- readPrimArray count 0+  writePrimArray count 0 (c + 1)+{-# INLINE markNull #-}+++-- | The validity of @n@ rows: 'Nothing' when no row was null.+freezeNullMask :: NullMask s -> Int -> ST s (Maybe AC.Validity)+freezeNullMask (NullMask _ bits count) n = do+  c <- readPrimArray count 0+  if c == 0+    then pure Nothing+    else unsafeIOToST $ do+      let !len = bitBytes n+      valid <- BSI.mallocByteString len+      unsafeWithForeignPtr bits $ \src -> unsafeWithForeignPtr valid $ \dst ->+        let go j+              | j >= len = pure ()+              | otherwise = do+                  b <- peekByteOff src j :: IO Word8+                  pokeByteOff dst j (complement b)+                  go (j + 1)+        in go 0+      pure $ case mkBitmap (BSI.BS valid len) 0 n of+        Right bm -> mkValidity bm+        Left e -> errorWithoutStackTrace ("Arrow.Record: validity bitmap: " ++ e)+++bitBytes :: Int -> Int+bitBytes n = (n + 7) `unsafeShiftR` 3+{-# INLINE bitBytes #-}+++-- | A zeroed, pinned bitmap with room for @n@ bits.+newBits :: Int -> ST s (ForeignPtr Word8)+newBits n = unsafeIOToST $ do+  let !len = bitBytes n+  fp <- BSI.mallocByteString len+  unsafeWithForeignPtr fp $ \p -> fillBytes p 0 len+  pure fp+++setBit' :: Ptr Word8 -> Int -> ST s ()+setBit' p i = unsafeIOToST $ do+  let !byte = i `unsafeShiftR` 3+  b <- peekByteOff p byte :: IO Word8+  pokeByteOff p byte (b .|. (1 `unsafeShiftL` (i .&. 7)))+{-# INLINE setBit' #-}+++-- ============================================================+-- Primitive encoders+-- ============================================================++int8E :: Encoder Int8+int8E = primE (AInt 8 True) PInt8+{-# INLINE int8E #-}+++int16E :: Encoder Int16+int16E = primE (AInt 16 True) PInt16+{-# INLINE int16E #-}+++int32E :: Encoder Int32+int32E = primE (AInt 32 True) PInt32+{-# INLINE int32E #-}+++int64E :: Encoder Int64+int64E = primE (AInt 64 True) PInt64+{-# INLINE int64E #-}+++word8E :: Encoder Word8+word8E = primE (AInt 8 False) PUInt8+{-# INLINE word8E #-}+++word16E :: Encoder Word16+word16E = primE (AInt 16 False) PUInt16+{-# INLINE word16E #-}+++word32E :: Encoder Word32+word32E = primE (AInt 32 False) PUInt32+{-# INLINE word32E #-}+++word64E :: Encoder Word64+word64E = primE (AInt 64 False) PUInt64+{-# INLINE word64E #-}+++floatE :: Encoder Float+floatE = primE (AFloatingPoint Single) PFloat+{-# INLINE floatE #-}+++doubleE :: Encoder Double+doubleE = primE (AFloatingPoint DoublePrecision) PDouble+{-# INLINE doubleE #-}+++{- | One bit per row, set in place in a pinned bitmap sized for the+row count.+-}+boolE :: Encoder Bool+boolE =+  Encoder+    { encoderType = ABool+    , encoderNullable = False+    , encoderSink = \n -> do+        bits <- newBits n+        nulls <- newNullMask n+        let !bp = unsafeForeignPtrToPtr bits+            push i b = if b then setBit' bp (I# i) else pure ()+            freeze = do+              valid <- freezeNullMask nulls n+              pure $ case mkBitmap (BSI.BS bits (bitBytes n)) 0 n of+                Right bm -> encoded "boolE" (mkBool valid bm)+                Left e -> errorWithoutStackTrace ("Arrow.Record.boolE: " ++ e)+        pure (ColSink push (\i -> markNull nulls (I# i)) freeze)+    }+{-# INLINE boolE #-}+++utf8E :: Encoder Text+utf8E = varE AUtf8 I.ColUtf8 TF.lengthWord8 TF.unsafeCopyToPtr+{-# INLINE utf8E #-}+++binaryE :: Encoder ByteString+binaryE = varE ABinary I.ColBinary BS.length $ \bs dst ->+  BSU.unsafeUseAsCStringLen bs $ \(src, l) -> copyBytes dst (castPtr src) l+{-# INLINE binaryE #-}+++{- | Var-length encoder (32-bit offsets). The row count is known, so+each row's end offset is stored straight into its slot of one buffer+sized for the rows; the bytes go into a growable pinned buffer whose+used size, capacity and address live in one small mutable array. An+append is three loads, the copy and two stores, and the append action+captures only that array, the offsets' address and the null mask (a+column builder keeps several mutable cells per buffer). Nulls are+recorded as in 'primE'.+-}+varE+  :: ArrowType+  -> (Maybe AC.Validity -> VS.Vector Int32 -> ByteString -> ColumnArray)+  -> (a -> Int)+  -> (a -> Ptr Word8 -> IO ())+  -> Encoder a+varE ty mk lenOf copyTo =+  Encoder+    { encoderType = ty+    , encoderNullable = False+    , encoderSink = \n -> do+        offs <- VSM.unsafeNew (n + 1)+        VSM.unsafeWrite offs 0 0+        let !cap0 = max 64 (n * 8)+        dat0 <- unsafeIOToST (mallocAligned cap0)+        ref <- newMutVar dat0+        st <- newPrimArray 3+        writePrimArray st varUsed 0+        writePrimArray st varCap cap0+        writePrimArray st varAddr (addrOf dat0)+        nulls <- newNullMask n+        let !op = mvectorPtr offs+            grow used need = do+              cap <- readPrimArray st varCap+              let !cap' = max need (cap * 2)+              old <- readMutVar ref+              new <- unsafeIOToST $ do+                fp <- mallocAligned cap'+                unsafeWithForeignPtr fp $ \d -> unsafeWithForeignPtr old $ \s -> copyBytes d s used+                pure fp+              writeMutVar ref new+              writePrimArray st varCap cap'+              writePrimArray st varAddr (addrOf new)+            push i x = do+              used <- readPrimArray st varUsed+              cap <- readPrimArray st varCap+              let !end = used + lenOf x+              when (end > maxOffset) $+                errorWithoutStackTrace "Arrow.Record: var-length column data exceeds 2^31 - 1 bytes"+              when (end > cap) (grow used end)+              base <- readPrimArray st varAddr+              unsafeIOToST (copyTo x (nullPtr `plusPtr` (base + used)))+              unsafeIOToST (pokeElemOff op (I# i + 1) (fromIntegral end))+              writePrimArray st varUsed end+            pushNull i = do+              used <- readPrimArray st varUsed+              unsafeIOToST (pokeElemOff op (I# i + 1) (fromIntegral used))+              markNull nulls (I# i)+            freeze = do+              used <- readPrimArray st varUsed+              fp <- readMutVar ref+              o <- VS.unsafeFreeze offs+              valid <- freezeNullMask nulls n+              pure (mk valid o (BSI.BS fp used))+        pure (ColSink push pushNull freeze)+    }+  where+    addrOf fp = unsafeForeignPtrToPtr fp `minusPtr` nullPtr+    maxOffset = fromIntegral (maxBound :: Int32)+{-# INLINE varE #-}+++-- | Slots of a 'varE' state array.+varUsed, varCap, varAddr :: Int+varUsed = 0+varCap = 1+varAddr = 2+++-- | Days since Unix epoch (INT32). Arrow logical @Date(DateDay)@.+date32E :: Encoder Int32+date32E = primE (ADate DateDay) PDate32+{-# INLINE date32E #-}+++{- | Microseconds since Unix epoch (INT64, no timezone). Arrow+logical @Timestamp(Microsecond, None)@.+-}+timestampE :: Encoder Int64+timestampE = primE (ATimestamp Microsecond Nothing) PTimestamp+{-# INLINE timestampE #-}+++-- ============================================================+-- Decoder+-- ============================================================++{- | A bound column (or row) reader: row index to value. The index is+unboxed so that calling a reader through an unknown closure passes it+in a register instead of allocating a box per row.+-}+type Reader a = Int# -> a+++{- | Reads Haskell values out of an Arrow column.++A decoder binds a column once (type check, dictionary handling) and+yields a reader that indexes row @i@ straight out of the column's+buffers. 'nullableD' switches to the second binder, which yields a+@Maybe a@ reader and accepts columns with nulls.++Dictionary-encoded columns of the decoder's value type are accepted:+nullable decoders read the values through the keys without expanding+the dictionary; required decoders expand it first ('AC.expandDictionary'),+which keeps the plain-column path a single, fully specialisable branch.+-}+data Decoder a = Decoder+  { decoderType :: ArrowType+  , decoderBind :: ColumnArray -> Either String (Reader a)+  , decoderBindMaybe :: ColumnArray -> Either String (Reader (Maybe a))+  }+++instance Functor Decoder where+  fmap f (Decoder ty b bm) =+    Decoder+      { decoderType = ty+      , decoderBind = \c -> case b c of+          Left e -> Left e+          Right g -> Right (\i -> f (g i))+      , decoderBindMaybe = \c -> case bm c of+          Left e -> Left e+          Right g -> Right (\i -> fmap f (g i))+      }+  {-# INLINE fmap #-}+++-- | Decode a column without nulls into a vector of values.+runDecoder :: Decoder a -> ColumnArray -> Either String (V.Vector a)+runDecoder d col = case decoderBind d col of+  Left e -> Left e+  Right g -> Right (V.generate (columnLength col) (\(I# i) -> g i))+++{- | Lift a 'Decoder' to read nullable columns: null rows become+'Nothing'. Arrow has no nested-null representation, so decoding+through a second 'nullableD' wrap returns 'Left'.+-}+nullableD :: Decoder a -> Decoder (Maybe a)+nullableD d =+  Decoder+    { decoderType = decoderType d+    , decoderBind = decoderBindMaybe d+    , decoderBindMaybe = \_ ->+        Left+          "Arrow.Record.nullableD: Arrow has no nested-null \+          \representation; don't wrap 'nullableD' twice"+    }+{-# INLINE nullableD #-}+++{- | Build a decoder from binders for the plain (non-dictionary)+column shapes, adding dictionary support. Required decoders expand a+dictionary column first (so a key selecting a null value is a null+row, rejected like any other); nullable decoders read the values+through the keys without materialising anything.+-}+dictD+  :: ArrowType+  -> (ColumnArray -> Either String (Reader a))+  -> (ColumnArray -> Either String (Reader (Maybe a)))+  -> Decoder a+dictD ty req opt = Decoder ty bindReq bindOpt+  where+    bindReq col = case undictionary col of+      Left e -> Left e+      Right plain -> req plain+    bindOpt col = case col of+      ColDictionary _ keys vals -> do+        checkResolved keys vals+        gv <- opt vals+        g <- throughKeys keys gv+        Right $ case validity keys of+          Nothing -> g+          kv -> \i -> if isValidAt kv (I# i) then g i else Nothing+      _ -> opt col+{-# INLINE dictD #-}+++{- | A dictionary column expanded to its value type; other columns+unchanged. Kept out of line so that a required decoder's binder calls+its plain-column reader once, letting GHC specialise the reader into+the row function.+-}+{-# NOINLINE undictionary #-}+undictionary :: ColumnArray -> Either String ColumnArray+undictionary col = case col of+  ColDictionary _ keys vals -> do+    checkResolved keys vals+    expandDictionary col+  _ -> Right col+++{- | An unresolved dictionary column (values still the empty+placeholder) can only be read when every key is null.+-}+checkResolved :: ColumnArray -> ColumnArray -> Either String ()+checkResolved keys vals+  | columnLength vals == 0 && nullCount keys < columnLength keys =+      Left "Arrow.Record: dictionary column has unresolved (empty) values"+  | otherwise = Right ()+++{- | Compose a value-row reader with a dictionary's keys. Keys are an+integer column at wire width; every valid key is in range of the+values (a 'ColDictionary' invariant, checked by 'checkResolved' for+the placeholder case). The result reads key slots without consulting+key validity; callers mask nulls.+-}+throughKeys :: forall b. ColumnArray -> Reader b -> Either String (Reader b)+throughKeys keys gv = case keys of+  ColPrim t _ ks -> case t of+    PInt8 -> Right (via ks)+    PInt16 -> Right (via ks)+    PInt32 -> Right (via ks)+    PInt64 -> Right (via ks)+    PUInt8 -> Right (via ks)+    PUInt16 -> Right (via ks)+    PUInt32 -> Right (via ks)+    PUInt64 -> Right (via ks)+    _ -> Left ("Arrow.Record: dictionary keys must be integers, got " ++ columnTag keys)+  _ -> Left ("Arrow.Record: dictionary keys must be integers, got " ++ columnTag keys)+  where+    via :: (Storable k, Integral k) => VS.Vector k -> Reader b+    via ks i = case fromIntegral (VS.unsafeIndex ks (I# i)) of I# k -> gv k+    {-# INLINE via #-}+++expectErr :: String -> ColumnArray -> String+expectErr want got =+  "Arrow.Record: expected " ++ want ++ ", got " ++ columnTag got+++nullsErr :: ColumnArray -> String+nullsErr col =+  "Arrow.Record: "+    ++ columnTag col+    ++ " column has "+    ++ show (nullCount col)+    ++ " null rows; decode it with nullableD"+++{- | Fixed-width decoder: the column must carry the given tag+(dictionary columns of that value type are accepted, see 'dictD').+-}+primD :: Storable a => ArrowType -> PrimType a -> String -> Decoder a+primD ty t want = dictD ty req opt+  where+    req col = case asPrim t col of+      Just (PrimArray Nothing xs) -> Right (\i -> VS.unsafeIndex xs (I# i))+      Just _ -> Left (nullsErr col)+      Nothing -> Left (expectErr want col)+    opt col = case asPrim t col of+      Just (PrimArray Nothing xs) -> Right (\i -> let !x = VS.unsafeIndex xs (I# i) in Just x)+      Just arr -> Right $ \i -> case unsafePrimAt arr (I# i) of+        Nothing -> Nothing+        Just x -> x `seq` Just x+      Nothing -> Left (expectErr want col)+{-# INLINE primD #-}+++-- ============================================================+-- Primitive decoders+-- ============================================================++int8D :: Decoder Int8+int8D = primD (AInt 8 True) PInt8 "ColInt8"+{-# INLINE int8D #-}+++int16D :: Decoder Int16+int16D = primD (AInt 16 True) PInt16 "ColInt16"+{-# INLINE int16D #-}+++int32D :: Decoder Int32+int32D = primD (AInt 32 True) PInt32 "ColInt32"+{-# INLINE int32D #-}+++int64D :: Decoder Int64+int64D = primD (AInt 64 True) PInt64 "ColInt64"+{-# INLINE int64D #-}+++word8D :: Decoder Word8+word8D = primD (AInt 8 False) PUInt8 "ColUInt8"+{-# INLINE word8D #-}+++word16D :: Decoder Word16+word16D = primD (AInt 16 False) PUInt16 "ColUInt16"+{-# INLINE word16D #-}+++word32D :: Decoder Word32+word32D = primD (AInt 32 False) PUInt32 "ColUInt32"+{-# INLINE word32D #-}+++word64D :: Decoder Word64+word64D = primD (AInt 64 False) PUInt64 "ColUInt64"+{-# INLINE word64D #-}+++floatD :: Decoder Float+floatD = primD (AFloatingPoint Single) PFloat "ColFloat"+{-# INLINE floatD #-}+++doubleD :: Decoder Double+doubleD = primD (AFloatingPoint DoublePrecision) PDouble "ColDouble"+{-# INLINE doubleD #-}+++boolD :: Decoder Bool+boolD = dictD ABool req opt+  where+    req col = case asBool col of+      Just (BoolArray Nothing bits) -> Right (\i -> bitAt bits (I# i))+      Just _ -> Left (nullsErr col)+      Nothing -> Left (expectErr "ColBool" col)+    opt col = case asBool col of+      Just (BoolArray Nothing bits) -> Right (\i -> let !b = bitAt bits (I# i) in Just b)+      Just arr -> Right (\i -> boolArrayAt arr (I# i))+      Nothing -> Left (expectErr "ColBool" col)+{-# INLINE boolD #-}+++{- | Text from a utf8, large utf8 or utf8 view column (the column was+validated as UTF-8 when it was built or decoded, so nothing is+re-validated).++Binding a utf8 or large utf8 column copies the string bytes its rows+reference once into one text array, and every row's 'Text' is a slice+of it: no per-row allocation beyond the 'Text' itself, and nothing+keeps the column's (or the IPC input's) buffers alive. Retention rule:+a 'Text' you keep keeps that column's copied string bytes alive; use+'T.copy' to detach it. Utf8 view rows are each copied into a fresh+'Text'.+-}+utf8D :: Decoder Text+utf8D = dictD AUtf8 req opt+  where+    req col+      | Just arr <- asUtf8 col = noNulls col (let !rows = textRows arr in \i -> textRowValue rows (I# i))+      | Just arr <- asLargeUtf8 col = noNulls col (let !rows = textRows arr in \i -> textRowValue rows (I# i))+      | ColUtf8View {} <- col = noNulls col (\i -> orEmpty T.empty (anyTextAt col (I# i)))+      | otherwise = Left (expectErr "ColUtf8" col)+    opt col+      | Just arr <- asUtf8 col = Right (let !rows = textRows arr in \i -> textRowAt rows (I# i))+      | Just arr <- asLargeUtf8 col = Right (let !rows = textRows arr in \i -> textRowAt rows (I# i))+      | ColUtf8View {} <- col = Right (\i -> forceJust (anyTextAt col (I# i)))+      | otherwise = Left (expectErr "ColUtf8" col)+{-# INLINE utf8D #-}+++{- | A utf8 column's rows over one text array holding a copy of the+bytes they reference: validity, offsets, the array, and the data offset+of the array's first byte.+-}+data TextRows o = TextRows !(Maybe AC.Validity) !(VS.Vector o) !TA.Array {-# UNPACK #-} !Int+++-- | Copy the referenced string bytes of a utf8 column once.+textRows :: (Storable o, Integral o) => Utf8Array o -> TextRows o+textRows (Utf8Array (BytesArray v o d))+  | len <= 0 = TextRows v o TA.empty base+  | otherwise = TextRows v o copied base+  where+    !n = VS.length o - 1+    !base = if n < 1 then 0 else fromIntegral (VS.unsafeIndex o 0)+    !len = if n < 1 then 0 else fromIntegral (VS.unsafeIndex o n) - base+    copied = unsafeDupablePerformIO $ BSU.unsafeUseAsCString d $ \p -> stToIO $ do+      ma <- TA.new len+      TA.copyFromPointer ma 0 (castPtr p `plusPtr` base) len+      TA.unsafeFreeze ma+{-# INLINE textRows #-}+++-- | Row @i@ (in range) as a slice, ignoring validity.+textRowValue :: (Storable o, Integral o) => TextRows o -> Int -> Text+textRowValue (TextRows _ o arr base) i =+  let !s = fromIntegral (VS.unsafeIndex o i)+      !len = fromIntegral (VS.unsafeIndex o (i + 1)) - s+  in if len == 0 then T.empty else TI.Text arr (s - base) len+{-# INLINE textRowValue #-}+++-- | Row @i@ (in range) as a slice; 'Nothing' when null.+textRowAt :: (Storable o, Integral o) => TextRows o -> Int -> Maybe Text+textRowAt rows@(TextRows v _ _ _) i+  | unsafeIsValidAt v i = let !t = textRowValue rows i in Just t+  | otherwise = Nothing+{-# INLINE textRowAt #-}+++-- | A required decoder's row reader, provided the column has no nulls.+noNulls :: ColumnArray -> Reader a -> Either String (Reader a)+noNulls col g+  | nullCount col > 0 = Left (nullsErr col)+  | otherwise = Right g+{-# INLINE noNulls #-}+++{- | The value of a row known to be valid (the accessors only return+'Nothing' for null rows, which 'noNulls' excluded).+-}+orEmpty :: a -> Maybe a -> a+orEmpty def = \case+  Just x -> x+  Nothing -> def+{-# INLINE orEmpty #-}+++-- | Evaluate the payload of a 'Just' so the row holds no thunk.+forceJust :: Maybe a -> Maybe a+forceJust = \case+  Nothing -> Nothing+  Just x -> x `seq` Just x+{-# INLINE forceJust #-}+++{- | Bytes from a binary, large binary, binary view or fixed-size+binary column (utf8 columns are accepted too). Rows are zero-copy+slices that keep the column's buffer alive.+-}+binaryD :: Decoder ByteString+binaryD = dictD ABinary req opt+  where+    req col+      | Just arr <- asBinary col = noNulls col (\i -> orEmpty BS.empty (unsafeBytesAt arr (I# i)))+      | Just arr <- asLargeBinary col = noNulls col (\i -> orEmpty BS.empty (unsafeBytesAt arr (I# i)))+      | isViewOrFixed col = noNulls col (\i -> orEmpty BS.empty (anyBytesAt col (I# i)))+      | otherwise = Left (expectErr "ColBinary" col)+    opt col+      | Just arr <- asBinary col = Right (\i -> forceJust (unsafeBytesAt arr (I# i)))+      | Just arr <- asLargeBinary col = Right (\i -> forceJust (unsafeBytesAt arr (I# i)))+      | isViewOrFixed col = Right (\i -> forceJust (anyBytesAt col (I# i)))+      | otherwise = Left (expectErr "ColBinary" col)+{-# INLINE binaryD #-}+++isViewOrFixed :: ColumnArray -> Bool+isViewOrFixed = \case+  ColBinaryView {} -> True+  ColUtf8View {} -> True+  ColFixedSizeBinary {} -> True+  _ -> False+++date32D :: Decoder Int32+date32D = primD (ADate DateDay) PDate32 "ColDate32"+{-# INLINE date32D #-}+++timestampD :: Decoder Int64+timestampD = primD (ATimestamp Microsecond Nothing) PTimestamp "ColTimestamp"+{-# INLINE timestampD #-}+++-- ============================================================+-- RowEncoder+-- ============================================================++{- | A record-level encoder: produces one 'ColumnArray' per field,+plus the matching 'Field' list, from a 'V.Vector' of records.++'RowEncoder' is 'Contravariant' and a 'Semigroup' / 'Monoid'.+Combine 'fieldE' calls with @<>@:++@+enc = 'fieldE' "sym" sym utf8E <> 'fieldE' "qty" qty int32E+@+-}+data RowEncoder r = RowEncoder+  { rowEncoderFields :: [Field]+  -- ^ 'Field' entries in declaration order.+  , rowEncoderSink :: forall s x. V.Vector x -> (x -> r) -> ST s (RowSink s r)+  {- ^ Allocate one builder per field for the given input rows (seen+  through a projection, for encoders that need to look at every row+  up front) and return the per-row append action and the freeze.+  -}+  }+++{- | Rows under construction: append record @i@ to every field's+builder, freeze every builder (one column per field).+-}+data RowSink s r = RowSink (Int# -> r -> ST s ()) (ST s [ColumnArray])+++instance Contravariant RowEncoder where+  contramap f (RowEncoder fields mk) = RowEncoder fields $ \v g -> do+    RowSink push freeze <- mk v (\x -> f (g x))+    pure (RowSink (\i r -> push i (f r)) freeze)+  {-# INLINE contramap #-}+++instance Semigroup (RowEncoder r) where+  RowEncoder fl ml <> RowEncoder fr mr = RowEncoder (fl ++ fr) $ \v g -> do+    RowSink pl zl <- ml v g+    RowSink pr zr <- mr v g+    pure (RowSink (\i r -> pl i r >> pr i r) ((++) <$> zl <*> zr))+  {-# INLINE (<>) #-}+++instance Monoid (RowEncoder r) where+  mempty = RowEncoder [] (\_ _ -> pure (RowSink (\_ _ -> pure ()) (pure [])))+  {-# INLINE mempty #-}+++{- | One 'ColumnArray' per field, parallel to 'rowEncoderFields'. The+rows are traversed once; each record is appended to every field's+builder.+-}+runRowEncoder :: RowEncoder r -> V.Vector r -> [ColumnArray]+runRowEncoder e v = runST $ do+  RowSink push freeze <- rowEncoderSink e v id+  V.imapM_ (\(I# i) r -> push i r) v+  freeze+++leafField :: Text -> Bool -> ArrowType -> V.Vector Field -> Field+leafField name nu ty children =+  Field+    { fieldName = name+    , fieldNullable = nu+    , fieldType = ty+    , fieldChildren = children+    , fieldDictionary = Nothing+    , fieldMetadata = V.empty+    }+++{- | Build a 'RowEncoder' for a single field: name + selector ++column encoder. The selector runs inside the encoder's builder loop.++@+fieldE "sym" tradeSym utf8E  :: RowEncoder Trade+@+-}+fieldE :: Text -> (r -> a) -> Encoder a -> RowEncoder r+fieldE name sel enc =+  RowEncoder [leafField name (encoderNullable enc) (encoderType enc) V.empty] $ \v _ -> do+    ColSink push _ freeze <- encoderSink enc (V.length v)+    pure (RowSink (\i r -> push i (sel r)) (fmap (: []) freeze))+{-# INLINE fieldE #-}+++{- | Embed a nested record as a struct column.++Lifts a 'RowEncoder' for a child record type @c@ into a+'RowEncoder' for the parent @r@ that emits the child's+column tree under one named struct field. The struct's+children are exactly the child encoder's fields (in the+order they were declared with '<>').++@+data Address = Address { city :: Text, zip :: Text }+data Customer = Customer { name :: Text, addr :: Address }++addressEnc :: 'RowEncoder' Address+addressEnc = 'fieldE' "city" city utf8E+          <> 'fieldE' "zip"  zip  utf8E++customerEnc :: 'RowEncoder' Customer+customerEnc = 'fieldE'  "name" name  utf8E+           <> 'structE' "addr" addr  addressEnc+@+-}+structE :: Text -> (r -> c) -> RowEncoder c -> RowEncoder r+structE name sel inner =+  RowEncoder [leafField name False AStruct (V.fromList (rowEncoderFields inner))] $ \v g -> do+    RowSink push freeze <- rowEncoderSink inner v (\x -> sel (g x))+    pure $+      RowSink+        (\i r -> push i (sel r))+        (do cols <- freeze; pure [buildStruct "structE" (V.length v) Nothing (childNames inner) cols])+{-# INLINE structE #-}+++{- | Like 'structE' but the parent rows are @Maybe c@: emits a+struct column with a validity bitmap. Child slots under a null+parent are arbitrary on the wire (Arrow spec, Layout.rst, "Struct+Layout"), so they repeat the first present row's value; if every+row is 'Nothing' the children are 'AC.fillerColumn' rows of the+child encoders' column shapes.++Pair with 'structDMaybe' on the read side.+-}+structEMaybe :: forall r c. Text -> (r -> Maybe c) -> RowEncoder c -> RowEncoder r+structEMaybe name sel inner =+  RowEncoder [leafField name True AStruct (V.fromList (rowEncoderFields inner))] $ \v g -> do+    let !n = V.length v+        pick x = sel (g x)+        assemble valid cols = [buildStruct "structEMaybe" n valid (childNames inner) cols]+    case V.find (isJust . pick) v >>= pick of+      Nothing -> do+        let !cols = map (fillerColumn n) (runRowEncoder inner (V.empty :: V.Vector c))+            !col = assemble (validityGenerate n (const False)) cols+        pure (RowSink (\_ _ -> pure ()) (pure col))+      Just present -> do+        let orPresent = fromMaybe present+            !valid = validityGenerate n (\i -> isJust (pick (V.unsafeIndex v i)))+        RowSink push freeze <- rowEncoderSink inner v (\x -> orPresent (pick x))+        pure (RowSink (\i r -> push i (orPresent (sel r))) (assemble valid <$> freeze))+{-# INLINE structEMaybe #-}+++childNames :: RowEncoder c -> [Text]+childNames = map fieldName . rowEncoderFields+++{- | Assemble an encoded struct. The children were built from the same+rows, so 'mkStruct' can only fail on a broken builder invariant.+-}+buildStruct :: String -> Int -> Maybe AC.Validity -> [Text] -> [ColumnArray] -> ColumnArray+buildStruct who n valid names cols =+  case mkStruct n valid (V.fromList (zip names cols)) of+    Right c -> c+    Left e -> errorWithoutStackTrace ("Arrow.Record." ++ who ++ ": " ++ e)+++-- ============================================================+-- RowDecoder+-- ============================================================++{- | Row decoder. Looks up named columns in a+'V.Vector ColumnArray' (keyed by the schema's field names), binds+the matching 'Decoder' to each once, and composes the resulting+row readers.++'RowDecoder' is an 'Applicative': combine several 'columnD'+calls with @<$>@ + @<*>@ to build a record. Field values are+evaluated when the record is built (no thunk per field).+-}+data RowDecoder r = RowDecoder+  { rowDecoderRequiredColumns :: [Text]+  {- ^ Names of columns the decoder consults when run.+  Order matches first appearance in the applicative chain;+  duplicates removed. Useful for column projection: a+  caller can ask the source format to only materialise+  these columns rather than the whole record batch.+  -}+  , rowDecoderBind :: Int -> V.Vector Text -> V.Vector ColumnArray -> Either String (Reader r)+  {- ^ Given the number of rows that will be read, the column names+  and the columns (parallel), check and bind every column the+  decoder needs and return the row reader.+  -}+  }+++instance Functor RowDecoder where+  fmap f (RowDecoder cs b) = RowDecoder cs $ \n names cols ->+    case b n names cols of+      Left e -> Left e+      Right g -> Right (\i -> let !x = g i in f x)+  {-# INLINE fmap #-}+++instance Applicative RowDecoder where+  pure x = RowDecoder [] (\_ _ _ -> Right (\_ -> x))+  {-# INLINE pure #-}+  RowDecoder cF bF <*> RowDecoder cX bX = RowDecoder+    (mergeRequired cF cX)+    $ \n names cols -> case bF n names cols of+      Left e -> Left e+      Right gf -> case bX n names cols of+        Left e -> Left e+        Right gx -> Right (\i -> let !x = gx i in gf i x)+  {-# INLINE (<*>) #-}+++{- | Order-preserving union of two 'rowDecoderRequiredColumns'+lists. Used by the Applicative instance to maintain the+"first appearance" order as decoders are combined.+-}+mergeRequired :: [Text] -> [Text] -> [Text]+mergeRequired xs ys = xs ++ filter (`notElem` xs) ys+++{- | Decode a batch: field names and columns are parallel and every+column has the same number of rows.+-}+runRowDecoder :: RowDecoder r -> V.Vector Field -> V.Vector ColumnArray -> Either String (V.Vector r)+runRowDecoder d fields cols+  | V.length fields /= V.length cols =+      Left $+        "Arrow.Record: schema has "+          ++ show (V.length fields)+          ++ " fields but the batch has "+          ++ show (V.length cols)+          ++ " columns"+  | V.any ((/= n) . columnLength) cols =+      Left $+        "Arrow.Record: column lengths differ ("+          ++ show (V.toList (V.map columnLength cols))+          ++ ")"+  | otherwise = case rowDecoderBind d n (V.map fieldName fields) cols of+      Left e -> Left e+      Right g -> Right (V.generate n (\(I# i) -> g i))+  where+    !n = if V.null cols then 0 else columnLength (V.unsafeHead cols)+++lookupColumn :: Text -> V.Vector Text -> V.Vector ColumnArray -> Maybe ColumnArray+lookupColumn name names cols = V.findIndex (== name) names >>= (cols V.!?)+++-- | Bind a decoder to a column that must hold at least @n@ rows.+bindColumn :: String -> Text -> Decoder a -> Int -> ColumnArray -> Either String (Reader a)+bindColumn who name d n col+  | columnLength col < n =+      Left $+        prefix+          ++ "column has "+          ++ show (columnLength col)+          ++ " rows, expected "+          ++ show n+  | otherwise = case decoderBind d col of+      Left e -> Left (prefix ++ e)+      Right g -> Right g+  where+    prefix = "Arrow.Record." ++ who ++ " " ++ show name ++ ": "+{-# INLINE bindColumn #-}+++{- | Decode the named column via the supplied 'Decoder'. Looks+the column up by 'Field' name in the schema the caller passes+to 'runRowDecoder'; returns 'Left' if the name isn't present.+-}+columnD :: Text -> Decoder a -> RowDecoder a+columnD name d = RowDecoder [name] $ \n names cols ->+  case lookupColumn name names cols of+    Nothing -> Left $ "Arrow.Record.columnD: no column named " ++ show name+    Just col -> bindColumn "columnD" name d n col+{-# INLINE columnD #-}+++{- | Like 'columnD' but supplies a default value if the+column is missing from the source schema. Useful for+schema-evolution: an older Parquet file dropped a column+that the Haskell record still wants; instead of failing,+the decoder substitutes the default.++Decoding errors on a /present/ column still propagate+(e.g. wrong type); only "no such column" falls back.+-}+columnDWithDefault :: Text -> a -> Decoder a -> RowDecoder a+columnDWithDefault name def d = RowDecoder [name] $ \n names cols ->+  case lookupColumn name names cols of+    Nothing -> Right (\_ -> def)+    Just col -> bindColumn "columnDWithDefault" name d n col+{-# INLINE columnDWithDefault #-}+++{- | Strategy for converting a record's selector name to its+on-the-wire column name. Mirrors the @renameStyle@ modifier+vocabulary in "Wireform.Derive".+-}+data NameStrategy+  = -- | Use the selector name unchanged.+    NameAsIs+  | {- | @userId@ → @user_id@. Inserts an underscore before any+    uppercase letter that follows a lowercase one and+    lower-cases the result.+    -}+    NameSnakeCase+  | {- | @user_id@ → @userId@. Drops underscores and+    upper-cases the following character.+    -}+    NameCamelCase+  | -- | @userId@ → @USER_ID@. snake-case + upper-case.+    NameUpperSnakeCase+  deriving (Show, Eq)+++-- | Apply a 'NameStrategy' to a 'Text' selector name.+applyNameStrategy :: NameStrategy -> Text -> Text+applyNameStrategy NameAsIs = id+applyNameStrategy NameSnakeCase = T.toLower . toSnake+  where+    toSnake t = T.pack (go ' ' (T.unpack t))+    -- Walk char-by-char carrying the previous character so we+    -- can decide whether to insert an underscore before an+    -- uppercase letter:+    --+    --   * insert when the previous char was lowercase (the+    --     usual word-boundary case: userId -> user_id)+    --   * insert when the next char is lowercase /and/ the+    --     previous char was uppercase (acronym→word boundary:+    --     userIDValue -> user_id_value, where the _ before V+    --     comes from this rule)+    --+    -- Simple, deterministic, and matches what 'inflection' /+    -- ActiveSupport / serde do.+    go _ [] = []+    go prev (c : cs)+      | isUp c+      , isLow prev+          || ( isUp prev && case cs of+                 (n : _) -> isLow n+                 [] -> False+             ) =+          '_' : c : go c cs+      | otherwise = c : go c cs+    isUp c = c >= 'A' && c <= 'Z'+    isLow c = c >= 'a' && c <= 'z'+applyNameStrategy NameCamelCase = toCamel+  where+    toCamel t =+      let parts = T.splitOn (T.pack "_") t+      in case parts of+           [] -> T.empty+           (p : ps) -> T.concat (T.toLower p : map cap ps)+    cap t+      | T.null t = t+      | otherwise = T.cons (toUpper1 (T.head t)) (T.tail t)+    toUpper1 c+      | c >= 'a' && c <= 'z' = toEnum (fromEnum c - 32)+      | otherwise = c+applyNameStrategy NameUpperSnakeCase =+  T.toUpper . applyNameStrategy NameSnakeCase+++{- | Inverse of 'structE': decode a struct column at the given name as+a record using the supplied inner 'RowDecoder'. The inner decoder+sees the struct's children by name; the outer decoder threads the+nested record into the parent record's applicative chain like any+other column. A struct with null rows needs 'structDMaybe'.+-}+structD :: Text -> RowDecoder c -> RowDecoder c+structD name inner = RowDecoder [name] $ \n names cols ->+  bindStruct "structD" name inner n names cols $ \col v g -> case v of+    Nothing -> Right g+    Just _ -> Left $ "Arrow.Record.structD " ++ show name ++ ": " ++ nullsErr col+{-# INLINE structD #-}+++{- | Like 'structD' but the struct may have null rows, which decode+to 'Nothing'. Structs without nulls are accepted too (every row+becomes 'Just').+-}+structDMaybe :: Text -> RowDecoder c -> RowDecoder (Maybe c)+structDMaybe name inner = RowDecoder [name] $ \n names cols ->+  bindStruct "structDMaybe" name inner n names cols $ \_ v g -> case v of+    Nothing -> Right (\i -> let !x = g i in Just x)+    Just _ -> Right (\i -> if isValidAt v (I# i) then let !x = g i in Just x else Nothing)+{-# INLINE structDMaybe #-}+++{- | Find the struct column, bind the inner decoder to its children+(each holds at least as many rows as the struct), and hand the+struct's validity and the child row reader to the continuation.+-}+bindStruct+  :: String+  -> Text+  -> RowDecoder c+  -> Int+  -> V.Vector Text+  -> V.Vector ColumnArray+  -> (ColumnArray -> Maybe AC.Validity -> Reader c -> Either String (Reader r))+  -> Either String (Reader r)+bindStruct who name inner n names cols k =+  case lookupColumn name names cols of+    Nothing -> Left $ prefix ++ "no column named " ++ show name+    Just col -> case col of+      ColStruct m v children+        | m < n -> Left $ prefix ++ "struct has " ++ show m ++ " rows, expected " ++ show n+        | otherwise ->+            case rowDecoderBind inner n (V.map fst children) (V.map snd children) of+              Left e -> Left (prefix ++ e)+              Right g -> k col v g+      other -> Left $ prefix ++ "expected ColStruct, got " ++ columnTag other+  where+    prefix = "Arrow.Record." ++ who ++ " " ++ show name ++ ": "+{-# INLINE bindStruct #-}+++-- ============================================================+-- Table+-- ============================================================++{- | Pairs a 'RowEncoder' with a 'RowDecoder' for one Haskell+record type. This is the handle you pass to the top-level+encode / decode helpers below.+-}+data Table r = Table+  { tableEncode :: RowEncoder r+  , tableDecode :: RowDecoder r+  }+++{- | Smart constructor. Equivalent to @Table enc dec@ but reads+better in call-site positions.+-}+table :: RowEncoder r -> RowDecoder r -> Table r+table = Table+{-# INLINE table #-}+++-- | Schema implied by the 'RowEncoder'.+tableSchema :: Table r -> Schema+tableSchema t =+  Schema+    { arrowFields = V.fromList (rowEncoderFields (tableEncode t))+    , arrowEndianness = Little+    , arrowMetadata = V.empty+    , arrowFeatures = V.empty+    }+++{- | Names of the columns the 'Table''s decoder needs. Equivalent+to @'rowDecoderRequiredColumns' . 'tableDecode'@; surfaced+here so callers can drive column projection through a+'Table' without unpacking the inner 'RowDecoder'.+-}+tableRequiredColumns :: Table r -> [Text]+tableRequiredColumns = rowDecoderRequiredColumns . tableDecode+++{- | Encode a vector of records as an Arrow batch + its schema.+The schema comes from 'tableSchema'; the batch is parallel to+'arrowFields' of that schema. Every column is built before the+pair is returned.+-}+encodeTable :: Table r -> V.Vector r -> (Schema, V.Vector ColumnArray)+encodeTable t rs =+  let !cols = V.fromList (runRowEncoder (tableEncode t) rs)+  in (tableSchema t, cols)+++{- | Decode an Arrow batch into a vector of records. Looks up+columns by schema field name; returns 'Left' on missing+columns, type mismatches, nulls in a non-nullable decoder, or a+malformed batch (field and column counts or column lengths differ).+-}+decodeTable+  :: Table r+  -> Schema+  -> V.Vector ColumnArray+  -> Either String (V.Vector r)+decodeTable t sch = runRowDecoder (tableDecode t) (arrowFields sch)+++-- ============================================================+-- Subset / projection+-- ============================================================++{- | Build a 'Table' for a subset of columns by name. The+resulting decoder ignores columns not in @keep@; the encoder+only emits the kept ones. Useful+for callers that have a single 'Table' and want to read or write+only a slice without writing a parallel @Table SubsetRecord@.++Returns 'Nothing' if any name in @keep@ isn't present in the+original table.+-}+subsetTable :: [Text] -> Table r -> Maybe (Table r)+subsetTable keep tbl = do+  let RowEncoder fields mk = tableEncode tbl+      byName nm = findIndex ((== nm) . fieldName) fields+  idxs <- traverse byName keep+  let !fieldVec = V.fromList fields+      pickCols cols = let !cv = V.fromList cols in map (V.unsafeIndex cv) idxs+  Just+    Table+      { tableEncode = RowEncoder (map (V.unsafeIndex fieldVec) idxs) $ \v g -> do+          RowSink push freeze <- mk v g+          pure (RowSink push (pickCols <$> freeze))+      , tableDecode = tableDecode tbl -- decoder looks columns up by name, so the subset is automatic+      }+++{- | Project an existing batch by column name, in the order+listed. Returns 'Nothing' if any name is missing.++Together with @'subsetTable'@ this lets callers reuse one+'Table' definition across read paths that materialise+different column subsets.+-}+projectTable+  :: [Text]+  -> Schema+  -> V.Vector ColumnArray+  -> Maybe (Schema, V.Vector ColumnArray)+projectTable keep sch cols = do+  let !fields = arrowFields sch+  idxs <- V.fromList <$> traverse (\nm -> V.findIndex ((== nm) . fieldName) fields) keep+  newCols <- traverse (cols V.!?) idxs+  pure (sch {arrowFields = V.map (V.unsafeIndex fields) idxs}, newCols)
+ src/Arrow/Record/Generic.hs view
@@ -0,0 +1,369 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE DefaultSignatures #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeOperators #-}++{- | @GHC.Generics@-backed 'Table' derivation for records whose+fields all have primitive Arrow encoders.++This module bridges "Arrow.Record"'s combinator API with+Haskell's 'Generic' machinery so that a record with+@deriving stock Generic@ can get an Arrow 'Table' instance in+one line:++@+data Trade = Trade { sym :: Text, qty :: Int32, note :: Maybe Text }+  deriving stock ('Generic')++tradeTable :: 'Table' Trade+tradeTable = 'genericTable'+@++Under the hood, 'genericTable' walks the 'Rep' to collect the+selector name + the 'HasEncoder' / 'HasDecoder' instance for+each field. Users who want a non-default representation for a+field type (e.g. a newtype wrapping @Int64@ that should+serialise as its underlying type) write a one-line+'HasEncoder' / 'HasDecoder' instance via+'Data.Functor.Contravariant.contramap' + 'fmap':++@+newtype UserId = UserId { unUserId :: Int64 }++instance 'HasEncoder' UserId where+  hasEncoder = 'contramap' unUserId 'int64E'+instance 'HasDecoder' UserId where+  hasDecoder = UserId \<$\> 'int64D'+@++For records that don't want to depend on 'Generic' (or that+want explicit column-name control), see "Arrow.Record.TH".+-}+module Arrow.Record.Generic (+  -- * Type classes for Generic deriving+  HasEncoder (..),+  HasDecoder (..),++  -- * Row-level (nested-record) classes+  HasRowEncoder (..),+  HasRowDecoder (..),++  -- * Deriver+  genericTable,+  genericRowEncoder,+  genericRowDecoder,++  -- * Re-exports for convenience+  Generic,+) where++import Arrow.Record (+  Decoder,+  Encoder,+  RowDecoder,+  RowEncoder,+  Table (..),+  binaryD,+  binaryE,+  boolD,+  boolE,+  columnD,+  doubleD,+  doubleE,+  fieldE,+  floatD,+  floatE,+  int16D,+  int16E,+  int32D,+  int32E,+  int64D,+  int64E,+  int8D,+  int8E,+  nullable,+  nullableD,+  utf8D,+  utf8E,+  word16D,+  word16E,+  word32D,+  word32E,+  word64D,+  word64E,+  word8D,+  word8E,+ )+import Data.ByteString (ByteString)+import Data.Functor.Contravariant (Contravariant (contramap))+import Data.Int (Int16, Int32, Int64, Int8)+import Data.Kind (Type)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Word (Word16, Word32, Word64, Word8)+import GHC.Generics+++-- ============================================================+-- HasEncoder / HasDecoder+-- ============================================================++{- | Bridge class between Haskell primitive types and+"Arrow.Record" 'Encoder' values. Only needed when deriving a+'Table' with 'genericTable'; hand-written 'RowEncoder's don't+touch this class.+-}+class HasEncoder a where+  hasEncoder :: Encoder a+++-- | Mirror of 'HasEncoder' on the decoding side.+class HasDecoder a where+  hasDecoder :: Decoder a+++{- | Per-record class used when a nested record needs to be+embedded as an Arrow @ColStruct@ column.++Default implementation uses 'genericRowEncoder', so any+'Generic' record whose fields all have 'HasEncoder' or+'HasEncoder (Maybe a)' instances picks up the instance for+free:++@+data Address = Address { city :: Text, zip :: Text }+  deriving stock ('Generic')++instance 'HasRowEncoder' Address+instance 'HasRowDecoder' Address+@++Once the inner record has these, lift it into a parent+'RowEncoder' / 'RowDecoder' with @'structE' "addr" addr+'hasRowEncoder'@ / @'structD' "addr" 'hasRowDecoder'@.+-}+class HasRowEncoder a where+  hasRowEncoder :: RowEncoder a+  default hasRowEncoder+    :: (Generic a, GRowEncoder (Rep a)) => RowEncoder a+  hasRowEncoder = genericRowEncoder+++-- | Mirror of 'HasRowEncoder' on the decoding side.+class HasRowDecoder a where+  hasRowDecoder :: RowDecoder a+  default hasRowDecoder+    :: (Generic a, GRowDecoder (Rep a)) => RowDecoder a+  hasRowDecoder = genericRowDecoder+++-- Primitive instances just delegate to the corresponding+-- combinator from "Arrow.Record".++instance HasEncoder Int8 where hasEncoder = int8E+++instance HasEncoder Int16 where hasEncoder = int16E+++instance HasEncoder Int32 where hasEncoder = int32E+++instance HasEncoder Int64 where hasEncoder = int64E+++instance HasEncoder Word8 where hasEncoder = word8E+++instance HasEncoder Word16 where hasEncoder = word16E+++instance HasEncoder Word32 where hasEncoder = word32E+++instance HasEncoder Word64 where hasEncoder = word64E+++instance HasEncoder Float where hasEncoder = floatE+++instance HasEncoder Double where hasEncoder = doubleE+++instance HasEncoder Bool where hasEncoder = boolE+++instance HasEncoder Text where hasEncoder = utf8E+++instance HasEncoder ByteString where hasEncoder = binaryE+++instance HasDecoder Int8 where hasDecoder = int8D+++instance HasDecoder Int16 where hasDecoder = int16D+++instance HasDecoder Int32 where hasDecoder = int32D+++instance HasDecoder Int64 where hasDecoder = int64D+++instance HasDecoder Word8 where hasDecoder = word8D+++instance HasDecoder Word16 where hasDecoder = word16D+++instance HasDecoder Word32 where hasDecoder = word32D+++instance HasDecoder Word64 where hasDecoder = word64D+++instance HasDecoder Float where hasDecoder = floatD+++instance HasDecoder Double where hasDecoder = doubleD+++instance HasDecoder Bool where hasDecoder = boolD+++instance HasDecoder Text where hasDecoder = utf8D+++instance HasDecoder ByteString where hasDecoder = binaryD+++{- | @'Maybe' a@ lifts through 'nullable' / 'nullableD'. The+@{\-# OVERLAPPING #-\}@ is needed because the instance head+'HasEncoder (Maybe a)' overlaps with the base instance when+GHC considers the final record field type.+-}+instance {-# OVERLAPPING #-} HasEncoder a => HasEncoder (Maybe a) where+  hasEncoder = nullable hasEncoder+  {-# INLINE hasEncoder #-}+++instance {-# OVERLAPPING #-} HasDecoder a => HasDecoder (Maybe a) where+  hasDecoder = nullableD hasDecoder+  {-# INLINE hasDecoder #-}+++-- ============================================================+-- Generic derivation+-- ============================================================++{- | Derive a 'Table' for any record whose field types all have+'HasEncoder' + 'HasDecoder' instances. Column names come from+the record selector names.++@+data Trade = Trade { sym :: Text, qty :: Int32 } deriving (Generic)+tradeTable :: 'Table' Trade+tradeTable = 'genericTable'+@+-}+genericTable+  :: forall r+   . ( Generic r+     , GRowEncoder (Rep r)+     , GRowDecoder (Rep r)+     )+  => Table r+genericTable =+  Table+    { tableEncode = genericRowEncoder @r+    , tableDecode = genericRowDecoder @r+    }+{-# INLINE genericTable #-}+++genericRowEncoder+  :: forall r+   . (Generic r, GRowEncoder (Rep r))+  => RowEncoder r+genericRowEncoder =+  contramap from (gRowEncoder :: RowEncoder (Rep r ()))+{-# INLINE genericRowEncoder #-}+++genericRowDecoder+  :: forall r+   . (Generic r, GRowDecoder (Rep r))+  => RowDecoder r+genericRowDecoder =+  to <$> (gRowDecoder :: RowDecoder (Rep r ()))+{-# INLINE genericRowDecoder #-}+++-- ============================================================+-- Generic Rep walkers+-- ============================================================++class GRowEncoder (f :: Type -> Type) where+  gRowEncoder :: RowEncoder (f p)+++class GRowDecoder (f :: Type -> Type) where+  gRowDecoder :: RowDecoder (f p)+++-- Datatype wrapper: transparent.+instance GRowEncoder f => GRowEncoder (M1 D m f) where+  gRowEncoder = unM1 `contramap` gRowEncoder+  {-# INLINE gRowEncoder #-}+++instance GRowDecoder f => GRowDecoder (M1 D m f) where+  gRowDecoder = M1 <$> gRowDecoder+  {-# INLINE gRowDecoder #-}+++-- Constructor wrapper: transparent.+instance GRowEncoder f => GRowEncoder (M1 C m f) where+  gRowEncoder = unM1 `contramap` gRowEncoder+  {-# INLINE gRowEncoder #-}+++instance GRowDecoder f => GRowDecoder (M1 C m f) where+  gRowDecoder = M1 <$> gRowDecoder+  {-# INLINE gRowDecoder #-}+++-- Product: concatenate via RowEncoder's Semigroup.+instance (GRowEncoder f, GRowEncoder g) => GRowEncoder (f :*: g) where+  gRowEncoder =+    ((\(l :*: _) -> l) `contramap` gRowEncoder)+      <> ((\(_ :*: r) -> r) `contramap` gRowEncoder)+  {-# INLINE gRowEncoder #-}+++instance (GRowDecoder f, GRowDecoder g) => GRowDecoder (f :*: g) where+  gRowDecoder = (:*:) <$> gRowDecoder <*> gRowDecoder+  {-# INLINE gRowDecoder #-}+++-- Record selector: pull the selector name via 'Selector' and+-- pair it with the field type's HasEncoder / HasDecoder instance.+instance (Selector s, HasEncoder a) => GRowEncoder (S1 s (Rec0 a)) where+  gRowEncoder =+    let selectorName = T.pack (selName (undefined :: S1 s (Rec0 a) p))+    in (unK1 . unM1)+         `contramap` fieldE selectorName id (hasEncoder :: Encoder a)+  {-# INLINE gRowEncoder #-}+++instance (Selector s, HasDecoder a) => GRowDecoder (S1 s (Rec0 a)) where+  gRowDecoder =+    let selectorName = T.pack (selName (undefined :: S1 s (Rec0 a) p))+    in (M1 . K1) <$> columnD selectorName (hasDecoder :: Decoder a)+  {-# INLINE gRowDecoder #-}
+ src/Arrow/Record/TH.hs view
@@ -0,0 +1,97 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TemplateHaskell #-}++{- | Template Haskell deriver for 'Arrow.Record.Table'.++When a record type doesn't have a 'GHC.Generics.Generic'+instance, or callers want explicit control over column names+(e.g. @"user_id"@ for a selector named @userId@), use the+'deriveTable' splice instead of @'Arrow.Record.Generic.genericTable'@.++@+{\-# LANGUAGE TemplateHaskell #-\}++import Arrow.Record.TH++data Trade = Trade { sym :: Text, qty :: Int32, note :: Maybe Text }++tradeTable :: 'Table' Trade+tradeTable = $('deriveTable' ''Trade)+@++or with a custom column-name transformer (e.g. snake_case):++@+tradeTable :: 'Table' Trade+tradeTable = $('deriveTableWith' toSnakeCase ''Trade)+  where+    toSnakeCase = ... :: String -> String+@++The splice emits code equivalent to a hand-written+'Arrow.Record.table' expression: 'Arrow.Record.fieldE' for+every record selector, 'Arrow.Record.columnD' on the decoder+side. Field types must have 'HasEncoder' / 'HasDecoder'+instances from "Arrow.Record.Generic".+-}+module Arrow.Record.TH (+  deriveTable,+  deriveTableWith,+) where++import Arrow.Record (+  columnD,+  fieldE,+  table,+ )+import Arrow.Record.Generic (HasDecoder (hasDecoder), HasEncoder (hasEncoder))+import Data.Text qualified as T+import Language.Haskell.TH+++{- | Derive a 'Table' for the named record type, using the+record selectors as column names verbatim.+-}+deriveTable :: Name -> Q Exp+deriveTable = deriveTableWith id+++{- | Like 'deriveTable' but renames columns via the supplied+pure function. Common choices: @'map' 'Data.Char.toLower'@,+a snake-case converter, or a full alias map via+@\\n -> 'Data.Map.findWithDefault' n n aliasMap@.+-}+deriveTableWith :: (String -> String) -> Name -> Q Exp+deriveTableWith renameField tyName = do+  info <- reify tyName+  fields <- case info of+    TyConI (DataD _ _ _ _ [RecC _ fs] _) -> pure fs+    TyConI (NewtypeD _ _ _ _ (RecC _ fs) _) -> pure fs+    _ ->+      fail $+        "Arrow.Record.TH.deriveTable: "+          ++ show tyName+          ++ " must be a single-constructor record type"++  -- Build the RowEncoder: mconcat of fieldE "name" selector hasEncoder.+  let mkEncoderPart (selName', _, _) = do+        let colName = renameField (nameBase selName')+        [|fieldE (T.pack colName) $(varE selName') hasEncoder|]+  encoderParts <- mapM mkEncoderPart fields+  let encoderExp = foldr1 (\a b -> [|$a <> $b|]) (map pure encoderParts)++  -- Build the RowDecoder: Ctor <$> columnD "name1" hasDecoder <*> ...+  let ctorName = case info of+        TyConI (DataD _ _ _ _ [RecC c _] _) -> c+        TyConI (NewtypeD _ _ _ _ (RecC c _) _) -> c+        _ -> error "unreachable: validated above"+  let mkDecoderPart (selName', _, _) =+        let colName = renameField (nameBase selName')+        in [|columnD (T.pack colName) hasDecoder|]+      initial = [|pure $(conE ctorName)|]+      combine acc step = [|$acc <*> $step|]+  decoderExp <- foldl combine initial (map mkDecoderPart fields)++  encoderExp' <- encoderExp+  -- Build: table encoderExp decoderExp+  [|table $(pure encoderExp') $(pure decoderExp)|]
+ src/Arrow/Stream.hs view
@@ -0,0 +1,1188 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}++{- | High-level, pyarrow-shaped Arrow IPC API.++95% of callers should reach for this module. It hides the+record-batch + dictionary-batch plumbing of "Arrow.FlatBufferIPC"+behind a single-call shape that mirrors+@pyarrow.ipc.new_stream@ / @pyarrow.ipc.open_stream@:++@+-- Write+case 'encodeArrowStream' 'defaultWriteOptions' schema batches of+  Right bytes -> ...+  Left  err   -> ...   -- the batches do not fit the schema++-- Read+case 'decodeArrowStream' bytes of+  Right (schema, batches) -> ...+  Left  err               -> ...+@++@batches@ is a list of @V.'V.Vector' 'ColumnArray'@, one column+per schema field, repeated once per record batch.++The eager encoders lay the whole output out in one allocation and copy+every body byte once. The lazy encoders ('encodeArrowStreamLazy',+'encodeArrowFileLazy') return a lazy 'BL.ByteString' whose chunks are+the message headers and the column buffers themselves (no body copy),+for 'BL.hPut' and socket writers. Decoded columns alias the input+bytes (see 'Arrow.Column.copyColumn' to detach them).++'ColDictionary' columns are handled automatically, at any depth and+inside dictionary values: the writer collects the dictionaries from+the input columns and emits a 'DictBatch' per id ahead of the first+record batch that references it (a dictionary nested inside another+dictionary's values ahead of the outer one); the reader resolves the+placeholder values column in returned 'ColDictionary' nodes against+the dictionary batches it saw earlier in the stream.++The writers check every batch against the schema first+('Arrow.Write.Columns.validateColumns') and report a mismatch, or a+dictionary they cannot write, as 'Left'.++For the file format ('encodeArrowFile' / 'decodeArrowFile') the+semantics are identical: same input shape, same output shape,+same dictionary handling, just an additional @ARROW1@ wrapper ++'Footer' index appended.++For lower-level control (custom dict ids, manual variadic+buffer counts, raw 'RecordBatchDef' construction, delta+dictionaries) drop down to "Arrow.FlatBufferIPC".+-}+module Arrow.Stream (+  -- * Encoding / decoding++  -- ** Streams (eager)+  encodeArrowStream,+  decodeArrowStream,++  -- ** Streams (lazy, zero-copy bodies)+  encodeArrowStreamLazy,++  -- ** Files (eager)+  encodeArrowFile,+  decodeArrowFile,++  -- ** Files (lazy, zero-copy bodies)+  encodeArrowFileLazy,++  -- ** Streams (incremental / iterator)+  StreamReader,+  openStreamReader,+  streamReaderSchema,+  streamReaderNext,+  streamReaderToList,+  streamReaderIter,+  streamReaderProjected,+  streamReaderProjectedIter,++  -- ** Files (incremental / iterator)+  FileReader,+  openFileReader,+  fileReaderSchema,+  fileReaderNext,+  fileReaderToList,+  fileReaderIter,++  -- ** Handles (semi-streaming, one section live at a time)+  fileBatchesHandle,+  streamBatchesHandle,++  -- * Options+  WriteOptions (..),+  DictHandling (..),+  defaultWriteOptions,+  ReadOptions (..),+  defaultReadOptions,+  BodyCompressionCodec (..),+  bodyCompressionAvailable,+) where++import Arrow.Column (+  ColumnArray,+  columnLength,+  concatColumnArray,+  concatColumnArrays,+  resolveDictionaryColumn,+  sliceColumnArray,+ )+import Arrow.Column.Internal (IntegralPrim (..), PrimType (..), SomePrimType (..), integralPrim, primTypeFor)+import Arrow.Column.Internal qualified as I+import Arrow.FlatBufferIPC (+  ArrowBlock (..),+  DictBatch (..),+  FileIndex (..),+  StreamFrame (..),+  decodeMessageFrame,+  parseFileIndex,+  readArrowFileFBWithDicts,+  readArrowStreamFBInterleaved,+  readFileIndex,+ )+import Arrow.FlatBufferIPC.Write (+  Frame,+  batchFrame,+  bodyCompressionAvailable,+  dictionaryFrame,+  encodeBatchWithEndianness,+  renderFile,+  renderFileLazy,+  renderStream,+  renderStreamLazy,+  schemaFrame,+ )+import Arrow.Read.Columns (ReadOptions (..), decodeDictionaryBatch, decodeRecordBatch, defaultReadOptions)+import Arrow.Types (+  ArrowType (..),+  BodyCompressionCodec (..),+  DictionaryEncoding (..),+  Field (..),+  Message (..),+  RecordBatchDef,+  Schema (..),+ )+import Arrow.Write.Columns (validateColumns, validateDecimalFields)+import Columnar.Stream qualified as IS+import Data.Bits (complement, shiftL, (.&.), (.|.))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Lazy qualified as BL+import Data.Foldable (foldlM)+import Data.IORef (IORef, newIORef, readIORef, writeIORef)+import Data.Int (Int32, Int64)+import Data.Map.Strict qualified as Map+import Data.Set qualified as Set+import Data.Text (Text)+import Data.Traversable (mapAccumM)+import Data.Vector qualified as V+import Data.Vector.Storable qualified as VS+import System.IO (Handle, SeekMode (..), hFileSize, hSeek)+++-- ============================================================+-- Write options+-- ============================================================++{- | Arrow IPC writer configuration. Construct one with+'defaultWriteOptions' and override the fields you care about.+-}+data WriteOptions = WriteOptions+  { writeBodyCompression :: !(Maybe BodyCompressionCodec)+  {- ^ When 'Just', record and dictionary batch buffers are+  compressed per Arrow's 'BodyCompression' table. 'Nothing' (the+  default) leaves buffers uncompressed. A codec that was not+  compiled in (cabal flags @zstd@ / @lz4@; check with+  'bodyCompressionAvailable') is ignored and the batches are+  written uncompressed, so the output is always readable.+  -}+  , writeDictHandling :: !DictHandling+  {- ^ How the writer treats repeated dictionary ids across+  batches: emit a single dict up front, or emit a fresh+  @isDelta=false@ replacement before any batch whose+  dictionary differs. Default: 'DictEmitOnce'. PyArrow's+  default matches 'DictReplaceOnChange'; use that when+  encoding streams destined for pyarrow-with-varying-dicts.+  -}+  }+  deriving (Show, Eq)+++-- | Dictionary-emission strategy for the high-level writer.+data DictHandling+  = {- | Emit one dictionary batch per id ahead of the first record+    batch. When batches carry different value columns for one id,+    the distinct value columns are concatenated into one dictionary+    and every batch's keys are shifted into it. Fails ('Left')+    when the combined dictionary has more values than the field's+    index type can address, or when the value columns cannot be+    concatenated. This is the default, and the only strategy the+    file format allows.+    -}+    DictEmitOnce+  | {- | Emit a fresh @isDelta=false@ replacement dict batch+    before any record batch whose dictionary values differ+    from the most-recently-emitted dict for that id (and before+    it, the dictionaries nested in its values that changed; an+    outer dictionary is re-sent whenever a dictionary nested in its+    values is). Matches the semantics pyarrow uses when a stream+    writer sees a changed dict.+    -}+    DictReplaceOnChange+  deriving (Show, Eq)+++-- | Defaults: no body compression, emit-once dictionaries.+defaultWriteOptions :: WriteOptions+defaultWriteOptions =+  WriteOptions+    { writeBodyCompression = Nothing+    , writeDictHandling = DictEmitOnce+    }+++-- ============================================================+-- Streams+-- ============================================================++{- | Encode a sequence of column-major batches as a self-contained+Arrow IPC /stream/. Equivalent to pyarrow's+@ipc.new_stream + write_batch + close@; see 'WriteOptions' for the+knobs and 'defaultWriteOptions' for pyarrow-like defaults.++Every batch is first checked against the schema+('Arrow.Write.Columns.validateColumns': column count, equal column+lengths, column shapes, nulls only under nullable fields, nested+lengths); a mismatch is a 'Left'.++Dictionary-encoded columns ('ColDictionary' anywhere in the column+tree, including inside dictionary values) are handled automatically,+see 'DictHandling'. Dictionary batches use the value field the schema+declares for the id (type, children), and a dictionary nested in+another dictionary's values is written before the outer one. Keys are+written at the width of the field's index type.++The output is one allocation holding every body byte exactly once.+-}+encodeArrowStream+  :: WriteOptions+  -> Schema+  -> [V.Vector ColumnArray]+  -> Either String ByteString+encodeArrowStream opts sch batches = renderStream <$> streamFrames opts sch batches+++{- | 'encodeArrowStream' as a lazy 'BL.ByteString'. Its chunks are the+message headers, the column buffers themselves (aliased, not copied)+and shared padding, so the work is proportional to the number of+buffers, not the data size. Same checks and dictionary handling.+-}+encodeArrowStreamLazy+  :: WriteOptions+  -> Schema+  -> [V.Vector ColumnArray]+  -> Either String BL.ByteString+encodeArrowStreamLazy opts sch batches = renderStreamLazy <$> streamFrames opts sch batches+++-- | Every frame of the stream except the end-of-stream marker.+streamFrames :: WriteOptions -> Schema -> [V.Vector ColumnArray] -> Either String [Frame]+streamFrames opts sch batches = do+  validateDecimalFields (arrowFields sch)+  mapM_ (validateColumns (arrowFields sch)) batches+  case writeDictHandling opts of+    DictEmitOnce -> do+      Compiled dicts rbs <- compileBatches opts sch batches+      Right (schemaFrame sch : dicts ++ rbs)+    DictReplaceOnChange -> (schemaFrame sch :) <$> replaceDictFrames opts sch batches+++{- | Frames for the 'DictReplaceOnChange' strategy. Each record batch+is preceded by fresh @isDelta=false@ dict batches for any dictionary+id whose values have changed since the last emission, plus every+dictionary whose values contain one of those.+-}+replaceDictFrames :: WriteOptions -> Schema -> [V.Vector ColumnArray] -> Either String [Frame]+replaceDictFrames opts sch = go Map.empty []+  where+    !mCodec = writeBodyCompression opts+    !plan = dictPlan sch+    go _ acc [] = Right (concat (reverse acc))+    go !lastDicts acc (cols0 : rest) = do+      (curDicts, rebased) <- unifyDictionaries plan [cols0]+      let !cols = case rebased of+            [c] -> c+            _ -> cols0+          !changed = Map.keysSet (Map.filterWithKey (\did vs -> Map.lookup did lastDicts /= Just vs) curDicts)+          !emitDicts = Map.restrictKeys curDicts (withDependents plan changed)+      dictFrames <- buildDictFrames mCodec sch plan emitDicts+      eb <- encodeBatchWithEndianness (arrowEndianness sch) mCodec (arrowFields sch) cols+      go (Map.union curDicts lastDicts) ((dictFrames ++ [batchFrame eb]) : acc) rest+++{- | Decode an Arrow IPC stream back into 'Schema' + record-batch+column vectors. Dictionary references are resolved+transparently: the returned 'ColumnArray' values contain the+actual dictionary values, not just keys. Columns alias the input.+-}+decodeArrowStream+  :: ReadOptions+  -> ByteString+  -> Either String (Schema, [V.Vector ColumnArray])+decodeArrowStream opts bs = do+  (sch, frames) <- readArrowStreamFBInterleaved bs+  decodeInterleavedFrames opts sch frames+++-- ============================================================+-- Files+-- ============================================================++{- | Encode batches as an Arrow IPC /file/ (@ARROW1@ header, the+stream payload, a FlatBuffers 'Footer' indexing every dictionary and+record batch, @ARROW1@ trailer). Batches are checked against the+schema as in 'encodeArrowStream'. Dictionaries are always unified+into one dictionary batch per id ('DictEmitOnce', whatever+'writeDictHandling' says), because the file format does not allow+replacement dictionaries.+-}+encodeArrowFile+  :: WriteOptions+  -> Schema+  -> [V.Vector ColumnArray]+  -> Either String ByteString+encodeArrowFile opts sch batches = do+  validateDecimalFields (arrowFields sch)+  mapM_ (validateColumns (arrowFields sch)) batches+  Compiled dicts rbs <- compileBatches opts sch batches+  Right (renderFile sch dicts rbs)+++-- | 'encodeArrowFile' as a lazy 'BL.ByteString' aliasing the column buffers (see 'encodeArrowStreamLazy').+encodeArrowFileLazy+  :: WriteOptions+  -> Schema+  -> [V.Vector ColumnArray]+  -> Either String BL.ByteString+encodeArrowFileLazy opts sch batches = do+  validateDecimalFields (arrowFields sch)+  mapM_ (validateColumns (arrowFields sch)) batches+  Compiled dicts rbs <- compileBatches opts sch batches+  Right (renderFileLazy sch dicts rbs)+++{- | Decode an Arrow IPC file. Dictionary-resolved 'ColumnArray'+values are returned just like 'decodeArrowStream'.+-}+decodeArrowFile+  :: ReadOptions+  -> ByteString+  -> Either String (Schema, [V.Vector ColumnArray])+decodeArrowFile opts bs = do+  (sch, dicts, frames) <- readArrowFileFBWithDicts bs+  decodeBatches opts sch dicts frames+++-- ============================================================+-- Internal: dictionary planning and unification+-- ============================================================++-- | Dictionary frames (in 'dpOrder') and record batch frames.+data Compiled = Compiled ![Frame] ![Frame]+++{- | Unify the dictionaries of every batch (see 'unifyDictionaries'),+then lay out one dictionary batch per id (nested dictionaries first,+see 'dictPlan') and every record batch.+-}+compileBatches :: WriteOptions -> Schema -> [V.Vector ColumnArray] -> Either String Compiled+compileBatches opts sch batches = do+  let !mCodec = writeBodyCompression opts+      !plan = dictPlan sch+  (dictMap, rebased) <- unifyDictionaries plan batches+  dicts <- buildDictFrames mCodec sch plan dictMap+  rbs <- traverse (fmap batchFrame . encodeBatchWithEndianness (arrowEndianness sch) mCodec (arrowFields sch)) rebased+  Right (Compiled dicts rbs)+++{- | What the writer needs to know about the schema's dictionaries:+the order to emit them in, the narrowest index type declared for each+id, and the ids nested inside each id's value type.+-}+data DictPlan = DictPlan+  { dpOrder :: ![Int64]+  -- ^ Every dictionary id, a dictionary nested inside another's values before the outer one.+  , dpIndexType :: !(Map.Map Int64 ArrowType)+  -- ^ The narrowest index type over the fields sharing the id.+  , dpNested :: !(Map.Map Int64 (Set.Set Int64))+  -- ^ Ids declared anywhere inside the id's value type.+  }+++-- | Walk the schema in post-order (children before parents) collecting the 'DictPlan'.+dictPlan :: Schema -> DictPlan+dictPlan sch =+  let ((order, ixTypes, nested), _) = goFields ([], Map.empty, Map.empty) (arrowFields sch)+  in DictPlan {dpOrder = reverse order, dpIndexType = ixTypes, dpNested = nested}+  where+    goFields st fs = V.foldl' (\(s, ids) f -> let (s', ids') = goField s f in (s', Set.union ids ids')) (st, Set.empty) fs+    goField st f =+      let ((order, ixTypes, nested), inner) = goFields st (fieldChildren f)+      in case fieldDictionary f of+           Nothing -> ((order, ixTypes, nested), inner)+           Just de ->+             let !did = deId de+             in (+                  ( if did `elem` order then order else did : order+                  , Map.insertWith narrower did (deIndexType de) ixTypes+                  , Map.insertWith Set.union did inner nested+                  )+                , Set.insert did inner+                )+    narrower new old = if maxIndexFor new < maxIndexFor old then new else old+++-- | Largest dictionary key an index type holds (signed 32-bit when the type is not an integer).+maxIndexFor :: ArrowType -> Int64+maxIndexFor = \case+  AInt 8 True -> 127+  AInt 8 False -> 255+  AInt 16 True -> 32767+  AInt 16 False -> 65535+  AInt 32 False -> 4294967295+  AInt 64 _ -> maxBound+  _ -> 2147483647+++-- | The key element type for an index type (signed 32-bit when the type is not an integer).+keyPrimFor :: ArrowType -> SomePrimType+keyPrimFor ty = case primTypeFor ty of+  Just (SomePrimType t) | Just IntegralPrim <- integralPrim t -> SomePrimType t+  _ -> SomePrimType PInt32+++-- | The ids together with every dictionary whose values contain one of them.+withDependents :: DictPlan -> Set.Set Int64 -> Set.Set Int64+withDependents plan ids =+  Set.union ids (Map.keysSet (Map.filter (not . Set.disjoint ids) (dpNested plan)))+++{- | One dictionary batch frame per id present in the map, in+'dpOrder'. The value field is the dictionary field itself minus its+encoding (same type and children).+-}+buildDictFrames :: Maybe BodyCompressionCodec -> Schema -> DictPlan -> Map.Map Int64 ColumnArray -> Either String [Frame]+buildDictFrames mCodec sch plan dictMap = traverse one (filter (`Map.member` dictMap) (dpOrder plan))+  where+    one did = case findDictField did (arrowFields sch) of+      Nothing -> Left ("Arrow.Stream: dictionary id " ++ show did ++ " has no field in the schema")+      Just f -> do+        let !valuesField = f {fieldDictionary = Nothing, fieldNullable = True}+        eb <- encodeBatchWithEndianness (arrowEndianness sch) mCodec (V.singleton valuesField) (V.singleton (Map.findWithDefault (I.ColNull 0) did dictMap))+        Right (dictionaryFrame did False eb)+++-- | Distinct value columns seen for one id (newest first, with their offsets) and their total length.+data DictAcc = DictAcc ![(ColumnArray, Int64)] !Int64+++{- | Collect the dictionaries of a run of batches into one values+column per id. Walking the batches in order, each dictionary column's+values are appended to its id's combined dictionary unless an equal+values column was already appended, and the column's keys are+shifted by the position of its values inside the combined dictionary+(written at the id's narrowest index type), so every batch can be+written against a single dictionary batch.++Dictionary columns inside dictionary values are collected the same+way, before their outer value column is placed, so the outer+dictionary is written with keys into the combined inner one. The+rebased columns carry an empty values column of the right type (the+writer only needs the keys), which keeps the concatenation of+outer value columns from shifting already rebased inner keys.++Fails when the value columns of one id cannot be concatenated, or+when the combined dictionary has more values than the id's index+type can address.+-}+unifyDictionaries :: DictPlan -> [V.Vector ColumnArray] -> Either String (Map.Map Int64 ColumnArray, [V.Vector ColumnArray])+unifyDictionaries plan batches = do+  (seen, rebased) <- mapAccumM (mapAccumM (mapAccumDictionaries shiftOne)) Map.empty batches+  dicts <- Map.traverseWithKey combine seen+  Right (dicts, rebased)+  where+    indexType did = Map.findWithDefault (AInt 32 True) did (dpIndexType plan)+    shiftOne st col = case col of+      I.ColDictionary did keys vals -> do+        (st1, vals') <- mapAccumDictionaries shiftOne st vals+        (st2, off) <- place st1 did vals'+        keys' <- if off == 0 then Right keys else shiftKeys (keyPrimFor (indexType did)) off keys+        Right (st2, I.ColDictionary did keys' (sliceColumnArray 0 0 vals'))+      _ -> Right (st, col)+    place st did vals =+      let !len = fromIntegral (columnLength vals) :: Int64+      in case Map.lookup did st of+           Nothing -> Right (Map.insert did (DictAcc [(vals, 0)] len) st, 0)+           Just (DictAcc entries total) -> case lookup vals entries of+             Just off -> Right (st, off)+             Nothing -> Right (Map.insert did (DictAcc ((vals, total) : entries) (total + len)) st, total)+    combine did (DictAcc entries total)+      | total - 1 > maxIx =+          Left+            ( "Arrow.Stream: dictionary id "+                ++ show did+                ++ ": the batches' dictionaries combine to "+                ++ show total+                ++ " values but the field's index type addresses at most "+                ++ show (toInteger maxIx + 1)+                ++ " (use DictReplaceOnChange or a wider index type)"+            )+      | otherwise = case reverse entries of+          [(v, _)] -> Right v+          ordered -> case concatColumnArrays (map fst ordered) of+            Right v -> Right v+            Left e -> Left ("Arrow.Stream: dictionary id " ++ show did ++ ": value columns cannot be combined: " ++ e)+      where+        !maxIx = maxIndexFor (indexType did)+++{- | Keys plus an offset, at the given integer type (one pass). Valid+keys stay below the combined dictionary length, which 'combine' checks+against the index type, so they never wrap; null slots may.+-}+shiftKeys :: SomePrimType -> Int64 -> ColumnArray -> Either String ColumnArray+shiftKeys (SomePrimType tt) off keys = case keys of+  I.ColPrim kt v ks+    | Just IntegralPrim <- integralPrim kt+    , Just IntegralPrim <- integralPrim tt ->+        Right (I.ColPrim tt v (VS.map (\k -> fromIntegral (fromIntegral k + off)) ks))+  _ -> Left "Arrow.Stream: dictionary keys must be an integer column"+++{- | Thread a state through every dictionary column of a column tree+(pre-order, children left to right), rebuilding the tree. Dictionary+values are not entered; the callback decides what to do with them.+-}+mapAccumDictionaries+  :: (s -> ColumnArray -> Either String (s, ColumnArray))+  -> s+  -> ColumnArray+  -> Either String (s, ColumnArray)+mapAccumDictionaries f = go+  where+    one st c k = (\(st', c') -> (st', k c')) <$> go st c+    go st col = case col of+      I.ColDictionary {} -> f st col+      I.ColStruct n v cs -> (\(st', kids) -> (st', I.ColStruct n v kids)) <$> mapAccumM named st cs+      I.ColList v o c -> one st c (I.ColList v o)+      I.ColLargeList v o c -> one st c (I.ColLargeList v o)+      I.ColFixedSizeList w n v c -> one st c (I.ColFixedSizeList w n v)+      I.ColMap v o k x -> do+        (st1, k') <- go st k+        (st2, x') <- go st1 x+        Right (st2, I.ColMap v o k' x')+      I.ColDenseUnion ts o cs -> (\(st', kids) -> (st', I.ColDenseUnion ts o kids)) <$> mapAccumM go st cs+      I.ColSparseUnion ts cs -> (\(st', kids) -> (st', I.ColSparseUnion ts kids)) <$> mapAccumM go st cs+      I.ColRunEndEncoded off n re vs -> one st vs (I.ColRunEndEncoded off n re)+      I.ColListView v o s c -> one st c (I.ColListView v o s)+      I.ColLargeListView v o s c -> one st c (I.ColLargeListView v o s)+      _ -> Right (st, col)+    named st (nm, c) = (\(st', c') -> (st', (nm, c'))) <$> go st c+++{- | Locate the 'Field' whose 'fieldDictionary' carries the given+id, walking nested children depth-first.+-}+findDictField :: Int64 -> V.Vector Field -> Maybe Field+findDictField did = goVec+  where+    goVec fs = goList (V.toList fs)+    goList [] = Nothing+    goList (f : fs) = case fieldDictionary f of+      Just de+        | deId de == did ->+            Just f+      _ -> case goVec (fieldChildren f) of+        Just g -> Just g+        Nothing -> goList fs+++-- ============================================================+-- Internal: decoding+-- ============================================================++{- | Decode the file's dictionary batches (in file order, honouring+delta batches) and its record batches, resolving every dictionary+column against the final dictionaries.+-}+decodeBatches+  :: ReadOptions+  -> Schema+  -> [DictBatch]+  -> [(RecordBatchDef, ByteString)]+  -> Either String (Schema, [V.Vector ColumnArray])+decodeBatches opts sch dicts frames = do+  dictMap <- foldlM (applyDictBatch opts sch) Map.empty dicts+  resolved <- traverse (\(rb, body) -> decodeResolved opts sch dictMap rb body) frames+  Right (sch, resolved)+++-- | One record batch with its dictionary columns resolved against the dictionaries in force.+decodeResolved :: ReadOptions -> Schema -> Map.Map Int64 ColumnArray -> RecordBatchDef -> ByteString -> Either String (V.Vector ColumnArray)+decodeResolved opts sch dictMap rb body = do+  cols <- decodeRecordBatch opts sch rb body+  V.mapM (resolveDictionaryColumn (`Map.lookup` dictMap)) cols+++{- | Fold one dictionary batch into the id-to-values map: a delta+batch appends to the existing dictionary ('concatColumnArray'; a+delta whose values cannot be appended is an error), any other batch+replaces it.++Dictionary columns inside the values (a dictionary whose value type+contains dictionary-encoded fields) are resolved against the+dictionaries in force when the batch arrives, as Arrow C++ does: a+later replacement or delta of the inner dictionary affects the outer+dictionaries sent after it, not the ones already received.+-}+applyDictBatch :: ReadOptions -> Schema -> Map.Map Int64 ColumnArray -> DictBatch -> Either String (Map.Map Int64 ColumnArray)+applyDictBatch opts sch dictMap db = do+  (did, raw) <- decodeDictionaryBatch opts sch db+  vals <- resolveDictionaryColumn (`Map.lookup` dictMap) raw+  if dbIsDelta db+    then case Map.lookup did dictMap of+      Nothing -> Right (Map.insert did vals dictMap)+      Just old -> case concatColumnArray old vals of+        Right new -> Right (Map.insert did new dictMap)+        Left e -> Left ("Arrow.Stream: delta dictionary batch for id " ++ show did ++ ": " ++ e)+    else Right (Map.insert did vals dictMap)+++{- | Decode a stream-order frame list (dict + record batches+interleaved) while honouring replacement / delta dictionaries:+each record batch resolves against the dictionaries in force at its+position in the stream.+-}+decodeInterleavedFrames+  :: ReadOptions+  -> Schema+  -> [StreamFrame]+  -> Either String (Schema, [V.Vector ColumnArray])+decodeInterleavedFrames opts sch = go Map.empty []+  where+    go _ acc [] = Right (sch, reverse acc)+    go !dictMap acc (SFDict db : rest) = do+      newMap <- applyDictBatch opts sch dictMap db+      go newMap acc rest+    go !dictMap acc (SFBatch rb body : rest) = do+      resolved <- decodeResolved opts sch dictMap rb body+      go dictMap (resolved : acc) rest+++-- ============================================================+-- Streaming reader+-- ============================================================++{- | An iterator-style handle for reading an Arrow IPC stream+batch-by-batch, mirroring pyarrow's @ipc.RecordBatchStreamReader@:++@+case 'openStreamReader' bytes of+  Left  e  -> handleError e+  Right rd -> do+    let !sch = 'streamReaderSchema' rd+    loop rd+  where+    loop rd = case 'streamReaderNext' rd of+      Right (Just (cols, rd')) -> consume cols >> loop rd'+      Right Nothing            -> finish ()+      Left  e                  -> handleError e+@++Dictionary batches are applied (replacements and deltas) as+'streamReaderNext' walks past them, so each record batch sees the+dictionaries in force at its position. Use 'streamReaderToList' to+drain the iterator into a list (equivalent to 'decodeArrowStream' but+keeps the streaming shape for callers that want incremental+processing later).+-}+data StreamReader = StreamReader+  { srSchema :: !Schema+  , srOpts :: !ReadOptions+  -- ^ Decode budgets applied to every batch 'streamReaderNext' decodes.+  , srDictMap :: !(Map.Map Int64 ColumnArray)+  -- ^ Dictionaries in force at the current stream position.+  , srFrames :: ![StreamFrame]+  }+++{- | Initialise an iterator from raw stream bytes. Parses the schema+and splits the stream into frames; nothing is decoded until+'streamReaderNext'.+-}+openStreamReader :: ReadOptions -> ByteString -> Either String StreamReader+openStreamReader opts bs = do+  (sch, frames) <- readArrowStreamFBInterleaved bs+  Right+    StreamReader+      { srSchema = sch+      , srOpts = opts+      , srDictMap = Map.empty+      , srFrames = frames+      }+++-- | The schema decoded at 'openStreamReader' time.+streamReaderSchema :: StreamReader -> Schema+streamReaderSchema = srSchema+++{- | Pull the next record batch from the iterator. Returns:++  * @Right (Just (cols, rd'))@: a decoded batch with+    dictionary references resolved, plus a continuation+    reader for the remaining frames.+  * @Right Nothing@: the stream's EOS marker has been+    consumed; no more batches.+  * @Left e@: a parse / validation error.+-}+streamReaderNext+  :: StreamReader+  -> Either String (Maybe (V.Vector ColumnArray, StreamReader))+streamReaderNext rd = case srFrames rd of+  [] -> Right Nothing+  (SFDict db : rest) -> do+    dictMap <- applyDictBatch (srOpts rd) (srSchema rd) (srDictMap rd) db+    streamReaderNext rd {srDictMap = dictMap, srFrames = rest}+  (SFBatch rb body : rest) -> do+    resolved <- decodeResolved (srOpts rd) (srSchema rd) (srDictMap rd) rb body+    Right (Just (resolved, rd {srFrames = rest}))+++{- | Drain a 'StreamReader' into a list of batches. Equivalent+(modulo the schema bundling) to 'decodeArrowStream' but keeps+the streaming shape for callers that prefer to iterate.+-}+streamReaderToList+  :: StreamReader+  -> Either String [V.Vector ColumnArray]+streamReaderToList rd0 = go rd0 []+  where+    go rd acc = case streamReaderNext rd of+      Left e -> Left e+      Right Nothing -> Right (reverse acc)+      Right (Just (cols, rd')) -> go rd' (cols : acc)+++{- | Lift a 'StreamReader' into the cross-format 'IS.Iter' shape.+Each 'IS.iterStep' decodes one record batch and resolves dictionary+references, matching 'streamReaderNext' semantics without forcing+the caller to thread the continuation reader by hand.+-}+streamReaderIter :: StreamReader -> IS.Iter (V.Vector ColumnArray)+streamReaderIter rd0 = IS.iterUnfold rd0 streamReaderNext+++-- ============================================================+-- File reader (footer-indexed sections)+-- ============================================================++{- | An iterator-style handle for reading an Arrow IPC /file/+batch-by-batch through its 'FileIndex' (footer) sections: the file+counterpart of 'StreamReader'. Opening parses the footer and decodes+the dictionary batches (every record batch needs them); each+'fileReaderNext' decodes exactly one record batch section and+touches nothing else in the file.++The input is the whole file in memory (columns alias it). For files+too large for that, 'fileBatchesHandle' seeks and reads one section+at a time.+-}+data FileReader = FileReader+  { frSchema :: !Schema+  , frOpts :: !ReadOptions+  , frDictMap :: !(Map.Map Int64 ColumnArray)+  , frBytes :: !ByteString+  , frBlocks :: !(V.Vector ArrowBlock)+  , frIx :: !Int+  }+++{- | Open a whole in-memory Arrow IPC file for section-at-a-time+decoding.+-}+openFileReader :: ReadOptions -> ByteString -> Either String FileReader+openFileReader opts bs = do+  idx <- readFileIndex bs+  dictMap <- readDictBlocks opts idx bs+  Right+    FileReader+      { frSchema = fiSchema idx+      , frOpts = opts+      , frDictMap = dictMap+      , frBytes = bs+      , frBlocks = fiRecordBatches idx+      , frIx = 0+      }+++-- | The schema the file's footer declares.+fileReaderSchema :: FileReader -> Schema+fileReaderSchema = frSchema+++{- | Pull the next record batch section: @Right (Just (cols, fr'))@+for one decoded batch (dictionary references resolved), @Right+Nothing@ past the last section, @Left@ for a parse / validation+error. Only the section the next 'ArrowBlock' names is touched.+-}+fileReaderNext :: FileReader -> Either String (Maybe (V.Vector ColumnArray, FileReader))+fileReaderNext fr+  | frIx fr >= V.length (frBlocks fr) = Right Nothing+  | otherwise = do+      let !blk = V.unsafeIndex (frBlocks fr) (frIx fr)+      (msg, body) <- blockMessage (frBytes fr) blk+      rb <- case msg of+        RecordBatch def -> Right def+        _ -> Left "Arrow.Stream: file section is not a record batch"+      resolved <- decodeResolved (frOpts fr) (frSchema fr) (frDictMap fr) rb body+      Right (Just (resolved, fr {frIx = frIx fr + 1}))+++-- | 'fileReaderNext' as a cross-format 'IS.Iter'.+fileReaderIter :: FileReader -> IS.Iter (V.Vector ColumnArray)+fileReaderIter fr0 = IS.iterUnfold fr0 fileReaderNext+++-- | Drain a 'FileReader' into a list of batches.+fileReaderToList :: FileReader -> Either String [V.Vector ColumnArray]+fileReaderToList fr = IS.iterToList (fileReaderIter fr)+++-- | Decode the message and body one file section holds.+blockMessage :: ByteString -> ArrowBlock -> Either String (Message, ByteString)+blockMessage bs (ArrowBlock off metaLen bodyLen)+  | off < 0 || metaLen <= 0 || bodyLen < 0 || off > n - fromIntegral metaLen =+      Left "Arrow.Stream: file section runs past the end of the input"+  | bodyLen > n - off - fromIntegral metaLen =+      Left "Arrow.Stream: file section body runs past the end of the input"+  | otherwise = do+      (msg, declaredBody) <-+        decodeMessageFrame (BS.take (fromIntegral metaLen) (BS.drop (fromIntegral off) bs))+      if declaredBody < 0 || declaredBody > bodyLen+        then Left "Arrow.Stream: file section body is shorter than its message declares"+        else+          let !body = BS.take (fromIntegral bodyLen) (BS.drop (fromIntegral (off + fromIntegral metaLen)) bs)+          in Right (msg, body)+  where+    !n = fromIntegral (BS.length bs) :: Int64+++{- | The dictionary batches one file's dictionary sections hold,+folded into the id-to-values map in file order (replacement and+delta semantics as in 'applyDictBatch').+-}+readDictBlocks :: ReadOptions -> FileIndex -> ByteString -> Either String (Map.Map Int64 ColumnArray)+readDictBlocks opts idx bs = go Map.empty 0+  where+    !n = V.length (fiDictBlocks idx)+    go !m !i+      | i >= n = Right m+      | otherwise = do+          (msg, body) <- blockMessage bs (V.unsafeIndex (fiDictBlocks idx) i)+          db <- case msg of+            DictionaryBatch did isDelta def ->+              Right DictBatch {dbId = did, dbIsDelta = isDelta, dbData = def, dbBody = body}+            _ -> Left "Arrow.Stream: file dictionary section is not a dictionary batch"+          m' <- applyDictBatch opts (fiSchema idx) m db+          go m' (i + 1)+++-- ============================================================+-- Handle sources (semi-streaming)+-- ============================================================++{- | Read an Arrow IPC /file/ from a 'Handle' one section at a time:+the footer is fetched (a few hundred bytes), the dictionary batches+are read up front, and every 'IS.iterIOStep' seeks to the next+record batch section, reads and decodes just that section, leaving+the rest of the file untouched. Live memory is one section plus the+dictionaries, however large the file. The caller owns the handle and+must not use it concurrently with the iterator.+-}+fileBatchesHandle :: ReadOptions -> Handle -> IO (Either String (IS.IterIO (V.Vector ColumnArray)))+fileBatchesHandle opts h = do+  eidx <- readFileIndexHandle h+  case eidx of+    Left e -> pure (Left e)+    Right idx -> do+      edicts <- readDictBlocksHandle opts idx h+      case edicts of+        Left e -> pure (Left e)+        Right dictMap -> do+          ref <- newIORef 0+          let !blocks = fiRecordBatches idx+              !n = V.length blocks+              step = do+                i <- readIORef ref+                if i >= n+                  then pure (Right Nothing)+                  else do+                    writeIORef ref (i + 1)+                    esec <- readSection h (V.unsafeIndex blocks i)+                    case esec of+                      Left e -> pure (Left e)+                      Right (msg, body) -> case msg of+                        RecordBatch rb ->+                          pure $ case decodeResolved opts (fiSchema idx) dictMap rb body of+                            Left e -> Left e+                            Right cols -> Right (Just cols)+                        _ -> pure (Left "Arrow.Stream: file section is not a record batch")+          pure (Right (IS.iterIOFromAction step))+++{- | Read an Arrow IPC /stream/ from a 'Handle' one frame at a time:+every 'IS.iterIOStep' reads one frame, applies dictionary frames to+the dictionaries in force, and decodes one record batch. Live memory+is one frame. The caller owns the handle and must not use it+concurrently with the iterator.+-}+streamBatchesHandle :: ReadOptions -> Handle -> IO (Either String (IS.IterIO (V.Vector ColumnArray)))+streamBatchesHandle opts h = do+  efirst <- hReadFrame h+  case efirst of+    Left e -> pure (Left e)+    Right Nothing -> pure (Left "Arrow.Stream: empty input, no schema message")+    Right (Just (msg, _)) -> case msg of+      SchemaMessage sch -> do+        ref <- newIORef Map.empty+        pure (Right (IS.iterIOFromAction (streamStep opts h sch ref)))+      _ -> pure (Left "Arrow.Stream: first stream message is not a schema")+++-- | Pull stream frames until one record batch is decoded (or the stream ends).+streamStep :: ReadOptions -> Handle -> Schema -> IORef (Map.Map Int64 ColumnArray) -> IO (Either String (Maybe (V.Vector ColumnArray)))+streamStep opts h sch ref = do+  ef <- hReadFrame h+  case ef of+    Left e -> pure (Left e)+    Right Nothing -> pure (Right Nothing)+    Right (Just (msg, body)) -> case msg of+      RecordBatch rb -> do+        dictMap <- readIORef ref+        case decodeResolved opts sch dictMap rb body of+          Left e -> pure (Left e)+          Right cols -> pure (Right (Just cols))+      DictionaryBatch did isDelta def -> do+        dictMap <- readIORef ref+        let db = DictBatch {dbId = did, dbIsDelta = isDelta, dbData = def, dbBody = body}+        case applyDictBatch opts sch dictMap db of+          Left e -> pure (Left e)+          Right dictMap' -> do+            writeIORef ref dictMap'+            streamStep opts h sch ref+      _ -> pure (Left "Arrow.Stream: unsupported message in stream")+++-- | Read the footer index of the file a handle points at.+readFileIndexHandle :: Handle -> IO (Either String FileIndex)+readFileIndexHandle h = do+  size <- hFileSize h+  if size < 14+    then pure (Left "Arrow.Stream: input too small to be an Arrow file")+    else do+      hSeek h AbsoluteSeek 0+      magic <- BS.hGet h 6+      if magic /= "ARROW1"+        then pure (Left "Arrow.Stream: missing leading ARROW1 magic")+        else do+          hSeek h AbsoluteSeek (size - 10)+          trailer <- BS.hGet h 10+          if BS.length trailer < 10 || BS.drop 4 trailer /= "ARROW1"+            then pure (Left "Arrow.Stream: missing trailing ARROW1 magic")+            else do+              let !footerLen = leI32 trailer+              if footerLen <= 0 || fromIntegral footerLen > size - 16+                then pure (Left "Arrow.Stream: invalid footer length in file trailer")+                else do+                  hSeek h AbsoluteSeek (size - 10 - fromIntegral footerLen)+                  footer <- BS.hGet h (fromIntegral footerLen)+                  if BS.length footer < fromIntegral footerLen+                    then pure (Left "Arrow.Stream: truncated file footer")+                    else pure (parseFileIndex (fromIntegral size) footer)+++-- | Read one file section (message metadata + body) at its block.+readSection :: Handle -> ArrowBlock -> IO (Either String (Message, ByteString))+readSection h (ArrowBlock off metaLen bodyLen)+  | metaLen <= 0 || fromIntegral metaLen > maxMessageMetadataBytes =+      pure (Left "Arrow.Stream: file section metadata length is invalid or too large")+  | bodyLen < 0 =+      pure (Left "Arrow.Stream: file section body length is negative")+  | otherwise = do+      hSeek h AbsoluteSeek (fromIntegral off)+      meta <- BS.hGet h (fromIntegral metaLen)+      if BS.length meta < fromIntegral metaLen+        then pure (Left "Arrow.Stream: truncated file section header")+        else case decodeMessageFrame meta of+          Left e -> pure (Left e)+          Right (msg, declaredBody)+            | declaredBody < 0 || declaredBody > bodyLen ->+                pure (Left "Arrow.Stream: file section body is shorter than its message declares")+            | otherwise -> do+                body <- BS.hGet h (fromIntegral bodyLen)+                if BS.length body < fromIntegral bodyLen+                  then pure (Left "Arrow.Stream: truncated file section body")+                  else pure (Right (msg, body))+++-- | The dictionary batches a file's dictionary sections hold, read one section at a time.+readDictBlocksHandle :: ReadOptions -> FileIndex -> Handle -> IO (Either String (Map.Map Int64 ColumnArray))+readDictBlocksHandle opts idx h = go Map.empty 0+  where+    !n = V.length (fiDictBlocks idx)+    go !m !i+      | i >= n = pure (Right m)+      | otherwise = do+          esec <- readSection h (V.unsafeIndex (fiDictBlocks idx) i)+          case esec of+            Left e -> pure (Left e)+            Right (msg, body) -> case msg of+              DictionaryBatch did isDelta def -> do+                let db = DictBatch {dbId = did, dbIsDelta = isDelta, dbData = def, dbBody = body}+                case applyDictBatch opts (fiSchema idx) m db of+                  Left e -> pure (Left e)+                  Right m' -> go m' (i + 1)+              _ -> pure (Left "Arrow.Stream: file dictionary section is not a dictionary batch")+++{- | Read one encapsulated message frame from a handle:+@Right (Just (message, body))@ for a frame, @Right Nothing@ at a+clean end of stream or the EOS marker, @Left@ on truncation. The+metadata length a frame may declare is capped at+'maxMessageMetadataBytes' so a tiny input cannot name an+out-of-memory buffer.+-}+hReadFrame :: Handle -> IO (Either String (Maybe (Message, ByteString)))+hReadFrame h = do+  pfx <- BS.hGet h 4+  if BS.length pfx == 0+    then pure (Right Nothing)+    else+      if BS.length pfx < 4+        then pure (Left "Arrow.Stream: truncated stream frame header")+        else do+          let !first4 = leU32 pfx+          if first4 == 0+            then pure (Right Nothing)+            else+              if first4 == 0xFFFFFFFF+                then do+                  lp <- BS.hGet h 4+                  if BS.length lp < 4+                    then pure (Left "Arrow.Stream: truncated stream frame after continuation")+                    else do+                      let !mlen = leU32 lp+                      if mlen == 0+                        then pure (Right Nothing)+                        else readMetaBody (pfx <> lp) mlen+                else readMetaBody pfx first4+  where+    readMetaBody prefix mlen+      | mlen > maxMessageMetadataBytes =+          pure (Left "Arrow.Stream: stream message metadata length is too large")+      | otherwise = do+          meta <- BS.hGet h mlen+          if BS.length meta < mlen+            then pure (Left "Arrow.Stream: truncated stream message metadata")+            else case decodeMessageFrame (prefix <> meta) of+              Left e -> pure (Left e)+              Right (msg, bodyLen)+                | bodyLen < 0 -> pure (Left "Arrow.Stream: negative message body length")+                | otherwise -> do+                    let !want = (fromIntegral bodyLen + 7) .&. complement 7 :: Int+                    body <- BS.hGet h want+                    if BS.length body < fromIntegral bodyLen+                      then pure (Left "Arrow.Stream: truncated stream message body")+                      else pure (Right (Just (msg, BS.take (fromIntegral bodyLen) body)))+++{- | Largest message metadata a handle-based reader will buffer. IPC+metadata (schema / dictionary / record batch headers) is small even+for very wide schemas; anything larger is a framing error or a lie,+and buffering it is what a hostile input would want.+-}+maxMessageMetadataBytes :: Int+maxMessageMetadataBytes = 1 `shiftL` 28+++-- | Little-endian unsigned 32-bit int in the first 4 bytes.+leU32 :: ByteString -> Int+leU32 bs = case BS.unpack (BS.take 4 bs) of+  [a, b, c, d] ->+    fromIntegral a+      .|. (fromIntegral b `shiftL` 8)+      .|. (fromIntegral c `shiftL` 16)+      .|. (fromIntegral d `shiftL` 24)+  _ -> 0+++-- | Little-endian signed 32-bit int in the first 4 bytes.+leI32 :: ByteString -> Int32+leI32 bs = case BS.unpack (BS.take 4 bs) of+  [a, b, c, d] ->+    fromIntegral a+      .|. (fromIntegral b `shiftL` 8)+      .|. (fromIntegral c `shiftL` 16)+      .|. (fromIntegral d `shiftL` 24)+  _ -> 0+++-- ============================================================+-- Column projection+-- ============================================================++{- | Like 'streamReaderToList' but only returns the named+columns. Useful when the schema has dozens of columns and the+caller only needs a handful.++The returned 'Schema' is the projected schema (preserving the+order the caller asked for) and each batch is a vector of+'ColumnArray' values in the same order. Names not present in+the source schema produce a 'Left'.+-}+streamReaderProjected+  :: [Text]+  -> StreamReader+  -> Either String (Schema, [V.Vector ColumnArray])+streamReaderProjected names rd0 = do+  let !sourceSchema = srSchema rd0+  idxs <- resolveProjectionIndices names sourceSchema+  let !projSchema = projectSchema idxs sourceSchema+  IS.iterToList (streamReaderProjectedIter names rd0)+    >>= \batches -> Right (projSchema, batches)+++{- | Iterator-shaped variant of 'streamReaderProjected'. Each+'IS.iterStep' yields a vector of column arrays containing+/only/ the requested columns, in the requested order.++The whole record batch is still decoded (and validated), but decoding+aliases the input, so the unrequested columns cost only their+validation.+-}+streamReaderProjectedIter+  :: [Text]+  -> StreamReader+  -> IS.Iter (V.Vector ColumnArray)+streamReaderProjectedIter names rd0 =+  case resolveProjectionIndices names (srSchema rd0) of+    Left e -> IS.iterUnfold () (\_ -> Left e)+    Right idxs -> IS.iterMap (projectColumns idxs) (streamReaderIter rd0)+++resolveProjectionIndices :: [Text] -> Schema -> Either String (V.Vector Int)+resolveProjectionIndices names sch = do+  let !nameToIdx = Map.fromList (V.toList (V.imap (\i f -> (fieldName f, i)) (arrowFields sch)))+  V.fromList <$> traverse (lookupOne nameToIdx) names+  where+    lookupOne m nm = case Map.lookup nm m of+      Just i -> Right i+      Nothing ->+        Left $+          "Arrow.Stream: projected column not present in source schema: "+            ++ show nm+++projectSchema :: V.Vector Int -> Schema -> Schema+projectSchema idxs sch =+  sch+    { arrowFields = V.map (V.unsafeIndex (arrowFields sch)) idxs+    }+++projectColumns :: V.Vector Int -> V.Vector ColumnArray -> V.Vector ColumnArray+projectColumns idxs cols = V.map (V.unsafeIndex cols) idxs
+ src/Arrow/Types.hs view
@@ -0,0 +1,458 @@+{- | Apache Arrow IPC metadata types.++Arrow IPC uses FlatBuffers for metadata. These types mirror the Arrow+Schema.fbs definitions for schema, field, type, and record batch metadata.+-}+module Arrow.Types (+  Schema (..),+  Field (..),+  DictionaryEncoding (..),+  ArrowType (..),+  Endianness (..),+  Precision (..),+  DateUnit (..),+  TimeUnit (..),+  IntervalUnit (..),+  UnionMode (..),+  Message (..),+  RecordBatchDef (..),+  BodyCompressionCodec (..),+  Feature (..),+  FieldNode (..),+  Buffer (..),++  -- * Smart constructors++  --++  {- | These thin wrappers default the optional record fields+  ('arrowMetadata', 'fieldMetadata', 'fieldDictionary',+  'fieldChildren') so test fixtures and one-off encoders+  don't have to spell out every slot. Adding a new optional+  field to 'Schema' or 'Field' won't break callers that+  went through these constructors.+  -}+  defaultSchema,+  defaultField,+  defaultLeafField,++  -- * Decimal schema parameters+  validateDecimalType,++  -- * Schema fingerprint+  schemaFingerprint,+  schemaEquivalent,+) where++import Control.DeepSeq (NFData)+import Crypto.Hash qualified as Hash+import Data.Bits (shiftL, (.|.))+import Data.ByteArray qualified as BA+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Lazy qualified as BL+import Data.Int (Int32, Int64)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Text.Encoding qualified as TE+import Data.Vector (Vector)+import Data.Vector qualified as V+import Data.Vector.Storable qualified as VS+import Data.Word (Word64, Word8)+import Foreign.Storable (Storable (..))+import GHC.Generics (Generic)+import Wireform.Builder qualified as BB+++data Endianness = Little | Big+  deriving stock (Show, Eq, Enum, Bounded, Generic)+  deriving anyclass (NFData)+++data Precision = Half | Single | DoublePrecision+  deriving stock (Show, Eq, Enum, Bounded, Generic)+  deriving anyclass (NFData)+++data DateUnit = DateDay | DateMillisecond+  deriving stock (Show, Eq, Enum, Bounded, Generic)+  deriving anyclass (NFData)+++data TimeUnit = Second | Millisecond | Microsecond | Nanosecond+  deriving stock (Show, Eq, Enum, Bounded, Generic)+  deriving anyclass (NFData)+++data IntervalUnit = YearMonth | DayTime | MonthDayNano+  deriving stock (Show, Eq, Enum, Bounded, Generic)+  deriving anyclass (NFData)+++data UnionMode = Sparse | Dense+  deriving stock (Show, Eq, Enum, Bounded, Generic)+  deriving anyclass (NFData)+++data ArrowType+  = ANull+  | AInt !Int !Bool+  | AFloatingPoint !Precision+  | ABinary+  | AUtf8+  | ABool+  | -- | Arrow format 1.5 signed 32-bit decimal (precision, scale).+    ADecimal32 !Int !Int+  | -- | Arrow format 1.5 signed 64-bit decimal (precision, scale).+    ADecimal64 !Int !Int+  | {- | 128-bit 'Decimal' (precision, scale). Kept as @ADecimal@+    for source compatibility with the original schema; the+    256-bit variant is 'ADecimal256'.+    -}+    ADecimal !Int !Int+  | {- | 256-bit 'Decimal' (precision, scale). Arrow's on-wire+    @Decimal@ type carries a @bitWidth@ field; we keep the two+    widths as distinct 'ArrowType' constructors so the reader+    can dispatch to the right 'ColumnArray' variant.+    -}+    ADecimal256 !Int !Int+  | ADate !DateUnit+  | ATime !TimeUnit !Int+  | ATimestamp !TimeUnit !(Maybe Text)+  | AInterval !IntervalUnit+  | AList+  | AStruct+  | AUnion !UnionMode !(Vector Int32)+  | AFixedSizeBinary !Int+  | AFixedSizeList !Int+  | AMap !Bool+  | ADuration !TimeUnit+  | ALargeBinary+  | ALargeUtf8+  | ALargeList+  | -- The remaining constructors are post-V5 schema additions+    -- (Arrow format version >= 1.4). The metadata writer in+    -- "Arrow.FlatBufferIPC" understands them; a corresponding column+    -- materializer / writer in "Arrow.Column" / "Arrow.Write" is not+    -- yet wired up — the schema slot is enough for round-tripping+    -- type information from external readers.+    ARunEndEncoded+  | ABinaryView+  | AUtf8View+  | AListView+  | ALargeListView+  deriving stock (Show, Eq, Generic)+  deriving anyclass (NFData)+++{- | Check decimal precision and the signed 32-bit scale carried by+Schema.fbs. General Arrow scales may be negative or exceed precision.+Other types are unaffected; this does not validate an entire schema.+-}+validateDecimalType :: ArrowType -> Either String ()+validateDecimalType ty = case ty of+  ADecimal32 precision scale -> decimal 9 precision scale+  ADecimal64 precision scale -> decimal 18 precision scale+  ADecimal precision scale -> decimal 38 precision scale+  ADecimal256 precision scale -> decimal 76 precision scale+  _ -> Right ()+  where+    decimal maximumPrecision precision scale+      | precision < 1 || precision > maximumPrecision =+          Left ("Arrow.Types: decimal precision must be between 1 and " ++ show maximumPrecision)+      | scale < fromIntegral (minBound :: Int32) || scale > fromIntegral (maxBound :: Int32) =+          Left "Arrow.Types: decimal scale exceeds signed 32-bit schema field"+      | otherwise = Right ()+++data Field = Field+  { fieldName :: !Text+  , fieldNullable :: !Bool+  , fieldType :: !ArrowType+  , fieldChildren :: !(Vector Field)+  , fieldDictionary :: !(Maybe DictionaryEncoding)+  {- ^ When non-'Nothing', the column is dictionary-encoded: the+  record batch carries indices of type 'deIndexType', while+  @fieldType@ / @fieldChildren@ describe the dictionary /values/,+  which live in separate 'DictionaryBatch' messages keyed by 'deId'+  (as in the Arrow spec).+  -}+  , fieldMetadata :: !(Vector (Text, Text))+  {- ^ Arrow per-field @custom_metadata@ (Schema.fbs field 6).+  Free-form key/value pairs that travel with the field+  through encode + decode. Empty when the field carries+  no annotation (the common case).+  -}+  }+  deriving stock (Show, Eq, Generic)+  deriving anyclass (NFData)+++{- | Per @format/Schema.fbs@:++@+table DictionaryEncoding {+  id: long;+  indexType: Int;+  isOrdered: bool;+  dictionaryKind: DictionaryKind;+}+@+-}+data DictionaryEncoding = DictionaryEncoding+  { deId :: !Int64+  , deIndexType :: !ArrowType -- always an AInt; default Int32 signed+  , deIsOrdered :: !Bool+  }+  deriving stock (Show, Eq, Generic)+  deriving anyclass (NFData)+++{- | Arrow schema feature flags (Schema.fbs's @Feature@ enum).+'FeatureDictionaryReplacement' (1) advertises multiple full dictionaries+with the same id; 'FeatureCompressedBody' (2) advertises IPC body+compression. 'FeatureUnused' (0) is the FlatBuffers placeholder.+-}+data Feature+  = FeatureUnused+  | FeatureDictionaryReplacement+  | FeatureCompressedBody+  deriving stock (Show, Eq, Enum, Bounded, Generic)+  deriving anyclass (NFData)+++data Schema = Schema+  { arrowFields :: !(Vector Field)+  , arrowEndianness :: !Endianness+  , arrowMetadata :: !(Vector (Text, Text))+  {- ^ Arrow schema-level @custom_metadata@ (Schema.fbs slot+  2 of @table Schema@). Most files set zero or one+  (@pandas@-style) annotation. Empty by default.+  -}+  , arrowFeatures :: !(Vector Feature)+  {- ^ Arrow schema-level @features@ (Schema.fbs slot 3).+  Empty by default. Flags survive IPC stream and file round-trips.+  Producers emitting replacement dictionaries or compressed bodies can+  advertise them with 'FeatureDictionaryReplacement' or+  'FeatureCompressedBody'. Unknown feature values are rejected on read.+  -}+  }+  deriving stock (Show, Eq, Generic)+  deriving anyclass (NFData)+++{- | One array's length and null count (Message.fbs @struct FieldNode@,+16 bytes). 'Storable' with the wire layout, so a record batch's nodes+are one 'VS.Vector'.+-}+data FieldNode = FieldNode+  { fnLength :: !Int64+  , fnNullCount :: !Int64+  }+  deriving stock (Show, Eq, Generic)+  deriving anyclass (NFData)+++instance Storable FieldNode where+  sizeOf _ = 16+  alignment _ = 8+  peek p = FieldNode <$> peekByteOff p 0 <*> peekByteOff p 8+  poke p (FieldNode l n) = pokeByteOff p 0 l >> pokeByteOff p 8 n+  {-# INLINE peek #-}+  {-# INLINE poke #-}+++{- | One body buffer's offset and length (Message.fbs @struct Buffer@,+16 bytes). 'Storable' with the wire layout.+-}+data Buffer = Buffer+  { bufOffset :: !Int64+  , bufLength :: !Int64+  }+  deriving stock (Show, Eq, Generic)+  deriving anyclass (NFData)+++instance Storable Buffer where+  sizeOf _ = 16+  alignment _ = 8+  peek p = Buffer <$> peekByteOff p 0 <*> peekByteOff p 8+  poke p (Buffer o l) = pokeByteOff p 0 o >> pokeByteOff p 8 l+  {-# INLINE peek #-}+  {-# INLINE poke #-}+++data RecordBatchDef = RecordBatchDef+  { rbLength :: !Int64+  , rbNodes :: !(VS.Vector FieldNode)+  -- ^ Field nodes in schema pre-order (one memcpy from the message).+  , rbBuffers :: !(VS.Vector Buffer)+  -- ^ Body buffers in spec layout order.+  , rbVariadicBufferCounts :: !(Vector Int64)+  {- ^ Per Arrow @format/Message.fbs@: when the schema contains+  @Utf8View@ or @BinaryView@ fields each such field has a+  variable number of additional data buffers for out-of-line+  string payloads. The vector lists, in pre-order schema+  traversal order, the number of variadic data buffers per+  view column. Empty for schemas without view types.+  -}+  , rbBodyCompression :: !(Maybe BodyCompressionCodec)+  {- ^ Per Arrow @format/Message.fbs@'s 'BodyCompression' table.+  When 'Just', each buffer in the body is wrapped in an+  @<i64 uncompressedLength><compressed bytes>@ envelope using+  the named codec; readers see the original layout once+  decompressed. 'Nothing' = uncompressed buffers (the wire+  default).+  -}+  }+  deriving stock (Show, Eq, Generic)+  deriving anyclass (NFData)+++{- | Body-compression codecs supported by Arrow IPC's+'BodyCompression' table (per @format/Message.fbs@).+-}+data BodyCompressionCodec+  = {- | LZ4 frame format (@CompressionType = LZ4_FRAME = 0@).+    Most widely supported; the default for pyarrow when+    compression is enabled.+    -}+    LZ4Frame+  | -- | Zstandard (@CompressionType = ZSTD = 1@).+    BodyZstd+  deriving stock (Show, Eq, Generic)+  deriving anyclass (NFData)+++{- | One Arrow IPC message header: the schema, a dictionary batch+(dictionary id, is-delta flag, the batch holding the dictionary+values), or a record batch. Bodies travel separately.+-}+data Message+  = SchemaMessage !Schema+  | DictionaryBatch !Int64 !Bool !RecordBatchDef+  | RecordBatch !RecordBatchDef+  deriving stock (Show, Eq, Generic)+  deriving anyclass (NFData)+++{- | Build a 'Schema' with little-endian + no custom_metadata.+Equivalent to @Schema fields Little V.empty@.+-}+defaultSchema :: Vector Field -> Schema+defaultSchema fs =+  Schema+    { arrowFields = fs+    , arrowEndianness = Little+    , arrowMetadata = V.empty+    , arrowFeatures = V.empty+    }+++{- | Build a 'Field' with no children, no dictionary, and no+custom_metadata.+-}+defaultLeafField :: Text -> Bool -> ArrowType -> Field+defaultLeafField name nullable ty =+  Field+    { fieldName = name+    , fieldNullable = nullable+    , fieldType = ty+    , fieldChildren = V.empty+    , fieldDictionary = Nothing+    , fieldMetadata = V.empty+    }+++{- | Build a 'Field' with explicit children but no dictionary+encoding and no custom_metadata.+-}+defaultField :: Text -> Bool -> ArrowType -> Vector Field -> Field+defaultField name nullable ty children =+  Field+    { fieldName = name+    , fieldNullable = nullable+    , fieldType = ty+    , fieldChildren = children+    , fieldDictionary = Nothing+    , fieldMetadata = V.empty+    }+++{- | A 64-bit deterministic fingerprint of an Arrow 'Schema',+computed as the high 64 bits of SHA-256 of a canonical+byte-encoding of the schema's structural skeleton (Avro-+style \"CRC-64-AVRO\"-shaped construction, but with SHA-256+as the underlying hash for collision resistance).++Two schemas that compare 'schemaEquivalent' produce the same+fingerprint; the construction is byte-deterministic across+GHC versions, processes, and architectures (no+'Data.Hashable' randomisation), so the fingerprint is safe+to persist as a cache key or schema-version identifier.++Equivalence ignores 'arrowMetadata' / 'arrowFeatures' and+per-field 'fieldMetadata' — those are annotations, not+structure. Endianness, field order, types, nullability, and+the full nested-children tree all participate.+-}+schemaFingerprint :: Schema -> Word64+schemaFingerprint sch =+  let !canon = BL.toStrict (BB.toLazyByteString (buildSchema sch))+      !digest = Hash.hash canon :: Hash.Digest Hash.SHA256+      !bytes = BA.convert digest :: ByteString+  in word64BE (BS.take 8 bytes)+  where+    word64BE bs =+      let r i = fromIntegral (BS.index bs i) :: Word64+      in (r 0 `shiftL` 56)+           .|. (r 1 `shiftL` 48)+           .|. (r 2 `shiftL` 40)+           .|. (r 3 `shiftL` 32)+           .|. (r 4 `shiftL` 24)+           .|. (r 5 `shiftL` 16)+           .|. (r 6 `shiftL` 8)+           .|. r 7++    buildSchema s =+      BB.word8 (endianTag (arrowEndianness s))+        <> buildVecField (arrowFields s)++    buildVecField fs =+      BB.int32BE (fromIntegral (V.length fs))+        <> foldMap buildField (V.toList fs)++    buildField f =+      buildText (fieldName f)+        <> BB.word8 (if fieldNullable f then 1 else 0)+        <> buildText (T.pack (show (fieldType f)))+        <> buildVecField (fieldChildren f)++    buildText t =+      let bs = TE.encodeUtf8 t+      in BB.int32BE (fromIntegral (BS.length bs))+           <> BB.byteString bs++    endianTag :: Endianness -> Word8+    endianTag Little = 0+    endianTag Big = 1+++{- | Structural equivalence, mirroring 'schemaFingerprint'.+Annotation fields ('arrowMetadata', 'arrowFeatures',+per-field 'fieldMetadata') are ignored.+-}+schemaEquivalent :: Schema -> Schema -> Bool+schemaEquivalent a b =+  arrowEndianness a == arrowEndianness b+    && V.length (arrowFields a) == V.length (arrowFields b)+    && and (V.zipWith fieldEquivalent (arrowFields a) (arrowFields b))+++fieldEquivalent :: Field -> Field -> Bool+fieldEquivalent a b =+  fieldName a == fieldName b+    && fieldNullable a == fieldNullable b+    && fieldType a == fieldType b+    && V.length (fieldChildren a) == V.length (fieldChildren b)+    && and (V.zipWith fieldEquivalent (fieldChildren a) (fieldChildren b))
+ src/Arrow/Vector.hs view
@@ -0,0 +1,60 @@+{- | Vectors in Arrow layout behind the @vector@ package API.++@'Vector' a@ is a data family (like @Data.Vector.Unboxed.Vector@) with+an instance per element type, each storing its elements the way an+Arrow column does:++* fixed-width types ('Data.Int.Int8' to 'Data.Int.Int64',+  'Data.Word.Word8' to 'Data.Word.Word64', 'Float', 'Double', 'Int',+  'Word' and the Arrow element types 'Arrow.Column.Float16',+  'Arrow.Column.IntervalDayTime', 'Arrow.Column.IntervalMonthDayNano',+  'Arrow.Column.Decimal128', 'Arrow.Column.Decimal256'): a storable+  vector;+* 'Bool': a bitmap;+* @'Maybe' a@: a validity bitmap plus a @'Vector' a@ (a @Vector (Maybe+  Int64)@ costs 8 bytes and one bit per row, with no heap object per+  row);+* 'Data.Text.Text', 'Data.ByteString.ByteString' and @'Vector' a@+  (list rows): start and end offsets into one store, so a row is an+  O(1) slice.++Instances nest: @Vector (Maybe (Vector (Maybe Int32)))@ is a list+column. Use the generic API with element type @Maybe a@ as usual:++> import Data.Vector.Generic qualified as G+>+> total :: Vector (Maybe Int64) -> Int64+> total = G.foldl' (\acc m -> maybe acc (+ acc) m) 0++Fused consumers like this one compile to a loop over the bitmap and the+values that allocates nothing per row. 'Data.Vector.Generic.convert' turns a vector into+a boxed "Data.Vector" (or any other generic vector) and back.++The conversions in "Arrow.Column" ('Arrow.Column.toMaybeVector' and+friends) alias the column's buffers, so the result keeps the decoded+input alive like the column does; 'Data.Vector.Generic.force' copies into fresh buffers.++Mutable 'Data.Text.Text', 'Data.ByteString.ByteString' and list+vectors keep their rows in an append-only store shared by every slice:+writing a row appends its bytes (O(row size)), and 'Data.Vector.Generic.freeze' /+'Data.Vector.Generic.unsafeFreeze' keep every byte ever written. 'Data.Vector.Generic.force' (or+'Data.Vector.Generic.freeze' of a fresh vector) copies only what the rows reference.+-}+module Arrow.Vector (+  Vector,+  MVector,+  IOVector,+  STVector,+  Element,+  FixedWidth (..),+  maybeValues,+) where++import Arrow.Vector.Internal+import Control.Monad.ST (RealWorld)+++type IOVector = MVector RealWorld+++type STVector s = MVector s
+ src/Arrow/Vector/Internal.hs view
@@ -0,0 +1,980 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UnboxedTuples #-}++{- | Representation of "Arrow.Vector": the data family instances with+their constructors, the classes, and the bit and var-length helpers the+column conversions in "Arrow.Column" build on.++The constructors carry invariants the instances rely on (bitmaps cover+their bit range, every var-length row range lies inside its store).+Code outside this package should use "Arrow.Vector" and the generic+vector API instead.++Layouts:++* fixed-width element types: a storable vector ('VS.Vector'), so a+  decoded column converts by aliasing its value buffer;+* 'Bool': an LSB-first bitmap with a bit offset;+* @'Maybe' a@: a validity bitmap with a bit offset plus a @'Vector' a@+  whose null slots hold unspecified values;+* 'Text', 'ByteString' and @'Vector' a@ elements: per-row start and end+  offsets into one store (a text array, a byte buffer or a child+  vector), so indexing is an O(1) slice.+-}+module Arrow.Vector.Internal (+  -- * Families+  Vector (..),+  MVector (..),+  Element,+  FixedWidth (..),+  maybeValues,++  -- * Var-length rows+  VarElem (..),+  VarVec (..),+  MVarVec (..),++  -- * Bits+  bitBytes,+  indexBit,+  allSetBits,+  bytesToBits,+  bitsToBytes,+) where++import Arrow.Column.Buffer (Decimal128, Decimal256, Float16, IntervalDayTime, IntervalMonthDayNano)+import Columnar.SIMD qualified as K+import Control.DeepSeq (NFData (..))+import Control.Monad.ST (ST)+import Control.Monad.ST.Unsafe (unsafeIOToST)+import Data.Bits (complement, unsafeShiftL, unsafeShiftR, (.&.), (.|.))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Internal qualified as BSI+import Data.Int (Int16, Int32, Int64, Int8)+import Data.Primitive.ByteArray (+  ByteArray (..),+  MutableByteArray,+  copyByteArray,+  copyMutableByteArray,+  getSizeofMutableByteArray,+  newByteArray,+  sizeofByteArray,+  unsafeFreezeByteArray,+  unsafeThawByteArray,+ )+import Data.Primitive.MutVar (MutVar, newMutVar, readMutVar, writeMutVar)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Text.Array qualified as TA+import Data.Text.Internal qualified as TI+import Data.Vector.Fusion.Util (Box (..))+import Data.Vector.Generic qualified as G+import Data.Vector.Generic.Mutable qualified as GM+import Data.Vector.Storable qualified as VS+import Data.Vector.Storable.Mutable qualified as VSM+import Data.Word (Word16, Word32, Word64, Word8)+import Foreign.ForeignPtr (ForeignPtr, plusForeignPtr)+import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)+import Foreign.Marshal.Utils (copyBytes)+import Foreign.Ptr (minusPtr, nullPtr, plusPtr)+import Foreign.Storable (Storable)+import GHC.ForeignPtr (mallocPlainForeignPtrBytes, unsafeWithForeignPtr)+++-- ============================================================+-- Families and classes+-- ============================================================++-- | Immutable vectors in Arrow layout, indexed by element type.+data family Vector a+++-- | Mutable vectors in Arrow layout.+data family MVector s a+++type instance G.Mutable Vector = MVector+++{- | Element types with a "Arrow.Vector" layout (the analogue of+@Data.Vector.Unboxed.Unbox@).+-}+class (G.Vector Vector a, GM.MVector MVector a) => Element a+++{- | Fixed-width element types: the vector is a storable vector, and+both conversions are O(1).+-}+class (Element a, Storable a) => FixedWidth a where+  -- | O(1): view a storable vector as a 'Vector'.+  fromStorable :: VS.Vector a -> Vector a++  -- | O(1): the storable vector behind a 'Vector'.+  storableVector :: Vector a -> VS.Vector a+++instance (Show a, Element a) => Show (Vector a) where+  showsPrec = G.showsPrec+++instance (Eq a, Element a) => Eq (Vector a) where+  (==) = G.eq+  {-# INLINE (==) #-}+++instance (Ord a, Element a) => Ord (Vector a) where+  compare = G.cmp+  {-# INLINE compare #-}+++-- | Every layout is strict all the way down, so WHNF is normal form.+instance NFData (Vector a) where+  rnf v = v `seq` ()+++instance Element a => Semigroup (Vector a) where+  (<>) = (G.++)+  {-# INLINE (<>) #-}+++instance Element a => Monoid (Vector a) where+  mempty = G.empty+  {-# INLINE mempty #-}+++-- ============================================================+-- Fixed-width element types+-- ============================================================++#define FIXED(ty,con,mcon) \+newtype instance MVector s (ty) = mcon (VSM.MVector s (ty)); \+newtype instance Vector (ty) = con (VS.Vector (ty)); \+instance GM.MVector MVector (ty) where { \+  {-# INLINE basicLength #-} \+; {-# INLINE basicUnsafeSlice #-} \+; {-# INLINE basicOverlaps #-} \+; {-# INLINE basicUnsafeNew #-} \+; {-# INLINE basicInitialize #-} \+; {-# INLINE basicUnsafeReplicate #-} \+; {-# INLINE basicUnsafeRead #-} \+; {-# INLINE basicUnsafeWrite #-} \+; {-# INLINE basicClear #-} \+; {-# INLINE basicSet #-} \+; {-# INLINE basicUnsafeCopy #-} \+; {-# INLINE basicUnsafeMove #-} \+; {-# INLINE basicUnsafeGrow #-} \+; basicLength (mcon v) = GM.basicLength v \+; basicUnsafeSlice i n (mcon v) = mcon (GM.basicUnsafeSlice i n v) \+; basicOverlaps (mcon a) (mcon b) = GM.basicOverlaps a b \+; basicUnsafeNew n = mcon <$> GM.basicUnsafeNew n \+; basicInitialize (mcon v) = GM.basicInitialize v \+; basicUnsafeReplicate n x = mcon <$> GM.basicUnsafeReplicate n x \+; basicUnsafeRead (mcon v) i = GM.basicUnsafeRead v i \+; basicUnsafeWrite (mcon v) i x = GM.basicUnsafeWrite v i x \+; basicClear (mcon v) = GM.basicClear v \+; basicSet (mcon v) x = GM.basicSet v x \+; basicUnsafeCopy (mcon a) (mcon b) = GM.basicUnsafeCopy a b \+; basicUnsafeMove (mcon a) (mcon b) = GM.basicUnsafeMove a b \+; basicUnsafeGrow (mcon v) n = mcon <$> GM.basicUnsafeGrow v n }; \+instance G.Vector Vector (ty) where { \+  {-# INLINE basicUnsafeFreeze #-} \+; {-# INLINE basicUnsafeThaw #-} \+; {-# INLINE basicLength #-} \+; {-# INLINE basicUnsafeSlice #-} \+; {-# INLINE basicUnsafeIndexM #-} \+; {-# INLINE basicUnsafeCopy #-} \+; {-# INLINE elemseq #-} \+; basicUnsafeFreeze (mcon v) = con <$> G.basicUnsafeFreeze v \+; basicUnsafeThaw (con v) = mcon <$> G.basicUnsafeThaw v \+; basicLength (con v) = G.basicLength v \+; basicUnsafeSlice i n (con v) = con (G.basicUnsafeSlice i n v) \+; basicUnsafeIndexM (con v) i = G.basicUnsafeIndexM v i \+; basicUnsafeCopy (mcon m) (con v) = G.basicUnsafeCopy m v \+; elemseq _ = seq }; \+instance Element (ty); \+instance FixedWidth (ty) where { \+  {-# INLINE fromStorable #-} \+; {-# INLINE storableVector #-} \+; fromStorable = con \+; storableVector (con v) = v }++FIXED(Int8, V_Int8, MV_Int8)+FIXED(Int16, V_Int16, MV_Int16)+FIXED(Int32, V_Int32, MV_Int32)+FIXED(Int64, V_Int64, MV_Int64)+FIXED(Int, V_Int, MV_Int)+FIXED(Word8, V_Word8, MV_Word8)+FIXED(Word16, V_Word16, MV_Word16)+FIXED(Word32, V_Word32, MV_Word32)+FIXED(Word64, V_Word64, MV_Word64)+FIXED(Word, V_Word, MV_Word)+FIXED(Float, V_Float, MV_Float)+FIXED(Double, V_Double, MV_Double)+FIXED(Float16, V_Float16, MV_Float16)+FIXED(IntervalDayTime, V_IntervalDayTime, MV_IntervalDayTime)+FIXED(IntervalMonthDayNano, V_IntervalMonthDayNano, MV_IntervalMonthDayNano)+FIXED(Decimal128, V_Decimal128, MV_Decimal128)+FIXED(Decimal256, V_Decimal256, MV_Decimal256)+++-- ============================================================+-- Bits+-- ============================================================++-- | Bytes holding @n@ bits that start at bit @o@ of the first byte.+bitBytes :: Int -> Int -> Int+bitBytes o n = if n <= 0 then 0 else (o + n + 7) `unsafeShiftR` 3+{-# INLINE bitBytes #-}+++-- | Bit @k@ (LSB first).+indexBit :: VS.Vector Word8 -> Int -> Bool+indexBit bs k = (VS.unsafeIndex bs (k `unsafeShiftR` 3) `unsafeShiftR` (k .&. 7)) .&. 1 /= 0+{-# INLINE indexBit #-}+++-- | A fresh bitmap of @n@ set bits.+allSetBits :: Int -> VS.Vector Word8+allSetBits n = VS.replicate (bitBytes 0 n) 0xFF+++-- | O(1): bytes as a bitmap buffer.+bytesToBits :: ByteString -> VS.Vector Word8+bytesToBits (BSI.BS fp len) = VS.unsafeFromForeignPtr0 fp len+{-# INLINE bytesToBits #-}+++-- | O(1): a bitmap buffer as bytes.+bitsToBytes :: VS.Vector Word8 -> ByteString+bitsToBytes v = let !(fp, len) = VS.unsafeToForeignPtr0 v in BSI.BS fp len+{-# INLINE bitsToBytes #-}+++readBitM :: VSM.MVector s Word8 -> Int -> ST s Bool+readBitM bs k = do+  w <- VSM.unsafeRead bs (k `unsafeShiftR` 3)+  pure ((w `unsafeShiftR` (k .&. 7)) .&. 1 /= 0)+{-# INLINE readBitM #-}+++writeBitM :: VSM.MVector s Word8 -> Int -> Bool -> ST s ()+writeBitM bs k b = do+  let !j = k `unsafeShiftR` 3+      !m = 1 `unsafeShiftL` (k .&. 7) :: Word8+  w <- VSM.unsafeRead bs j+  VSM.unsafeWrite bs j (if b then w .|. m else w .&. complement m)+{-# INLINE writeBitM #-}+++-- | Set bits @[off, off + n)@ to @b@, leaving the neighbouring bits alone.+fillBitsM :: VSM.MVector s Word8 -> Int -> Int -> Bool -> ST s ()+fillBitsM bs off n b+  | n <= 0 = pure ()+  | j0 == j1 = patch j0 (headMask .&. tailMask)+  | otherwise = do+      patch j0 headMask+      VSM.set (VSM.unsafeSlice (j0 + 1) (j1 - j0 - 1) bs) (if b then 0xFF else 0)+      patch j1 tailMask+  where+    !end = off + n+    !j0 = off `unsafeShiftR` 3+    !j1 = (end - 1) `unsafeShiftR` 3+    headMask = 0xFF `unsafeShiftL` (off .&. 7) :: Word8+    tailMask = 0xFF `unsafeShiftR` (7 - ((end - 1) .&. 7)) :: Word8+    patch j m = do+      w <- VSM.unsafeRead bs j+      VSM.unsafeWrite bs j (if b then w .|. m else w .&. complement m)+++-- | Copy @n@ bits between non-overlapping ranges (destination neighbours kept).+copyBitsM :: VSM.MVector s Word8 -> Int -> VSM.MVector s Word8 -> Int -> Int -> ST s ()+copyBitsM (VSM.MVector _ fd) doff (VSM.MVector _ fs) soff n+  | n <= 0 = pure ()+  | otherwise =+      unsafeIOToST $ unsafeWithForeignPtr fd $ \pd -> unsafeWithForeignPtr fs $ \ps -> K.copyBits pd doff ps soff n+{-# INLINE copyBitsM #-}+++-- | 'copyBitsM' from an immutable bitmap.+copyBitsFrom :: VSM.MVector s Word8 -> Int -> VS.Vector Word8 -> Int -> Int -> ST s ()+copyBitsFrom (VSM.MVector _ fd) doff src soff n+  | n <= 0 = pure ()+  | otherwise =+      unsafeIOToST $ unsafeWithForeignPtr fd $ \pd -> VS.unsafeWith src $ \ps -> K.copyBits pd doff ps soff n+{-# INLINE copyBitsFrom #-}+++-- | 'copyBitsM' for ranges that may overlap (through a scratch buffer when they do).+moveBitsM :: VSM.MVector s Word8 -> Int -> VSM.MVector s Word8 -> Int -> Int -> ST s ()+moveBitsM dst doff src soff n+  | bitsOverlap dst doff n src soff n = do+      tmp <- VSM.unsafeNew (bitBytes 0 n)+      copyBitsM tmp 0 src soff n+      copyBitsM dst doff tmp 0 n+  | otherwise = copyBitsM dst doff src soff n+++-- | Whether two bit ranges share a bit (compared by absolute bit address).+bitsOverlap :: VSM.MVector s Word8 -> Int -> Int -> VSM.MVector s Word8 -> Int -> Int -> Bool+bitsOverlap (VSM.MVector _ fa) oa na (VSM.MVector _ fb) ob nb =+  na > 0 && nb > 0 && a < b + nb && b < a + na+  where+    !a = addrOf fa * 8 + oa+    !b = addrOf fb * 8 + ob+    addrOf fp = unsafeForeignPtrToPtr fp `minusPtr` nullPtr+++-- ============================================================+-- Bool+-- ============================================================++-- | Bit offset (below 8 after slicing), length, bytes.+data instance MVector s Bool = MV_Bool {-# UNPACK #-} !Int {-# UNPACK #-} !Int {-# UNPACK #-} !(VSM.MVector s Word8)+++-- | Bit offset (below 8 after slicing), length, bytes.+data instance Vector Bool = V_Bool {-# UNPACK #-} !Int {-# UNPACK #-} !Int {-# UNPACK #-} !(VS.Vector Word8)+++instance GM.MVector MVector Bool where+  {-# INLINE basicLength #-}+  basicLength (MV_Bool _ n _) = n+  {-# INLINE basicUnsafeSlice #-}+  basicUnsafeSlice i n (MV_Bool o _ bs) =+    let !k = o + i+        !o' = k .&. 7+    in MV_Bool o' n (VSM.unsafeSlice (k `unsafeShiftR` 3) (bitBytes o' n) bs)+  {-# INLINE basicOverlaps #-}+  basicOverlaps (MV_Bool oa na a) (MV_Bool ob nb b) = bitsOverlap a oa na b ob nb+  {-# INLINE basicUnsafeNew #-}+  basicUnsafeNew n = MV_Bool 0 n <$> VSM.unsafeNew (bitBytes 0 n)+  {-# INLINE basicInitialize #-}+  basicInitialize (MV_Bool o n bs) = fillBitsM bs o n False+  {-# INLINE basicUnsafeRead #-}+  basicUnsafeRead (MV_Bool o _ bs) i = readBitM bs (o + i)+  {-# INLINE basicUnsafeWrite #-}+  basicUnsafeWrite (MV_Bool o _ bs) i b = writeBitM bs (o + i) b+  {-# INLINE basicSet #-}+  basicSet (MV_Bool o n bs) b = fillBitsM bs o n b+  {-# INLINE basicUnsafeCopy #-}+  basicUnsafeCopy (MV_Bool od n dst) (MV_Bool os _ src) = copyBitsM dst od src os n+  {-# INLINE basicUnsafeMove #-}+  basicUnsafeMove (MV_Bool od n dst) (MV_Bool os _ src) = moveBitsM dst od src os n+  {-# INLINE basicUnsafeGrow #-}+  basicUnsafeGrow (MV_Bool o n bs) by = do+    nb <- VSM.unsafeNew (bitBytes 0 (n + by))+    copyBitsM nb 0 bs o n+    pure (MV_Bool 0 (n + by) nb)+++instance G.Vector Vector Bool where+  {-# INLINE basicUnsafeFreeze #-}+  basicUnsafeFreeze (MV_Bool o n bs) = V_Bool o n <$> VS.unsafeFreeze bs+  {-# INLINE basicUnsafeThaw #-}+  basicUnsafeThaw (V_Bool o n bs) = MV_Bool o n <$> VS.unsafeThaw bs+  {-# INLINE basicLength #-}+  basicLength (V_Bool _ n _) = n+  {-# INLINE basicUnsafeSlice #-}+  basicUnsafeSlice i n (V_Bool o _ bs) =+    let !k = o + i+        !o' = k .&. 7+    in V_Bool o' n (VS.unsafeSlice (k `unsafeShiftR` 3) (bitBytes o' n) bs)+  {-# INLINE basicUnsafeIndexM #-}+  basicUnsafeIndexM (V_Bool o _ bs) i = let !b = indexBit bs (o + i) in Box b+  {-# INLINE basicUnsafeCopy #-}+  basicUnsafeCopy (MV_Bool od n dst) (V_Bool os _ src) = copyBitsFrom dst od src os n+++instance Element Bool+++-- ============================================================+-- Maybe+-- ============================================================++-- | Validity bit offset (below 8 after slicing), validity bytes (bit set = 'Just'), values.+data instance MVector s (Maybe a) = MV_Maybe {-# UNPACK #-} !Int {-# UNPACK #-} !(VSM.MVector s Word8) !(MVector s a)+++-- | Validity bit offset (below 8 after slicing), validity bytes (bit set = 'Just'), values.+data instance Vector (Maybe a) = V_Maybe {-# UNPACK #-} !Int {-# UNPACK #-} !(VS.Vector Word8) !(Vector a)+++instance Element a => GM.MVector MVector (Maybe a) where+  {-# INLINE basicLength #-}+  basicLength (MV_Maybe _ _ xs) = GM.basicLength xs+  {-# INLINE basicUnsafeSlice #-}+  basicUnsafeSlice i n (MV_Maybe o bs xs) =+    let !k = o + i+        !o' = k .&. 7+    in MV_Maybe o' (VSM.unsafeSlice (k `unsafeShiftR` 3) (bitBytes o' n) bs) (GM.basicUnsafeSlice i n xs)+  {-# INLINE basicOverlaps #-}+  basicOverlaps (MV_Maybe _ _ a) (MV_Maybe _ _ b) = GM.basicOverlaps a b+  {-# INLINE basicUnsafeNew #-}+  basicUnsafeNew n = MV_Maybe 0 <$> VSM.unsafeNew (bitBytes 0 n) <*> GM.basicUnsafeNew n+  {-# INLINE basicInitialize #-}+  basicInitialize (MV_Maybe o bs xs) = do+    fillBitsM bs o (GM.basicLength xs) False+    GM.basicInitialize xs+  {-# INLINE basicUnsafeRead #-}+  basicUnsafeRead (MV_Maybe o bs xs) i = do+    ok <- readBitM bs (o + i)+    if ok then Just <$> GM.basicUnsafeRead xs i else pure Nothing+  {-# INLINE basicUnsafeWrite #-}+  basicUnsafeWrite (MV_Maybe o bs xs) i m = case m of+    Nothing -> writeBitM bs (o + i) False+    Just x -> do+      writeBitM bs (o + i) True+      GM.basicUnsafeWrite xs i x+  {-# INLINE basicClear #-}+  basicClear (MV_Maybe _ _ xs) = GM.basicClear xs+  {-# INLINE basicSet #-}+  basicSet (MV_Maybe o bs xs) m = case m of+    Nothing -> fillBitsM bs o (GM.basicLength xs) False+    Just x -> do+      fillBitsM bs o (GM.basicLength xs) True+      GM.basicSet xs x+  {-# INLINE basicUnsafeCopy #-}+  basicUnsafeCopy (MV_Maybe od dbs dxs) (MV_Maybe os sbs sxs) = do+    copyBitsM dbs od sbs os (GM.basicLength sxs)+    GM.basicUnsafeCopy dxs sxs+  {-# INLINE basicUnsafeMove #-}+  basicUnsafeMove (MV_Maybe od dbs dxs) (MV_Maybe os sbs sxs) = do+    moveBitsM dbs od sbs os (GM.basicLength sxs)+    GM.basicUnsafeMove dxs sxs+  {-# INLINE basicUnsafeGrow #-}+  basicUnsafeGrow (MV_Maybe o bs xs) by = do+    let !n = GM.basicLength xs+    nb <- VSM.unsafeNew (bitBytes 0 (n + by))+    copyBitsM nb 0 bs o n+    MV_Maybe 0 nb <$> GM.basicUnsafeGrow xs by+++instance Element a => G.Vector Vector (Maybe a) where+  {-# INLINE basicUnsafeFreeze #-}+  basicUnsafeFreeze (MV_Maybe o bs xs) = V_Maybe o <$> VS.unsafeFreeze bs <*> G.basicUnsafeFreeze xs+  {-# INLINE basicUnsafeThaw #-}+  basicUnsafeThaw (V_Maybe o bs xs) = MV_Maybe o <$> VS.unsafeThaw bs <*> G.basicUnsafeThaw xs+  {-# INLINE basicLength #-}+  basicLength (V_Maybe _ _ xs) = G.basicLength xs+  {-# INLINE basicUnsafeSlice #-}+  basicUnsafeSlice i n (V_Maybe o bs xs) =+    let !k = o + i+        !o' = k .&. 7+    in V_Maybe o' (VS.unsafeSlice (k `unsafeShiftR` 3) (bitBytes o' n) bs) (G.basicUnsafeSlice i n xs)+  {-# INLINE basicUnsafeIndexM #-}+  basicUnsafeIndexM (V_Maybe o bs xs) i+    | indexBit bs (o + i) = case G.basicUnsafeIndexM xs i of Box x -> Box (Just x)+    | otherwise = Box Nothing+  {-# INLINE basicUnsafeCopy #-}+  basicUnsafeCopy (MV_Maybe od dbs dxs) (V_Maybe os sbs sxs) = do+    copyBitsFrom dbs od sbs os (G.basicLength sxs)+    G.basicUnsafeCopy dxs sxs+  {-# INLINE elemseq #-}+  elemseq _ m y = case m of+    Nothing -> y+    Just x -> G.elemseq (undefined :: Vector a) x y+++instance Element a => Element (Maybe a)+++{- | O(1): the values behind a vector of 'Maybe's. Slots whose row is+'Nothing' hold unspecified values.+-}+maybeValues :: Vector (Maybe a) -> Vector a+maybeValues (V_Maybe _ _ xs) = xs+{-# INLINE maybeValues #-}+++-- ============================================================+-- Var-length rows: Text, ByteString, Vector a+-- ============================================================++{- | Immutable var-length rows: row @i@ is the slice+@[starts ! i, ends ! i)@ of the store. Rows may share or skip bytes of+the store; every range lies inside it.+-}+data VarVec a = VarVec {-# UNPACK #-} !(VS.Vector Int) {-# UNPACK #-} !(VS.Vector Int) !(Store a)+++{- | Mutable var-length rows: per-row start and end offsets into a+growable, append-only store shared by every slice of the vector.+Writing a row appends its bytes (or child elements) and points the row+at them; bytes already written are never overwritten, so a row read+from the vector stays valid after later writes and growth.+-}+data MVarVec s a = MVarVec {-# UNPACK #-} !(VSM.MVector s Int) {-# UNPACK #-} !(VSM.MVector s Int) {-# UNPACK #-} !(MutVar s (Grow s a))+++-- | Element types stored as slices of one store.+class VarElem a where+  -- | The immutable store.+  type Store a++  -- | The growable store and how much of it is used.+  data Grow s a++  -- | Row @[s, e)@ of a store, O(1).+  storeRow :: Store a -> Int -> Int -> a++  -- | Store units (bytes or child elements) a row occupies.+  rowLength :: a -> Int++  -- | Append a row; it lands at 'growUsed' of the old store.+  appendRow :: Grow s a -> a -> ST s (Grow s a)++  -- | Append @[s, e)@ of an immutable store.+  appendStore :: Grow s a -> Store a -> Int -> Int -> ST s (Grow s a)++  -- | Append @[s, e)@ of another growable store.+  appendGrow :: Grow s a -> Grow s a -> Int -> Int -> ST s (Grow s a)++  growUsed :: Grow s a -> Int++  -- | Row @[s, e)@ of a growable store, O(1).+  growRow :: Grow s a -> Int -> Int -> ST s a++  growEmpty :: ST s (Grow s a)++  -- | O(1); keeps the whole buffer, unused capacity included.+  growFreeze :: Grow s a -> ST s (Store a)++  -- | O(1); the store counts as full, so the first append copies it.+  growThaw :: Store a -> ST s (Grow s a)++  toVar :: Vector a -> VarVec a+  fromVar :: VarVec a -> Vector a+++varLength :: VarVec a -> Int+varLength (VarVec st _ _) = VS.length st+{-# INLINE varLength #-}+++varSlice :: Int -> Int -> VarVec a -> VarVec a+varSlice i n (VarVec st en d) = VarVec (VS.unsafeSlice i n st) (VS.unsafeSlice i n en) d+{-# INLINE varSlice #-}+++varIndex :: VarElem a => VarVec a -> Int -> Box a+varIndex (VarVec st en d) i = let !x = storeRow d (VS.unsafeIndex st i) (VS.unsafeIndex en i) in Box x+{-# INLINE varIndex #-}+++varFreeze :: VarElem a => MVarVec s a -> ST s (VarVec a)+varFreeze (MVarVec st en ref) = do+  st' <- VS.unsafeFreeze st+  en' <- VS.unsafeFreeze en+  d <- readMutVar ref >>= growFreeze+  pure (VarVec st' en' d)+{-# INLINE varFreeze #-}+++{- | A vector converted from Arrow offsets keeps starts and ends as two+overlapping slices of one offsets buffer; writing row @i@ would then+move the start of row @i + 1@, so such a vector gets separate copies.+-}+varThaw :: VarElem a => VarVec a -> ST s (MVarVec s a)+varThaw (VarVec st en d) = do+  st' <- VS.unsafeThaw st+  en' <- if storableOverlap st en then VS.thaw en else VS.unsafeThaw en+  ref <- growThaw d >>= newMutVar+  pure (MVarVec st' en' ref)+{-# INLINE varThaw #-}+++storableOverlap :: VS.Vector Int -> VS.Vector Int -> Bool+storableOverlap a b =+  let !(fa, na) = VS.unsafeToForeignPtr0 a+      !(fb, nb) = VS.unsafeToForeignPtr0 b+      !pa = unsafeForeignPtrToPtr fa `minusPtr` nullPtr+      !pb = unsafeForeignPtrToPtr fb `minusPtr` nullPtr+  in na > 0 && nb > 0 && pa < pb + nb * 8 && pb < pa + na * 8+++mvarLength :: MVarVec s a -> Int+mvarLength (MVarVec st _ _) = VSM.length st+{-# INLINE mvarLength #-}+++mvarSlice :: Int -> Int -> MVarVec s a -> MVarVec s a+mvarSlice i n (MVarVec st en ref) = MVarVec (VSM.unsafeSlice i n st) (VSM.unsafeSlice i n en) ref+{-# INLINE mvarSlice #-}+++mvarOverlaps :: MVarVec s a -> MVarVec s a -> Bool+mvarOverlaps (MVarVec a _ _) (MVarVec b _ _) = VSM.overlaps a b+{-# INLINE mvarOverlaps #-}+++-- | Offsets start zeroed even for 'GM.unsafeNew': an unwritten row is the empty row, never a wild range.+mvarNew :: VarElem a => Int -> ST s (MVarVec s a)+mvarNew n = do+  st <- zeroInts n+  en <- zeroInts n+  ref <- growEmpty >>= newMutVar+  pure (MVarVec st en ref)+{-# INLINE mvarNew #-}+++zeroInts :: Int -> ST s (VSM.MVector s Int)+zeroInts n = do+  v <- VSM.unsafeNew n+  VSM.set v 0+  pure v+++mvarInitialize :: MVarVec s a -> ST s ()+mvarInitialize (MVarVec st en _) = VSM.set st 0 >> VSM.set en 0+{-# INLINE mvarInitialize #-}+++mvarRead :: VarElem a => MVarVec s a -> Int -> ST s a+mvarRead (MVarVec st en ref) i = do+  s <- VSM.unsafeRead st i+  e <- VSM.unsafeRead en i+  g <- readMutVar ref+  growRow g s e+{-# INLINE mvarRead #-}+++mvarWrite :: VarElem a => MVarVec s a -> Int -> a -> ST s ()+mvarWrite (MVarVec st en ref) i x = do+  g <- readMutVar ref+  let !s = growUsed g+  g' <- appendRow g x+  writeMutVar ref g'+  VSM.unsafeWrite st i s+  VSM.unsafeWrite en i (s + rowLength x)+{-# INLINE mvarWrite #-}+++-- | One append, then every row points at it.+mvarSet :: VarElem a => MVarVec s a -> a -> ST s ()+mvarSet (MVarVec st en ref) x = do+  g <- readMutVar ref+  let !s = growUsed g+  g' <- appendRow g x+  writeMutVar ref g'+  VSM.set st s+  VSM.set en (s + rowLength x)+{-# INLINE mvarSet #-}+++-- | Slices of one vector share their store: only offsets move. Otherwise the rows' bytes are appended.+mvarCopy :: VarElem a => MVarVec s a -> MVarVec s a -> ST s ()+mvarCopy dst@(MVarVec dst' den dref) (MVarVec sst sen sref)+  | dref == sref = VSM.unsafeCopy dst' sst >> VSM.unsafeCopy den sen+  | otherwise = do+      sg <- readMutVar sref+      copyRows dst (VSM.length sst) (VSM.unsafeRead sst) (VSM.unsafeRead sen) (\g s e -> appendGrow g sg s e)+{-# INLINE mvarCopy #-}+++mvarMove :: VarElem a => MVarVec s a -> MVarVec s a -> ST s ()+mvarMove dst@(MVarVec dst' den dref) src@(MVarVec sst sen sref)+  | dref == sref = VSM.unsafeMove dst' sst >> VSM.unsafeMove den sen+  | otherwise = mvarCopy dst src+{-# INLINE mvarMove #-}+++-- | Grows the offsets (new rows empty); the store is shared, not copied.+mvarGrow :: MVarVec s a -> Int -> ST s (MVarVec s a)+mvarGrow (MVarVec st en ref) by = do+  let !n = VSM.length st+  st' <- VSM.unsafeGrow st by+  en' <- VSM.unsafeGrow en by+  VSM.set (VSM.unsafeSlice n by st') 0+  VSM.set (VSM.unsafeSlice n by en') 0+  pure (MVarVec st' en' ref)+{-# INLINE mvarGrow #-}+++-- | Copy rows from an immutable vector (appending the bytes they reference).+copyFromVar :: VarElem a => MVarVec s a -> VarVec a -> ST s ()+copyFromVar dst (VarVec st en d) =+  copyRows dst (VS.length st) (\i -> pure (VS.unsafeIndex st i)) (\i -> pure (VS.unsafeIndex en i)) (\g s e -> appendStore g d s e)+{-# INLINE copyFromVar #-}+++-- | Span of the non-empty source rows and their total length.+data Span = Span {-# UNPACK #-} !Int {-# UNPACK #-} !Int {-# UNPACK #-} !Int+++{- | Copy @n@ rows into @dst@. When the non-empty rows reference a span+no longer than their total length (contiguous rows, or rows sharing+bytes such as dictionary values), the span is appended once and the+rows are rebased into it; otherwise each row is appended on its own, so+the copy never carries bytes no row references.+-}+copyRows ::+  VarElem a =>+  MVarVec s a ->+  Int ->+  (Int -> ST s Int) ->+  (Int -> ST s Int) ->+  (Grow s a -> Int -> Int -> ST s (Grow s a)) ->+  ST s ()+copyRows (MVarVec dst den ref) n srcStart srcEnd append = do+  Span lo hi total <- scan 0 maxBound minBound 0+  g0 <- readMutVar ref+  let !base = growUsed g0+  if total == 0+    then VSM.set dst base >> VSM.set den base+    else+      if hi - lo <= total+        then do+          g1 <- append g0 lo hi+          writeMutVar ref g1+          let rebase !i+                | i >= n = pure ()+                | otherwise = do+                    s <- srcStart i+                    e <- srcEnd i+                    if e > s+                      then VSM.unsafeWrite dst i (s - lo + base) >> VSM.unsafeWrite den i (e - lo + base)+                      else VSM.unsafeWrite dst i base >> VSM.unsafeWrite den i base+                    rebase (i + 1)+          rebase 0+        else do+          let each !i g+                | i >= n = writeMutVar ref g+                | otherwise = do+                    s <- srcStart i+                    e <- srcEnd i+                    let !u = growUsed g+                    VSM.unsafeWrite dst i u+                    VSM.unsafeWrite den i (u + max 0 (e - s))+                    if e > s then append g s e >>= each (i + 1) else each (i + 1) g+          each 0 g0+  where+    scan !i !lo !hi !total+      | i >= n = pure (Span lo hi total)+      | otherwise = do+          s <- srcStart i+          e <- srcEnd i+          if e > s then scan (i + 1) (min lo s) (max hi e) (total + e - s) else scan (i + 1) lo hi total+{-# INLINE copyRows #-}+++#define VARLEN(ctx,ty,con,mcon) \+instance ctx => GM.MVector MVector (ty) where { \+  {-# INLINE basicLength #-} \+; {-# INLINE basicUnsafeSlice #-} \+; {-# INLINE basicOverlaps #-} \+; {-# INLINE basicUnsafeNew #-} \+; {-# INLINE basicInitialize #-} \+; {-# INLINE basicUnsafeRead #-} \+; {-# INLINE basicUnsafeWrite #-} \+; {-# INLINE basicSet #-} \+; {-# INLINE basicUnsafeCopy #-} \+; {-# INLINE basicUnsafeMove #-} \+; {-# INLINE basicUnsafeGrow #-} \+; basicLength (mcon v) = mvarLength v \+; basicUnsafeSlice i n (mcon v) = mcon (mvarSlice i n v) \+; basicOverlaps (mcon a) (mcon b) = mvarOverlaps a b \+; basicUnsafeNew n = mcon <$> mvarNew n \+; basicInitialize (mcon v) = mvarInitialize v \+; basicUnsafeRead (mcon v) i = mvarRead v i \+; basicUnsafeWrite (mcon v) i x = mvarWrite v i x \+; basicSet (mcon v) x = mvarSet v x \+; basicUnsafeCopy (mcon a) (mcon b) = mvarCopy a b \+; basicUnsafeMove (mcon a) (mcon b) = mvarMove a b \+; basicUnsafeGrow (mcon v) n = mcon <$> mvarGrow v n }; \+instance ctx => G.Vector Vector (ty) where { \+  {-# INLINE basicUnsafeFreeze #-} \+; {-# INLINE basicUnsafeThaw #-} \+; {-# INLINE basicLength #-} \+; {-# INLINE basicUnsafeSlice #-} \+; {-# INLINE basicUnsafeIndexM #-} \+; {-# INLINE basicUnsafeCopy #-} \+; {-# INLINE elemseq #-} \+; basicUnsafeFreeze (mcon v) = con <$> varFreeze v \+; basicUnsafeThaw (con v) = mcon <$> varThaw v \+; basicLength (con v) = varLength v \+; basicUnsafeSlice i n (con v) = con (varSlice i n v) \+; basicUnsafeIndexM (con v) i = varIndex v i \+; basicUnsafeCopy (mcon m) (con v) = copyFromVar m v \+; elemseq _ = seq }; \+instance ctx => Element (ty)+++-- Text ------------------------------------------------------------++-- | Rows are 'Text' slices of one text array (no copy, no re-validation).+newtype instance Vector Text = V_Text (VarVec Text)+++newtype instance MVector s Text = MV_Text (MVarVec s Text)+++instance VarElem Text where+  type Store Text = TA.Array+  data Grow s Text = TextGrow {-# UNPACK #-} !(MutableByteArray s) {-# UNPACK #-} !Int+  storeRow arr s e+    | e == s = T.empty+    | otherwise = TI.Text arr s (e - s)+  {-# INLINE storeRow #-}+  rowLength (TI.Text _ _ len) = len+  {-# INLINE rowLength #-}+  appendRow g (TI.Text arr off len) = appendStore g arr off (off + len)+  {-# INLINE appendRow #-}+  appendStore g (TA.ByteArray ba) s e = do+    TextGrow mba u <- reserveText g (e - s)+    copyByteArray mba u (ByteArray ba) s (e - s)+    pure (TextGrow mba (u + e - s))+  appendGrow g (TextGrow src _) s e = do+    TextGrow mba u <- reserveText g (e - s)+    copyMutableByteArray mba u src s (e - s)+    pure (TextGrow mba (u + e - s))+  growUsed (TextGrow _ u) = u+  {-# INLINE growUsed #-}+  growRow (TextGrow mba _) s e+    | e == s = pure T.empty+    | otherwise = do+        ByteArray ba <- unsafeFreezeByteArray mba+        pure (TI.Text (TA.ByteArray ba) s (e - s))+  {-# INLINE growRow #-}+  growEmpty = (\mba -> TextGrow mba 0) <$> newByteArray 0+  growFreeze (TextGrow mba _) = (\(ByteArray ba) -> TA.ByteArray ba) <$> unsafeFreezeByteArray mba+  growThaw (TA.ByteArray ba) = (\mba -> TextGrow mba (sizeofByteArray (ByteArray ba))) <$> unsafeThawByteArray (ByteArray ba)+  toVar (V_Text v) = v+  {-# INLINE toVar #-}+  fromVar = V_Text+  {-# INLINE fromVar #-}+++reserveText :: Grow s Text -> Int -> ST s (Grow s Text)+reserveText g@(TextGrow mba u) need = do+  cap <- getSizeofMutableByteArray mba+  if u + need <= cap+    then pure g+    else do+      let !cap' = max (u + need) (max 64 (2 * cap))+      mba' <- newByteArray cap'+      copyMutableByteArray mba' 0 mba 0 u+      pure (TextGrow mba' u)+++VARLEN((), Text, V_Text, MV_Text)+++-- ByteString ------------------------------------------------------++-- | Rows are zero-copy 'ByteString' slices of one buffer.+newtype instance Vector ByteString = V_Bytes (VarVec ByteString)+++newtype instance MVector s ByteString = MV_Bytes (MVarVec s ByteString)+++instance VarElem ByteString where+  type Store ByteString = ByteString+  -- buffer, capacity, used+  data Grow s ByteString = BytesGrow !(ForeignPtr Word8) {-# UNPACK #-} !Int {-# UNPACK #-} !Int+  storeRow (BSI.BS fp _) s e = BSI.BS (fp `plusForeignPtr` s) (e - s)+  {-# INLINE storeRow #-}+  rowLength = BS.length+  {-# INLINE rowLength #-}+  appendRow g bs = appendStore g bs 0 (BS.length bs)+  {-# INLINE appendRow #-}+  appendStore g (BSI.BS fp _) s e = appendBytes g fp s (e - s)+  appendGrow g (BytesGrow fp _ _) s e = appendBytes g fp s (e - s)+  growUsed (BytesGrow _ _ u) = u+  {-# INLINE growUsed #-}+  growRow (BytesGrow fp _ _) s e = pure (BSI.BS (fp `plusForeignPtr` s) (e - s))+  {-# INLINE growRow #-}+  growEmpty = pure (BytesGrow emptyBuffer 0 0)+  growFreeze (BytesGrow fp _ u) = pure (BSI.BS fp u)+  growThaw (BSI.BS fp len) = pure (BytesGrow fp len len)+  toVar (V_Bytes v) = v+  {-# INLINE toVar #-}+  fromVar = V_Bytes+  {-# INLINE fromVar #-}+++emptyBuffer :: ForeignPtr Word8+emptyBuffer = case BS.empty of BSI.BS fp _ -> fp+{-# NOINLINE emptyBuffer #-}+++appendBytes :: Grow s ByteString -> ForeignPtr Word8 -> Int -> Int -> ST s (Grow s ByteString)+appendBytes g src s len+  | len <= 0 = pure g+  | otherwise = do+      BytesGrow fp cap u <- reserveBytes g len+      unsafeIOToST $ unsafeWithForeignPtr fp $ \d -> unsafeWithForeignPtr src $ \p ->+        copyBytes (d `plusPtr` u) (p `plusPtr` s) len+      pure (BytesGrow fp cap (u + len))+++reserveBytes :: Grow s ByteString -> Int -> ST s (Grow s ByteString)+reserveBytes g@(BytesGrow fp cap u) need+  | u + need <= cap = pure g+  | otherwise = unsafeIOToST $ do+      let !cap' = max (u + need) (max 64 (2 * cap))+      fp' <- mallocPlainForeignPtrBytes cap'+      unsafeWithForeignPtr fp' $ \d -> unsafeWithForeignPtr fp $ \p -> copyBytes d p u+      pure (BytesGrow fp' cap' u)+++VARLEN((), ByteString, V_Bytes, MV_Bytes)+++-- Vector a (list rows) --------------------------------------------++-- | Rows are O(1) slices of one child vector.+newtype instance Vector (Vector a) = V_List (VarVec (Vector a))+++newtype instance MVector s (Vector a) = MV_List (MVarVec s (Vector a))+++instance Element a => VarElem (Vector a) where+  type Store (Vector a) = Vector a+  -- child elements (capacity = length), used+  data Grow s (Vector a) = ListGrow !(MVector s a) {-# UNPACK #-} !Int+  storeRow xs s e = G.basicUnsafeSlice s (e - s) xs+  {-# INLINE storeRow #-}+  rowLength = G.basicLength+  {-# INLINE rowLength #-}+  appendRow g xs = appendStore g xs 0 (G.basicLength xs)+  {-# INLINE appendRow #-}+  appendStore g xs s e = do+    ListGrow mv u <- reserveList g (e - s)+    G.basicUnsafeCopy (GM.basicUnsafeSlice u (e - s) mv) (G.basicUnsafeSlice s (e - s) xs)+    pure (ListGrow mv (u + e - s))+  appendGrow g (ListGrow src _) s e = do+    ListGrow mv u <- reserveList g (e - s)+    GM.basicUnsafeCopy (GM.basicUnsafeSlice u (e - s) mv) (GM.basicUnsafeSlice s (e - s) src)+    pure (ListGrow mv (u + e - s))+  growUsed (ListGrow _ u) = u+  {-# INLINE growUsed #-}+  growRow (ListGrow mv _) s e = G.basicUnsafeFreeze (GM.basicUnsafeSlice s (e - s) mv)+  {-# INLINE growRow #-}+  growEmpty = (\mv -> ListGrow mv 0) <$> GM.basicUnsafeNew 0+  growFreeze (ListGrow mv u) = G.basicUnsafeFreeze (GM.basicUnsafeSlice 0 u mv)+  growThaw xs = (\mv -> ListGrow mv (G.basicLength xs)) <$> G.basicUnsafeThaw xs+  toVar (V_List v) = v+  {-# INLINE toVar #-}+  fromVar = V_List+  {-# INLINE fromVar #-}+++reserveList :: Element a => Grow s (Vector a) -> Int -> ST s (Grow s (Vector a))+reserveList g@(ListGrow mv u) need+  | u + need <= cap = pure g+  | otherwise = do+      let !cap' = max (u + need) (max 16 (2 * cap))+      mv' <- GM.basicUnsafeGrow mv (cap' - cap)+      pure (ListGrow mv' u)+  where+    !cap = GM.basicLength mv+++VARLEN((Element a), Vector a, V_List, MV_List)
+ src/Arrow/Write.hs view
@@ -0,0 +1,39 @@+{- | Apache Arrow IPC whole-stream and whole-file writers. Both emit+the standard Arrow IPC format (FlatBuffers metadata, encapsulated+messages, spec footer for files) that pyarrow, arrow-cpp and arrow-rs+read, using 'Arrow.Stream.defaultWriteOptions' (no body compression,+one dictionary batch per dictionary id). Use "Arrow.Stream" directly+for compression, dictionary-replacement options, or the lazy+(zero-copy body) writers.+-}+module Arrow.Write (+  writeArrowStream,+  writeArrowFile,++  -- * Record batch layout (see "Arrow.FlatBufferIPC")+  validateColumns,+  BatchPlan (..),+  planBatch,+) where++import Arrow.Column (ColumnArray)+import Arrow.Stream (defaultWriteOptions, encodeArrowFile, encodeArrowStream)+import Arrow.Types (Schema)+import Arrow.Write.Columns (BatchPlan (..), planBatch, validateColumns)+import Data.ByteString (ByteString)+import Data.Vector qualified as V+++{- | Write a complete Arrow IPC stream (schema, dictionaries, record+batches, end-of-stream marker). 'Left' when a batch does not fit the+schema or a dictionary cannot be written (see 'Arrow.Stream.encodeArrowStream').+-}+writeArrowStream :: Schema -> V.Vector (V.Vector ColumnArray) -> Either String ByteString+writeArrowStream schema batches = encodeArrowStream defaultWriteOptions schema (V.toList batches)+++{- | Write a complete Arrow IPC file (magic, stream payload, FlatBuffers+footer, magic). 'Left' as for 'writeArrowStream'.+-}+writeArrowFile :: Schema -> V.Vector (V.Vector ColumnArray) -> Either String ByteString+writeArrowFile schema batches = encodeArrowFile defaultWriteOptions schema (V.toList batches)
+ src/Arrow/Write/Columns.hs view
@@ -0,0 +1,504 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE ScopedTypeVariables #-}++{- | Arrow IPC record batch layout: the plan pass that turns a batch of+'ColumnArray's into Arrow field nodes and body pieces (spec layout,+depth-first pre-order, a validity slot for every array that has one).++Pieces alias the column buffers. The only copies made here normalise+sliced columns, once per buffer and only when needed: offsets that do+not start at 0 are rebased, and bitmaps whose bit offset is not a+multiple of 8 are re-aligned. Var-length data and list children are+trimmed to the range the offsets reference, struct and sparse union+children to the parent's rows.++The message framing lives in "Arrow.FlatBufferIPC.Write"; the+high-level stream and file writers are in "Arrow.Stream" and+"Arrow.Write".+-}+module Arrow.Write.Columns (+  -- * Schema agreement+  validateColumns,+  validateDecimalFields,++  -- * Layout+  BatchPlan (..),+  planBatch,+  planBatchWithEndianness,+) where++import Arrow.Column (+  columnLength,+  columnTag,+  nullCount,+  rebaseRunEnds,+  sliceColumnArray,+ )+import Arrow.Column.Internal (+  Bitmap (..),+  ChildRange (..),+  ColumnArray,+  IntegralPrim (..),+  Offset,+  PrimType,+  SomePrimType (..),+  Validity (..),+  copyBitmap,+  integralPrim,+  primTypeFor,+  primWidth,+  rebaseToZero,+  samePrimType,+  storableToBytes,+  withPrim,+ )+import Arrow.Column.Internal qualified as I+import Arrow.Types+import Data.Bits (shiftR, (.&.))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.Int (Int16, Int32, Int64, Int8)+import Data.Maybe (isJust)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Vector qualified as V+import Data.Vector.Storable qualified as VS+import Foreign.Storable (Storable, sizeOf)+++-- ============================================================+-- Plan+-- ============================================================++{- | One record batch laid out for the wire: field nodes, one body+piece per 'Buffer' slot (unpadded; the framing pads each to 8 bytes)+and the variadic buffer counts of view columns, all in schema+pre-order.+-}+data BatchPlan = BatchPlan+  { bpLength :: !Int64+  , bpNodes :: !(VS.Vector FieldNode)+  , bpPieces :: !(V.Vector ByteString)+  , bpVariadic :: !(V.Vector Int64)+  }+++-- | Reversed field nodes, pieces and variadic counts, with the first two lengths.+data Acc = Acc ![FieldNode] {-# UNPACK #-} !Int ![ByteString] {-# UNPACK #-} !Int ![Int64]+++{- | Lay out a batch whose columns agree with the fields (see+'validateColumns'; this function only checks what it needs to stay+total). O(columns + buffers) apart from the slice normalisation+described in the module header.+-}+planBatch :: V.Vector Field -> V.Vector ColumnArray -> Either String BatchPlan+planBatch = planBatchWithEndianness Little+++{- | Lay out columns in the schema's declared byte order. Little-endian+pieces keep the normal zero-copy path; big-endian numeric pieces are+copied once, before any body compression. Bitmaps and string bytes alias.+-}+planBatchWithEndianness :: Endianness -> V.Vector Field -> V.Vector ColumnArray -> Either String BatchPlan+planBatchWithEndianness endian fields cols+  | V.length fields /= V.length cols = Left "Arrow.Write: column count differs from the schema's field count"+  | otherwise = do+      Acc ns nn ps np vs <- V.ifoldM' (\acc i f -> planCol endian f (V.unsafeIndex cols i) acc) (Acc [] 0 [] 0 []) fields+      Right+        BatchPlan+          { bpLength = if V.null cols then 0 else fromIntegral (columnLength (V.unsafeHead cols))+          , bpNodes = VS.fromListN nn (reverse ns)+          , bpPieces = V.fromListN np (reverse ps)+          , bpVariadic = V.fromList (reverse vs)+          }+++node :: Int -> Int -> Acc -> Acc+node len nc (Acc ns nn ps np vs) = Acc (FieldNode (fromIntegral len) (fromIntegral nc) : ns) (nn + 1) ps np vs+{-# INLINE node #-}+++piece :: ByteString -> Acc -> Acc+piece p (Acc ns nn ps np vs) = Acc ns nn (p : ps) (np + 1) vs+{-# INLINE piece #-}+++variadic :: Int -> Acc -> Acc+variadic c (Acc ns nn ps np vs) = Acc ns nn ps np (fromIntegral c : vs)+++-- | The node of an array that has a validity slot, followed by its validity piece.+withValidity :: ColumnArray -> Maybe Validity -> Acc -> Acc+withValidity col v acc = validityPiece v (node (columnLength col) (nullCount col) acc)+++-- | The validity piece; a validity without nulls is written like 'Nothing' (an empty slot).+validityPiece :: Maybe Validity -> Acc -> Acc+validityPiece = \case+  Just (Validity bm nc) | nc > 0 -> piece (bitmapPiece bm)+  _ -> piece BS.empty+++{- | The bytes of a bitmap starting at bit 0: an alias when the bit+offset is byte aligned, else one bit-copy.+-}+bitmapPiece :: Bitmap -> ByteString+bitmapPiece bm@(Bitmap bs off len)+  | off .&. 7 == 0 = BS.take nbytes (BS.drop (off `shiftR` 3) bs)+  | otherwise = I.bitmapBytes (copyBitmap bm)+  where+    !nbytes = (len + 7) `shiftR` 3+++-- | Offsets starting at 0 (aliased when they already do) and the child range they cover.+offsetsPiece :: Offset o => VS.Vector o -> (ByteString, ChildRange)+offsetsPiece o+  | not (VS.null o) && VS.unsafeHead o == 0 = (storableToBytes o, ChildRange 0 (fromIntegral (VS.unsafeLast o)))+  | otherwise = let !(o', r) = rebaseToZero o in (storableToBytes o', r)+++sliceRange :: ChildRange -> ColumnArray -> ColumnArray+sliceRange (ChildRange s l) = sliceColumnArray s l+++bytesRange :: ChildRange -> ByteString -> ByteString+bytesRange (ChildRange s l) = BS.take l . BS.drop s+++childField :: Field -> Int -> Either String Field+childField f i = case fieldChildren f V.!? i of+  Just c -> Right c+  Nothing -> Left ("Arrow.Write: field " ++ show (fieldName f) ++ " has no child " ++ show i)+++planCol :: Endianness -> Field -> ColumnArray -> Acc -> Either String Acc+planCol endian f col acc = case col of+  I.ColNull n -> Right (node n n acc)+  I.ColPrim t v xs -> Right (piece (endianPrim endian t (withPrim t (storableToBytes xs))) (withValidity col v acc))+  I.ColBool v bm -> Right (piece (bitmapPiece bm) (withValidity col v acc))+  I.ColUtf8 v o d -> Right (var v o d)+  I.ColBinary v o d -> Right (var v o d)+  I.ColLargeUtf8 v o d -> Right (var v o d)+  I.ColLargeBinary v o d -> Right (var v o d)+  I.ColFixedSizeBinary w n v d -> Right (piece (BS.take (w * n) d) (withValidity col v acc))+  I.ColUtf8View v views bufs -> Right (view v views bufs)+  I.ColBinaryView v views bufs -> Right (view v views bufs)+  I.ColStruct n v cs ->+    V.ifoldM'+      (\a i (_, c) -> childField f i >>= \cf -> planCol endian cf (sliceColumnArray 0 n c) a)+      (withValidity col v acc)+      cs+  I.ColList v o c -> list v o c+  I.ColLargeList v o c -> list v o c+  I.ColListView v o s c -> do+    cf <- childField f 0+    planCol endian cf c (piece (endianVector endian s) (piece (endianVector endian o) (withValidity col v acc)))+  I.ColLargeListView v o s c -> do+    cf <- childField f 0+    planCol endian cf c (piece (endianVector endian s) (piece (endianVector endian o) (withValidity col v acc)))+  I.ColFixedSizeList w n v c -> do+    cf <- childField f 0+    planCol endian cf (sliceColumnArray 0 (n * w) c) (withValidity col v acc)+  I.ColMap v o k x -> do+    entries <- childField f 0+    kf <- childField entries 0+    xf <- childField entries 1+    let !(offs, r) = offsetsPiece o+        !acc1 = piece BS.empty (node (childLength r) 0 (piece (endianOffsets o offs) (withValidity col v acc)))+    planCol endian kf (sliceRange r k) acc1 >>= planCol endian xf (sliceRange r x)+  I.ColDenseUnion t o cs ->+    let !acc1 = piece (endianVector endian o) (piece (unionTypes f t) (node (VS.length t) 0 acc))+    in V.ifoldM' (\a i c -> childField f i >>= \cf -> planCol endian cf c a) acc1 cs+  I.ColSparseUnion t cs ->+    let !n = VS.length t+        !acc1 = piece (unionTypes f t) (node n 0 acc)+    in V.ifoldM' (\a i c -> childField f i >>= \cf -> planCol endian cf (sliceColumnArray 0 n c) a) acc1 cs+  I.ColDictionary _ keys _ -> case keys of+    I.ColPrim kt v ks -> do+      bytes <- keysPiece f kt ks+      Right (piece (endianFixed endian (dictionaryWidth f) bytes) (withValidity col v acc))+    _ -> Left ("Arrow.Write: dictionary " ++ show (fieldName f) ++ " has non-integer keys " ++ columnTag keys)+  I.ColRunEndEncoded {} ->+    rebaseRunEnds col >>= \case+      I.ColRunEndEncoded _ n re vals -> do+        rf <- childField f 0+        vf <- childField f 1+        planCol endian rf re (node n 0 acc) >>= planCol endian vf vals+      other -> Left ("Arrow.Write: rebaseRunEnds returned " ++ columnTag other)+  where+    endianOffsets :: Offset o => VS.Vector o -> ByteString -> ByteString+    endianOffsets o = endianFixed endian (vectorWidth o)+    endianViews bytes = if endian == Big then I.swapViewCopy False bytes else bytes+    var :: Offset o => Maybe Validity -> VS.Vector o -> ByteString -> Acc+    var v o d =+      let !(offs, r) = offsetsPiece o+      in piece (bytesRange r d) (piece (endianOffsets o offs) (withValidity col v acc))+    list :: Offset o => Maybe Validity -> VS.Vector o -> ColumnArray -> Either String Acc+    list v o c = do+      cf <- childField f 0+      let !(offs, r) = offsetsPiece o+      planCol endian cf (sliceRange r c) (piece (endianOffsets o offs) (withValidity col v acc))+    view v views bufs =+      let !acc1 = piece (endianViews (BS.take (16 * columnLength col) views)) (withValidity col v acc)+      in variadic (V.length bufs) (V.foldl' (flip piece) acc1 bufs)+++-- Numeric conversions are per buffer, outside the normal row path.+endianPrim :: Endianness -> PrimType a -> ByteString -> ByteString+endianPrim Little _ bytes = bytes+endianPrim Big t bytes+  | width <= 1 = bytes+  | otherwise = I.swapCopy (I.swapPattern t) width (BS.length bytes `quot` width) bytes+  where+    width = primWidth t+++endianFixed :: Endianness -> Int -> ByteString -> ByteString+endianFixed Little _ bytes = bytes+endianFixed Big width bytes+  | width <= 1 = bytes+  | otherwise = I.swapCopy [width] width (BS.length bytes `quot` width) bytes+++vectorWidth :: forall a. Storable a => VS.Vector a -> Int+vectorWidth _ = sizeOf (undefined :: a)+++endianVector :: Storable a => Endianness -> VS.Vector a -> ByteString+endianVector endian values = endianFixed endian (vectorWidth values) (storableToBytes values)+++dictionaryWidth :: Field -> Int+dictionaryWidth field = case fieldDictionary field of+  Just encoding -> case deIndexType encoding of+    AInt width _ -> width `quot` 8+    _ -> 4+  Nothing -> 4+++{- | Union type ids on the wire: the column holds child positions, the+field may declare other type ids (@AUnion _ ids@). Aliased when the ids+are positional (or absent).+-}+unionTypes :: Field -> VS.Vector Int8 -> ByteString+unionTypes f t = case fieldType f of+  AUnion _ ids+    | not (V.and (V.imap (\i x -> fromIntegral i == x) ids)) ->+        storableToBytes (VS.map (\k -> maybe k fromIntegral (ids V.!? fromIntegral k)) t)+  _ -> storableToBytes t+++{- | Dictionary keys at the width of the field's index type (signed+32-bit when it declares none). Keys already at that width are aliased;+others are converted in one pass. Valid keys are below the dictionary+length, which the stream writer checks against the index type, so the+conversion never truncates a valid key.+-}+keysPiece :: Field -> PrimType a -> VS.Vector a -> Either String ByteString+keysPiece f kt ks = case integralPrim kt of+  Nothing -> Left ("Arrow.Write: dictionary " ++ show (fieldName f) ++ " keys are not integers")+  Just IntegralPrim+    | width == primWidth kt -> Right (storableToBytes ks)+    | otherwise -> Right $ case width of+        1 -> storableToBytes (VS.map fromIntegral ks :: VS.Vector Int8)+        2 -> storableToBytes (VS.map fromIntegral ks :: VS.Vector Int16)+        8 -> storableToBytes (VS.map fromIntegral ks :: VS.Vector Int64)+        _ -> storableToBytes (VS.map fromIntegral ks :: VS.Vector Int32)+  where+    !width = case deIndexType <$> fieldDictionary f of+      Just (AInt bits _) | bits == 8 || bits == 16 || bits == 64 -> bits `quot` 8+      _ -> 4+++-- ============================================================+-- Schema agreement+-- ============================================================++{- | Check that a batch agrees with the schema fields before it is+encoded: one column per field, every column with the same row count,+and every column (recursively) of its field's shape:++* the column type matches the field type (integer width and+  signedness, float precision, date / time / interval unit class,+  decimal precision and scale, fixed widths, list size), and a+  dictionary-encoded field holds a 'ColDictionary' with the field's+  dictionary id, integer keys, and values that match the value type;+* a column with nulls sits under a nullable field (dictionary rows+  count through their keys; nulls in rows the parent does not+  reference are ignored);+* nested columns have one child per child field, struct children have+  at least the struct's row count, a fixed-size list child has at+  least @rows * size@ elements, list and map offsets lie inside the+  child, list view offsets and sizes have the row count, and sparse+  union children cover every row.++Only shapes and lengths are inspected: O(columns), no per-row work.+-}+validateColumns :: V.Vector Field -> V.Vector ColumnArray -> Either String ()+validateColumns fields cols+  | V.length fields /= V.length cols =+      Left+        ( "Arrow.Write: batch has "+            ++ show (V.length cols)+            ++ " columns, the schema has "+            ++ show (V.length fields)+            ++ " fields"+        )+  | otherwise = do+      let !n = if V.null cols then 0 else columnLength (V.head cols)+      case V.findIndex (\c -> columnLength c /= n) cols of+        Just i ->+          Left+            ( "Arrow.Write: column "+                ++ show (fieldName (V.unsafeIndex fields i))+                ++ " has "+                ++ show (columnLength (V.unsafeIndex cols i))+                ++ " rows, the first column has "+                ++ show n+            )+        Nothing -> V.zipWithM_ (checkField []) fields cols+++{- | Check decimal parameters in every field, including nested/dictionary+values, without requiring a batch. Used for schema-only IPC output too.+-}+validateDecimalFields :: V.Vector Field -> Either String ()+validateDecimalFields = V.mapM_ $ \field -> do+  validateDecimalType (fieldType field)+  validateDecimalFields (fieldChildren field)+++-- | One column against its field; @path0@ names the enclosing fields (innermost first).+checkField :: [Text] -> Field -> ColumnArray -> Either String ()+checkField path0 f col = do+  validateDecimalType (fieldType f)+  case fieldDictionary f of+    Just de -> case col of+      I.ColDictionary did keys vals+        | did /= deId de -> bad ("dictionary id " ++ show did ++ ", the field declares " ++ show (deId de))+        | not (integerKeys keys) -> bad ("dictionary keys must be an integer column, not " ++ columnTag keys)+        | nullCount keys > 0 && not (fieldNullable f) -> bad "dictionary column with null rows under a non-nullable field"+        | otherwise -> checkField path0 f {fieldDictionary = Nothing, fieldNullable = True} vals+      _ -> mismatch+    Nothing+      | hasValidity && not (fieldNullable f) && nullCount col > 0 ->+          bad (columnTag col ++ " column with " ++ show (nullCount col) ++ " nulls under a non-nullable field")+      | otherwise -> shape+  where+    path = fieldName f : path0+    bad msg = Left ("Arrow.Write: column " ++ show (T.intercalate "." (reverse path)) ++ ": " ++ msg)+    mismatch = bad ("field type " ++ show (fieldType f) ++ " cannot hold a " ++ columnTag col ++ " column")+    ok = Right ()+    expect cond msg = if cond then ok else bad msg+    hasValidity = case fieldType f of+      ANull -> False+      AUnion _ _ -> False+      ARunEndEncoded -> False+      _ -> True+    kids = fieldChildren f+    rows = columnLength col+    integerKeys = \case+      I.ColPrim kt _ _ -> isJust (integralPrim kt)+      _ -> False+    one k = case V.toList kids of+      [c] -> k c+      _ -> bad (show (fieldType f) ++ " field must have exactly one child")+    sizeIs what n m = expect (n == m) (what ++ " has " ++ show m ++ " entries for " ++ show n ++ " rows")+    -- Offsets inside the child; the referenced child range.+    offsets :: Offset o => VS.Vector o -> Int -> Either String ChildRange+    offsets offs childLen+      | VS.null offs = Left "offsets are empty (a list needs rows + 1 offsets)"+      | VS.head offs < 0 = Left "offsets start below zero"+      | VS.last offs < VS.head offs = Left "offsets end before they start"+      | toInteger (VS.last offs) > toInteger childLen =+          Left ("offsets end at " ++ show (toInteger (VS.last offs)) ++ ", the child has " ++ show childLen ++ " rows")+      | otherwise = Right (ChildRange (fromIntegral (VS.head offs)) (fromIntegral (VS.last offs - VS.head offs)))+    withOffsets :: Offset o => VS.Vector o -> Int -> (ChildRange -> Either String ()) -> Either String ()+    withOffsets offs childLen k = either bad k (offsets offs childLen)+    listOf :: Offset o => VS.Vector o -> ColumnArray -> Either String ()+    listOf offs c = one $ \cf -> withOffsets offs (columnLength c) $ \r -> checkField path cf (sliceRange r c)+    varOf :: Offset o => VS.Vector o -> ByteString -> Either String ()+    varOf offs d = withOffsets offs (BS.length d) (const ok)+    viewOf :: Offset o => VS.Vector o -> VS.Vector o -> ColumnArray -> Either String ()+    viewOf offs sizes c = one $ \cf -> do+      sizeIs "list view offsets" rows (VS.length offs)+      sizeIs "list view sizes" rows (VS.length sizes)+      checkField path cf c+    struct n cs+      | V.length cs /= V.length kids =+          bad ("struct column has " ++ show (V.length cs) ++ " children, the field has " ++ show (V.length kids))+      | otherwise =+          V.zipWithM_+            ( \cf (_, c) -> do+                expect+                  (columnLength c >= n)+                  ("struct child " ++ show (fieldName cf) ++ " has " ++ show (columnLength c) ++ " rows, the struct has " ++ show n)+                checkField path cf (sliceColumnArray 0 n c)+            )+            kids+            cs+    fixed w w' n c+      | w /= w' = mismatch+      | otherwise = one $ \cf -> do+          expect+            (columnLength c >= n * w)+            ("fixed-size list child has " ++ show (columnLength c) ++ " elements, " ++ show n ++ " rows of size " ++ show w ++ " need " ++ show (n * w))+          checkField path cf (sliceColumnArray 0 (n * w) c)+    mapOf offs k v = case V.toList kids of+      [entries] | [kf, vf] <- V.toList (fieldChildren entries) -> do+        expect (columnLength k == columnLength v) ("map keys have " ++ show (columnLength k) ++ " rows, values " ++ show (columnLength v))+        withOffsets offs (columnLength k) $ \r -> do+          checkField (fieldName entries : path) kf (sliceRange r k)+          checkField (fieldName entries : path) vf (sliceRange r v)+      _ -> bad "map field must have one entries struct child with key and value"+    union cs check+      | V.length cs /= V.length kids =+          bad ("union column has " ++ show (V.length cs) ++ " children, the field has " ++ show (V.length kids))+      | otherwise = V.zipWithM_ check kids cs+    prim :: PrimType a -> Either String ()+    prim t = case primTypeFor (fieldType f) of+      Just (SomePrimType et)+        | samePrimType t et -> ok+        | otherwise -> mismatch+      Nothing -> mismatch+    shape = case (fieldType f, col) of+      (ANull, I.ColNull _) -> ok+      (_, I.ColPrim t _ _) -> prim t+      (ABool, I.ColBool {}) -> ok+      (AUtf8, I.ColUtf8 _ o d) -> varOf o d+      (ABinary, I.ColBinary _ o d) -> varOf o d+      (ALargeUtf8, I.ColLargeUtf8 _ o d) -> varOf o d+      (ALargeBinary, I.ColLargeBinary _ o d) -> varOf o d+      (AUtf8View, I.ColUtf8View _ vs _) -> expect (BS.length vs >= 16 * rows) "views buffer shorter than 16 bytes per row"+      (ABinaryView, I.ColBinaryView _ vs _) -> expect (BS.length vs >= 16 * rows) "views buffer shorter than 16 bytes per row"+      (AFixedSizeBinary w, I.ColFixedSizeBinary w' n _ d) -> do+        expect (w == w') ("fixed-size binary width " ++ show w' ++ ", the field declares " ++ show w)+        expect (BS.length d >= w * n) ("fixed-size binary data has " ++ show (BS.length d) ++ " bytes for " ++ show n ++ " rows")+      (AStruct, I.ColStruct n _ cs) -> struct n cs+      (AList, I.ColList _ o c) -> listOf o c+      (ALargeList, I.ColLargeList _ o c) -> listOf o c+      (AFixedSizeList w, I.ColFixedSizeList w' n _ c) -> fixed w w' n c+      (AMap _, I.ColMap _ o k v) -> mapOf o k v+      (AListView, I.ColListView _ o s c) -> viewOf o s c+      (ALargeListView, I.ColLargeListView _ o s c) -> viewOf o s c+      (AUnion Dense _, I.ColDenseUnion ts offs cs) -> do+        sizeIs "dense union offsets" (VS.length ts) (VS.length offs)+        union cs (checkField path)+      (AUnion Sparse _, I.ColSparseUnion ts cs) ->+        union cs $ \cf c -> do+          expect+            (columnLength c >= VS.length ts)+            ("sparse union child " ++ show (fieldName cf) ++ " has " ++ show (columnLength c) ++ " rows, the union has " ++ show (VS.length ts))+          checkField path cf (sliceColumnArray 0 (VS.length ts) c)+      (ARunEndEncoded, I.ColRunEndEncoded _ _ re vals) -> case V.toList kids of+        [rf, vf] -> do+          expect+            (columnLength vals >= columnLength re)+            ("run-end-encoded column has " ++ show (columnLength re) ++ " runs but " ++ show (columnLength vals) ++ " values")+          checkField path rf re+          checkField path vf vals+        _ -> bad "run-end-encoded field must have two children (run_ends, values)"+      _ -> mismatch
+ test-derive/Main.hs view
@@ -0,0 +1,13 @@+module Main (main) where++import Test.Arrow.Derive qualified+import Test.Syd+++main :: IO ()+main =+  sydTest $+    describe "wireform-arrow-derive" $+      sequence_+        [ Test.Arrow.Derive.tests+        ]
+ test-derive/Test/Arrow/Derive.hs view
@@ -0,0 +1,130 @@+{-# LANGUAGE OverloadedStrings #-}++module Test.Arrow.Derive (tests) where++import Arrow.Derive (hasTable)+import Arrow.Record (Table, decodeTable, encodeTable, tableSchema)+import Arrow.Types (Field (..), Schema (..))+import Data.Vector qualified as V+import Test.Arrow.Derive.Instances (outcomeIsRejected)+import Test.Arrow.Derive.Types+import Test.Syd+++tests :: Spec+tests =+  describe "Arrow.Derive" $+    sequence_+      [ recordTests+      , newtypeTests+      , maybeTests+      , coercedTests+      , spliceFailureTests+      ]+++recordTests :: Spec+recordTests =+  describe "record" $+    sequence_+      [ it "schema honours rename + renameStyle and drops skipped column" $ do+          let sch = tableSchema (hasTable :: Table Profile)+              names = V.toList (V.map fieldName (arrowFields sch))+          names `shouldBe` ["name", "profile_age", "email"]+      , it "round-trip restores skipped fields from defaults" $ do+          let rows =+                V.fromList+                  [ Profile "Alice" 30 "alice@x" "secret"+                  , Profile "Bob" 25 "bob@x" "hidden"+                  ]+              (sch, cols) = encodeTable hasTable rows+          (V.length cols == 3) `shouldBe` True+          case decodeTable hasTable sch cols of+            Right rs -> do+              V.length rs `shouldBe` V.length rows+              let r0 = V.head rs+                  r1 = V.last rs+              profileName r0 `shouldBe` "Alice"+              profileAge r0 `shouldBe` 30+              profileEmail r0 `shouldBe` "alice@x"+              profilePrivate r0 `shouldBe` defaultPrivate+              profileName r1 `shouldBe` "Bob"+              profileAge r1 `shouldBe` 25+              profileEmail r1 `shouldBe` "bob@x"+              profilePrivate r1 `shouldBe` defaultPrivate+            Left e -> expectationFailure e+      ]+++newtypeTests :: Spec+newtypeTests =+  describe "newtype-passthrough" $+    sequence_+      [ it "WithTag round-trips Tag column via passthrough instance" $ do+          let rows =+                V.fromList+                  [ WithTag (Tag 7) "alpha"+                  , WithTag (Tag 9) "beta"+                  ]+              (sch, cols) = encodeTable hasTable rows+              names = V.toList (V.map fieldName (arrowFields sch))+          -- Arrow's idiomatic name style is snake_case, applied+          -- automatically when no explicit rename is set.+          names `shouldBe` ["wt_id", "wt_name"]+          V.length cols `shouldBe` 2+          case decodeTable hasTable sch cols of+            Right rs -> rs `shouldBe` rows+            Left e -> expectationFailure e+      ]+++maybeTests :: Spec+maybeTests =+  describe "Maybe-fields" $+    sequence_+      [ it "Maybe Text column reflected in schema as nullable" $ do+          let sch = tableSchema (hasTable :: Table Event)+              flds = V.toList (arrowFields sch)+          map fieldName flds `shouldBe` ["event_id", "event_note"]+          map fieldNullable flds `shouldBe` [False, True]+      , it "Maybe Text round-trips with Just / Nothing values" $ do+          let rows =+                V.fromList+                  [ Event 1 (Just "first")+                  , Event 2 Nothing+                  , Event 3 (Just "third")+                  ]+              (sch, cols) = encodeTable hasTable rows+          case decodeTable hasTable sch cols of+            Right rs -> rs `shouldBe` rows+            Left e -> expectationFailure e+      ]+++coercedTests :: Spec+coercedTests =+  describe "coerced" $+    sequence_+      [ it "coerced field uses underlying-type schema and round-trips" $ do+          let sch = tableSchema (hasTable :: Table Result)+              names = V.toList (V.map fieldName (arrowFields sch))+          names `shouldBe` ["result_name", "result_score"]+          let rows =+                V.fromList+                  [ Result "alpha" (Score 100)+                  , Result "beta" (Score 200)+                  ]+              (sch', cols) = encodeTable hasTable rows+          case decodeTable hasTable sch' cols of+            Right rs -> rs `shouldBe` rows+            Left e -> expectationFailure e+      ]+++spliceFailureTests :: Spec+spliceFailureTests =+  describe "splice-time-rejection" $+    sequence_+      [ it "deriveArrow ''Outcome (a sum type) is rejected at splice time" $+          (outcomeIsRejected) `shouldBe` True+      ]
+ test-derive/Test/Arrow/Derive/Instances.hs view
@@ -0,0 +1,35 @@+{-# LANGUAGE TemplateHaskell #-}+{-# OPTIONS_GHC -Wno-orphans #-}++module Test.Arrow.Derive.Instances (+  outcomeIsRejected,+) where++import Arrow.Derive+import Language.Haskell.TH (recover)+import Test.Arrow.Derive.Types+++deriveArrow ''Tag+deriveArrow ''Profile+deriveArrow ''WithTag+deriveArrow ''Event+deriveArrow ''Result+++{- | Witness that 'deriveArrow' rejects sum types at splice time.++The splice is wrapped in 'recover': if @deriveArrow ''Outcome@+aborts (the expected outcome for a sum type), the recovery+branch returns 'True'; otherwise 'False'. The runtime test in+"Test.Arrow.Derive" then asserts the value is 'True'.+-}+outcomeIsRejected :: Bool+outcomeIsRejected =+  $( recover+       [|True|]+       ( do+           _ <- deriveArrow ''Outcome+           [|False|]+       )+   )
+ test-derive/Test/Arrow/Derive/Types.hs view
@@ -0,0 +1,107 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TemplateHaskell #-}++-- | Annotated fixture types for the Arrow deriver round-trip tests.+module Test.Arrow.Derive.Types (+  Profile (..),+  Tag (..),+  WithTag (..),+  Event (..),+  Score (..),+  Result (..),+  Outcome (..),+  defaultPrivate,+) where++import Data.Int (Int32, Int64)+import Data.Text (Text)+import Wireform.Derive.Backend+import Wireform.Derive.Modifier+import Wireform.Derive.NameStyle+++{- | A record with a mix of rename, renameStyle, and skip+defaults+modifiers.+-}+data Profile = Profile+  { profileName :: !Text+  , profileAge :: !Int32+  , profileEmail :: !Text+  , profilePrivate :: !Text+  }+  deriving (Eq, Show)+++defaultPrivate :: Text+defaultPrivate = "<redacted>"+++{-# ANN profileName (rename "name") #-}+++{-# ANN profileAge (renameStyle SnakeCase) #-}+++{-# ANN profileEmail (renameStyle (StripPrefix "profile" `andThen` SnakeCase)) #-}+++{-# ANN profilePrivate (forBackend backendArrow skip) #-}+{-# ANN profilePrivate (forBackend backendArrow (defaults 'defaultPrivate)) #-}+++{- | A newtype wrapping a primitive Arrow type. The deriver+generates 'HasEncoder' / 'HasDecoder' pass-through instances+for it so it can be used as a column inside another record.+-}+newtype Tag = Tag {unTag :: Int64}+  deriving (Eq, Show)+++{- | A record that embeds the 'Tag' newtype, exercising the+newtype's pass-through column instances at the row level.+-}+data WithTag = WithTag+  { wtId :: !Tag+  , wtName :: !Text+  }+  deriving (Eq, Show)+++{- | Exercises the 'Maybe' lift through 'nullable' / 'nullableD'+in the column codec. The deriver inherits this behaviour from+the @{\-# OVERLAPPING #-\}@ instances in "Arrow.Record.Generic".+-}+data Event = Event+  { eventId :: !Int32+  , eventNote :: !(Maybe Text)+  }+  deriving (Eq, Show)+++{- | A newtype the deriver does /not/ generate column instances+for. Instead, the field that carries it uses the 'coerced'+modifier to delegate to the underlying 'Int32' codec.+-}+newtype Score = Score {unScore :: Int32}+  deriving (Eq, Show)+++{- | Exercises 'coerced'. The 'resultScore' field uses+'Data.Coerce.coerce' to bridge between 'Score' and 'Int32' on+both encoder and decoder sides.+-}+data Result = Result+  { resultName :: !Text+  , resultScore :: !Score+  }+  deriving (Eq, Show)+++{-# ANN resultScore (coerced ''Int32) #-}+++{- | A sum type used to verify the deriver rejects non-record+shapes at splice time.+-}+data Outcome = OutcomeWin Int32 | OutcomeLose+  deriving (Eq, Show)
+ test-interop/Main.hs view
@@ -0,0 +1,1113 @@+{-# LANGUAGE GADTs #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PatternSynonyms #-}++{- | Bidirectional pyarrow interop suite for wireform-arrow.++Both sides define the same case matrix: this module as wireform-arrow+'Schema' plus 'ColumnArray' batches, @pyarrow_interop.py@ (next to this+file) as pyarrow record batches built with pyarrow's own constructors.+Case names tie the two together.++* Direction A: wireform-arrow writes every case as an IPC stream and an+  IPC file (plus zstd / lz4 body-compressed variants when the package is+  built with @-f+zstd@ / @-f+lz4@). pyarrow opens each one, runs+  @validate(full=True)@ on every batch, and compares schema (with custom+  metadata) and values against its own expected batches.+* Direction B: pyarrow writes every case (plain, sliced at a non-zero+  offset, zstd, lz4; stream and file). wireform-arrow decodes each with+  'decodeArrowStream' / 'decodeArrowFile' and compares the schema and the+  logical row values against the batches defined here.++Gating: the Python interpreter comes from @WIREFORM_ARROW_PYTHON@+(default @python3@). When @import pyarrow@ fails the suite prints SKIP+and exits 0, unless @WIREFORM_ARROW_REQUIRE_PYARROW=1@.++@--write DIR@ only writes the direction A files into @DIR@ (no Python+needed), for other readers such as @interop/arrow-rs@.+-}+module Main (main) where++import Arrow.Column (+  ChildRange (..),+  Validity,+  ColumnArray,+  Float16 (..),+  IntervalDayTime (..),+  IntervalMonthDayNano (..),+  PrimType (..),+  asPrim,+  columnLength,+  decimal128FromInteger,+  decimal128ToInteger,+  decimal256FromInteger,+  decimal256ToInteger,+  dictKeyAt,+  expandDictionary,+  fromBools,+  fromByteStrings,+  fromMaybeBinaryView,+  fromMaybeBools,+  fromMaybeByteStrings,+  fromMaybeFixedSizeBinary,+  fromMaybeLargeByteStrings,+  fromMaybeLargeTexts,+  fromMaybeTexts,+  fromMaybeUtf8View,+  fromMaybes,+  fromTexts,+  isValidAt,+  listRange,+  mkDenseUnion,+  mkDictionary,+  mkFixedSizeList,+  mkLargeList,+  mkLargeListView,+  mkList,+  mkListView,+  mkMap,+  mkRunEndEncoded,+  mkSparseUnion,+  mkStruct,+  primColumn,+  toBoolVector,+  toBytesVector,+  toMaybeVector,+  toTextVector,+  validityFromBools,+  withPrim,+  pattern ColBinary,+  pattern ColBinaryView,+  pattern ColBool,+  pattern ColDenseUnion,+  pattern ColDictionary,+  pattern ColFixedSizeBinary,+  pattern ColFixedSizeList,+  pattern ColLargeBinary,+  pattern ColLargeList,+  pattern ColLargeListView,+  pattern ColLargeUtf8,+  pattern ColList,+  pattern ColListView,+  pattern ColMap,+  pattern ColNull,+  pattern ColPrim,+  pattern ColRunEndEncoded,+  pattern ColSparseUnion,+  pattern ColStruct,+  pattern ColUtf8,+  pattern ColUtf8View,+ )+import Arrow.Read.Columns (defaultReadOptions)+import Arrow.Stream (+  DictHandling (..),+  WriteOptions (..),+  decodeArrowFile,+  decodeArrowStream,+  defaultWriteOptions,+  encodeArrowFile,+  encodeArrowStream,+ )+import Arrow.Types (+  ArrowType (..),+  BodyCompressionCodec (..),+  DateUnit (..),+  DictionaryEncoding (..),+  Endianness (..),+  Field (..),+  IntervalUnit (..),+  Precision (..),+  Schema (..),+  TimeUnit (..),+  UnionMode (..),+ )+import Arrow.Vector qualified as AV+import Control.DeepSeq (force)+import Control.Exception (SomeException, evaluate, try)+import Control.Monad (forM, forM_, unless, when)+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.Int (Int32, Int64, Int8)+import Data.IORef (IORef, modifyIORef', newIORef, readIORef)+import Data.List (isSuffixOf, sortOn, stripPrefix)+import Data.Either (partitionEithers)+import Data.Map.Strict qualified as Map+import Data.Maybe (catMaybes, fromMaybe, isJust, mapMaybe)+import Data.Set qualified as Set+import Data.Text (Text)+import Data.Vector qualified as V+import Data.Vector.Generic qualified as VG+import Data.Vector.Storable qualified as VS+import GHC.Float (float2Double)+import System.Directory (createDirectoryIfMissing, doesFileExist)+import System.Environment (getArgs, lookupEnv)+import System.Exit (ExitCode (..), exitFailure, exitSuccess)+import System.FilePath ((</>))+import System.IO (hSetEncoding, stdout, utf8)+import System.IO.Temp (withSystemTempDirectory)+import System.Process (readProcessWithExitCode)+++-- ============================================================+-- Case matrix. Keep in sync with pyarrow_interop.py.+-- ============================================================++data Case = Case+  { caseName :: String+  , caseSchema :: Schema+  , caseBatches :: [V.Vector ColumnArray]+  , caseStreamOnly :: Bool+  -- ^ Replacement / delta dictionaries are only legal in the stream format.+  , caseReplaceDicts :: Bool+  -- ^ Write with 'DictReplaceOnChange' instead of the default 'DictEmitOnce'.+  }+++field :: Text -> Bool -> ArrowType -> [Field] -> Field+field n nullable ty kids = Field n nullable ty (V.fromList kids) Nothing V.empty+++dictField :: Text -> Bool -> ArrowType -> Int64 -> ArrowType -> Field+dictField n nullable valueTy did indexTy =+  Field n nullable valueTy V.empty (Just (DictionaryEncoding did indexTy False)) V.empty+++schemaOf :: [Field] -> Schema+schemaOf fs = Schema (V.fromList fs) Little V.empty V.empty+++simple :: String -> Field -> [ColumnArray] -> Case+simple n f cols = Case n (schemaOf [f]) (map V.singleton cols) False False+++-- | Unwrap a checked fixture constructor.+ok :: Either String ColumnArray -> ColumnArray+ok = either (error . ("bad fixture: " ++)) id+++{- | A non-nullable case @name@ (nulls dropped from the first batch) and a+nullable case @name_nullable@ (first batch with nulls, second batch with+none, so pyarrow omits the validity bitmap there).+-}+pair+  :: String+  -> ArrowType+  -> [Field]+  -> ([a] -> ColumnArray)+  -> ([Maybe a] -> ColumnArray)+  -> [Maybe a]+  -> [a]+  -> [Case]+pair n ty kids plain nullable b1 b2 =+  [ simple n (field "x" False ty kids) [plain (catMaybes b1), plain b2]+  , simple (n ++ "_nullable") (field "x" True ty kids) [nullable b1, nullable (map Just b2)]+  ]+++-- | 'pair' for a fixed-width element type.+primPair :: String -> ArrowType -> PrimType a -> [Maybe a] -> [a] -> [Case]+primPair n ty p = withPrim p (pair n ty [] (primColumn p . VS.fromList) (fromMaybes p . V.fromList))+++-- | 'pair' for a column built from a boxed vector of rows.+boxedPair :: String -> ArrowType -> (V.Vector a -> ColumnArray) -> (V.Vector (Maybe a) -> ColumnArray) -> [Maybe a] -> [a] -> [Case]+boxedPair n ty plain nullable = pair n ty [] (plain . V.fromList) (nullable . V.fromList)+++-- | 'boxedPair' for constructors that only take optional rows.+maybePair :: String -> ArrowType -> (V.Vector (Maybe a) -> ColumnArray) -> [Maybe a] -> [a] -> [Case]+maybePair n ty nullable = boxedPair n ty (nullable . V.map Just) nullable+++i32, i16, i64 :: [Int32] -> ColumnArray+i32 = primColumn PInt32 . VS.fromList+i16 = primColumn PInt16 . VS.fromList . map fromIntegral+i64 = primColumn PInt64 . VS.fromList . map fromIntegral+++i32m :: [Maybe Int32] -> ColumnArray+i32m = fromMaybes PInt32 . V.fromList+++utf8m :: [Maybe Text] -> ColumnArray+utf8m = fromMaybeTexts . V.fromList+++-- | Validity from per-row flags (True = valid).+valid :: [Bool] -> Maybe Validity+valid = validityFromBools . V.fromList+++cases :: [Case]+cases =+  concat+    [ primPair "int8" (AInt 8 True) PInt8 [Just 0, Just 1, Nothing, Just (-1), Just 127, Just (-128)] [42, -42]+    , primPair "int16" (AInt 16 True) PInt16 [Just 0, Just 1, Nothing, Just (-1), Just 32767, Just (-32768)] [1000, -1000]+    , primPair "int32" (AInt 32 True) PInt32 [Just 0, Just 1, Nothing, Just (-1), Just maxBound, Just minBound] [7, -7]+    , primPair "int64" (AInt 64 True) PInt64 [Just 0, Just 1, Nothing, Just (-1), Just maxBound, Just minBound] [123456789012, -5]+    , primPair "uint8" (AInt 8 False) PUInt8 [Just 0, Just 1, Nothing, Just maxBound] [128]+    , primPair "uint16" (AInt 16 False) PUInt16 [Just 0, Just 1, Nothing, Just maxBound] [256]+    , primPair "uint32" (AInt 32 False) PUInt32 [Just 0, Just 1, Nothing, Just maxBound] [65536]+    , primPair "uint64" (AInt 64 False) PUInt64 [Just 0, Just 1, Nothing, Just maxBound] [4294967296]+    , primPair "float16" (AFloatingPoint Half) PFloat16 (map (fmap Float16) [Just 0x0000, Just 0x3C00, Nothing, Just 0xC000, Just 0x7BFF, Just 0x7C00]) [Float16 0x3800]+    , primPair "float32" (AFloatingPoint Single) PFloat [Just 0, Just 1.5, Nothing, Just (-2.5), Just (1 / 0), Just 3.4028234663852886e38] [0.25]+    , primPair "float64" (AFloatingPoint DoublePrecision) PDouble [Just 0, Just 1.5, Nothing, Just (-2.5), Just (-1 / 0), Just 1e300, Just 5e-324] [3.141592653589793]+    , boxedPair "bool" ABool fromBools fromMaybeBools (map Just [True, False] ++ [Nothing] ++ map Just [True, True, False, False, True, False]) [False, True]+    , boxedPair "utf8" AUtf8 fromTexts fromMaybeTexts strs ["second batch"]+    , maybePair "large_utf8" ALargeUtf8 fromMaybeLargeTexts strs ["second batch"]+    , boxedPair "binary" ABinary fromByteStrings fromMaybeByteStrings bins ["z"]+    , maybePair "large_binary" ALargeBinary fromMaybeLargeByteStrings bins ["z"]+    , maybePair "fixed_size_binary4" (AFixedSizeBinary 4) (ok . fromMaybeFixedSizeBinary 4) [Just "abcd", Just "\0\0\0\0", Nothing, Just "\xff\xfe\xfd\xfc"] ["wxyz"]+    , maybePair "utf8_view" AUtf8View fromMaybeUtf8View views ["inline", "an out of line string value"]+    , maybePair "binary_view" ABinaryView fromMaybeBinaryView viewBytes ["inline", BS.replicate 13 0]+    , primPair "date32" (ADate DateDay) PDate32 [Just 0, Just 18000, Nothing, Just (-1), Just 19000] [20000]+    , primPair "date64" (ADate DateMillisecond) PDate64 [Just 0, Just 1555200000000, Nothing, Just (-86400000)] [86400000]+    , primPair "time32_s" (ATime Second 32) PTime32 [Just 0, Just 3600, Nothing, Just 86399] [60]+    , primPair "time32_ms" (ATime Millisecond 32) PTime32 [Just 0, Just 1000, Nothing, Just 86399999] [1]+    , primPair "time64_us" (ATime Microsecond 64) PTime64 [Just 0, Just 12345000000, Nothing, Just 86399999999] [1]+    , primPair "time64_ns" (ATime Nanosecond 64) PTime64 [Just 0, Just 1, Nothing, Just 86399999999999] [2]+    , concatMap timeUnitCases [(Second, "s"), (Millisecond, "ms"), (Microsecond, "us"), (Nanosecond, "ns")]+    , primPair "timestamp_us_new_york" (ATimestamp Microsecond (Just "America/New_York")) PTimestamp [Just 0, Just 1700000000000000, Nothing] [5]+    , decimalPair "decimal32_9_2" (ADecimal32 9 2) (PDecimal32 9 2) fromInteger [Just (-999999999), Just 123, Nothing, Just 0, Just 999999999] [1]+    , decimalPair "decimal64_18_2" (ADecimal64 18 2) (PDecimal64 18 2) fromInteger [Just (-999999999999999999), Just 123, Nothing, Just 0, Just 999999999999999999] [1]+    , decimalPair "decimal128_18_2" (ADecimal 18 2) (PDecimal128 18 2) decimal128FromInteger [Just 0, Just 123, Nothing, Just (-123), Just 999999999999999999] [1]+    , decimalPair "decimal128_38_10" (ADecimal 38 10) (PDecimal128 38 10) decimal128FromInteger [Just 123456789012345678901234567890, Nothing, Just (-1)] [10000000000]+    , decimalPair "decimal256_76_5" (ADecimal256 76 5) (PDecimal256 76 5) decimal256FromInteger [Just 100000, Nothing, Just (-1234567890123456789012345678901234567890123456789012345)] [1]+    , primPair+        "interval_month_day_nano"+        (AInterval MonthDayNano)+        PIntervalMonthDayNano+        [Just (IntervalMonthDayNano 1 2 3), Just (IntervalMonthDayNano 0 0 0), Nothing, Just (IntervalMonthDayNano (-1) (-2) (-3))]+        [IntervalMonthDayNano 12 30 1000000000]+    , primPair "interval_year_month" (AInterval YearMonth) PIntervalYearMonth [Just 0, Just 14, Nothing, Just (-3)] [5]+    , primPair+        "interval_day_time"+        (AInterval DayTime)+        PIntervalDayTime+        [Just (IntervalDayTime 0 0), Just (IntervalDayTime 1 1000), Nothing, Just (IntervalDayTime (-1) (-1))]+        [IntervalDayTime 2 2]+    , [simple "null" (field "x" True ANull []) [ColNull 4, ColNull 1]]+    , structCases+    , listCases+    , listViewCases+    , unionCases+    , dictCases+    , reeCases+    , [metadataCase]+    , mixedCases+    ]+  where+    strs = [Just "", Just "a", Nothing, Just "h\233llo w\246rld", Just "\128512 emoji"]+    bins = [Just "", Just "\0\1\2", Nothing, Just (BS.replicate 20 0xff)]+    views =+      [ Just ""+      , Just "short"+      , Nothing+      , Just "exactly12byt"+      , Just "thirteen byte"+      , Just "this string is definitely longer than twelve bytes"+      ]+    viewBytes = [Just "", Just "short", Nothing, Just "exactly12byt", Just "thirteen byte", Just "this string is definitely longer than twelve bytes"]+    timeUnitCases (u, s) =+      concat+        [ primPair ("timestamp_" ++ s) (ATimestamp u Nothing) PTimestamp ts [42]+        , primPair ("timestamp_" ++ s ++ "_utc") (ATimestamp u (Just "UTC")) PTimestamp ts [42]+        , primPair ("duration_" ++ s) (ADuration u) PDuration [Just 0, Just 60, Nothing, Just (-1)] [1]+        ]+    ts = [Just 0, Just 1700000000, Nothing, Just (-1)]+    decimalPair :: String -> ArrowType -> PrimType d -> (Integer -> d) -> [Maybe Integer] -> [Integer] -> [Case]+    decimalPair n ty p conv b1 b2 = primPair n ty p (map (fmap conv) b1) (map conv b2)+++structCases :: [Case]+structCases =+  [ simple+      "struct"+      (field "x" False AStruct kids)+      [ ok (mkStruct 3 Nothing (V.fromList [("i", i32 [1, 2, 3]), ("s", utf8m [Just "a", Nothing, Just "c"])]))+      , ok (mkStruct 1 Nothing (V.fromList [("i", i32 [4]), ("s", utf8m [Just "d"])]))+      ]+  , simple+      "struct_nullable"+      (field "x" True AStruct kids)+      [ ok+          ( mkStruct+              4+              (valid [True, False, True, True])+              (V.fromList [("i", i32 [1, 0, 2, 3]), ("s", utf8m [Just "a", Nothing, Nothing, Just "c"])])+          )+      , ok (mkStruct 1 (valid [True]) (V.fromList [("i", i32 [4]), ("s", utf8m [Just "d"])]))+      ]+  , -- pa.struct([]): the row count is all a fieldless struct carries.+    simple "struct_empty" (field "x" False AStruct []) [ok (mkStruct 3 Nothing V.empty), ok (mkStruct 1 Nothing V.empty)]+  , simple+      "struct_empty_nullable"+      (field "x" True AStruct [])+      [ok (mkStruct 4 (valid [True, False, True, True]) V.empty), ok (mkStruct 1 (valid [True]) V.empty)]+  ]+  where+    kids = [field "i" False (AInt 32 True) [], field "s" True AUtf8 []]+++listCases :: [Case]+listCases =+  [ simple+      "list_int32"+      (field "x" False AList [item (AInt 32 True)])+      [ ok (mkList Nothing (VS.fromList [0, 2, 2, 5]) (i32m [Just 1, Just 2, Just 3, Nothing, Just 5]))+      , ok (mkList Nothing (VS.fromList [0, 1, 3]) (i32m [Just 6, Just 7, Just 8]))+      ]+  , simple+      "list_int32_nullable"+      (field "x" True AList [item (AInt 32 True)])+      [ ok (mkList (valid [True, False, True, True]) (VS.fromList [0, 2, 2, 2, 5]) (i32m [Just 1, Just 2, Just 3, Nothing, Just 5]))+      , ok (mkList (valid [True, True]) (VS.fromList [0, 1, 3]) (i32m [Just 6, Just 7, Just 8]))+      ]+  , simple+      "large_list_utf8"+      (field "x" False ALargeList [item AUtf8])+      [ ok (mkLargeList Nothing (VS.fromList [0, 2, 2, 3]) (utf8m [Just "a", Just "b", Just "c"]))+      , ok (mkLargeList Nothing (VS.fromList [0, 1]) (utf8m [Just "d"]))+      ]+  , simple+      "large_list_utf8_nullable"+      (field "x" True ALargeList [item AUtf8])+      [ ok (mkLargeList (valid [True, False, True, True]) (VS.fromList [0, 2, 2, 2, 3]) (utf8m [Just "a", Just "b", Just "c"]))+      , ok (mkLargeList (valid [True]) (VS.fromList [0, 1]) (utf8m [Just "d"]))+      ]+  , simple+      "fixed_size_list_int16_3"+      (field "x" False (AFixedSizeList 3) [item (AInt 16 True)])+      [ ok (mkFixedSizeList 3 2 Nothing (i16m [Just 1, Just 2, Just 3, Just 4, Nothing, Just 6]))+      , ok (mkFixedSizeList 3 1 Nothing (i16m [Just 7, Just 8, Just 9]))+      ]+  , simple+      "fixed_size_list_int16_3_nullable"+      (field "x" True (AFixedSizeList 3) [item (AInt 16 True)])+      [ ok (mkFixedSizeList 3 3 (valid [True, False, True]) (i16m [Just 1, Just 2, Just 3, Nothing, Nothing, Nothing, Just 4, Nothing, Just 6]))+      , ok (mkFixedSizeList 3 1 (valid [True]) (i16m [Just 7, Just 8, Just 9]))+      ]+  , -- pa.list_(pa.int32(), 0): rows of zero elements over an empty child.+    simple+      "fixed_size_list_int32_0"+      (field "x" False (AFixedSizeList 0) [item (AInt 32 True)])+      [ok (mkFixedSizeList 0 3 Nothing (i32m [])), ok (mkFixedSizeList 0 2 Nothing (i32m []))]+  , simple+      "fixed_size_list_int32_0_nullable"+      (field "x" True (AFixedSizeList 0) [item (AInt 32 True)])+      [ ok (mkFixedSizeList 0 4 (valid [True, False, True, True]) (i32m []))+      , ok (mkFixedSizeList 0 2 (valid [True, True]) (i32m []))+      ]+  , simple+      "map_utf8_int32"+      (field "x" False (AMap False) [entries])+      [ ok (mkMap Nothing (VS.fromList [0, 2, 2, 3]) (fromTexts (V.fromList ["a", "b", "c"])) (i32m [Just 1, Nothing, Just 3]))+      , ok (mkMap Nothing (VS.fromList [0, 1]) (fromTexts (V.fromList ["d"])) (i32m [Just 4]))+      ]+  , simple+      "map_utf8_int32_nullable"+      (field "x" True (AMap False) [entries])+      [ ok (mkMap (valid [True, False, True, True]) (VS.fromList [0, 2, 2, 2, 3]) (fromTexts (V.fromList ["a", "b", "c"])) (i32m [Just 1, Nothing, Just 3]))+      , ok (mkMap (valid [True]) (VS.fromList [0, 1]) (fromTexts (V.fromList ["d"])) (i32m [Just 4]))+      ]+  ]+  where+    i16m = fromMaybes PInt16 . V.fromList+    entries =+      field "entries" False AStruct [field "key" False AUtf8 [], field "value" True (AInt 32 True) []]+++item :: ArrowType -> Field+item ty = field "item" True ty []+++listViewCases :: [Case]+listViewCases =+  [ simple+      "list_view_int32"+      (field "x" False AListView [item (AInt 32 True)])+      [ ok (mkListView Nothing (VS.fromList [4, 0, 1, 0]) (VS.fromList [2, 3, 0, 1]) child)+      , ok (mkListView Nothing (VS.fromList [0]) (VS.fromList [1]) child2)+      ]+  , simple+      "list_view_int32_nullable"+      (field "x" True AListView [item (AInt 32 True)])+      [ ok (mkListView (valid [True, False, True, True, True]) (VS.fromList [4, 0, 1, 0, 2]) (VS.fromList [2, 0, 0, 1, 2]) child)+      , ok (mkListView (valid [True]) (VS.fromList [0]) (VS.fromList [1]) child2)+      ]+  , simple+      "large_list_view_int32"+      (field "x" False ALargeListView [item (AInt 32 True)])+      [ ok (mkLargeListView Nothing (VS.fromList [4, 0, 1, 0]) (VS.fromList [2, 3, 0, 1]) child)+      , ok (mkLargeListView Nothing (VS.fromList [0]) (VS.fromList [1]) child2)+      ]+  , simple+      "large_list_view_int32_nullable"+      (field "x" True ALargeListView [item (AInt 32 True)])+      [ ok (mkLargeListView (valid [True, False, True, True, True]) (VS.fromList [4, 0, 1, 0, 2]) (VS.fromList [2, 0, 0, 1, 2]) child)+      , ok (mkLargeListView (valid [True]) (VS.fromList [0]) (VS.fromList [1]) child2)+      ]+  ]+  where+    child = i32m [Just 10, Just 20, Just 30, Nothing, Just 50, Just 60]+    child2 = i32m [Just 99]+++{- | Union columns store per-row /child indices/; the field's type codes+map them to wire type ids.+-}+unionCases :: [Case]+unionCases =+  [ dense "dense_union" [0, 1] False+  , dense "dense_union_nullable" [0, 1] True+  , sparse "sparse_union" [0, 1] False+  , sparse "sparse_union_nullable" [0, 1] True+  , dense "dense_union_type_codes" [3, 7] True+  , sparse "sparse_union_type_codes" [5, 2] True+  ]+  where+    kids = [field "i" True (AInt 32 True) [], field "s" True AUtf8 []]+    ufield mode codes nullable = field "x" nullable (AUnion mode (V.fromList codes)) kids+    tids = VS.fromList :: [Int8] -> VS.Vector Int8+    dense n codes nullable =+      simple+        n+        (ufield Dense codes nullable)+        [ ok+            ( mkDenseUnion+                (tids [0, 1, 0, 0, 1])+                (VS.fromList [0, 0, 1, 2, 1])+                (V.fromList [i32m [Just 1, Just 2, Nothing], utf8m [Just "a", Just "b"]])+            )+        , ok (mkDenseUnion (tids [1]) (VS.fromList [0]) (V.fromList [i32m [], utf8m [Just "z"]]))+        ]+    sparse n codes nullable =+      simple+        n+        (ufield Sparse codes nullable)+        [ ok+            ( mkSparseUnion+                (tids [0, 1, 0, 1])+                (V.fromList [i32m [Just 1, Just 0, Nothing, Just 0], utf8m [Nothing, Just "a", Just "x", Just "b"]])+            )+        , ok (mkSparseUnion (tids [1]) (V.fromList [i32m [Just 0], utf8m [Just "z"]]))+        ]+++abc :: ColumnArray+abc = fromTexts (V.fromList ["a", "b", "c"])+++-- | A dictionary column over @vals@ with the given keys column (at the index type's wire width).+dict :: Int64 -> ColumnArray -> ColumnArray -> ColumnArray+dict did keys vals = ok (mkDictionary did keys vals)+++dictCases :: [Case]+dictCases =+  [ simple+      "dict_utf8"+      (dictField "x" False AUtf8 0 i32t)+      [dict 0 (ix [0, 1, 0, 2, 1]) abc, dict 0 (ix [2, 2, 0]) abc]+  , simple+      "dict_utf8_nullable"+      (dictField "x" True AUtf8 0 i32t)+      [ dict 0 (i32m [Just 0, Nothing, Just 1, Just 2]) abc+      , dict 0 (i32m [Just 1, Just 0]) abc+      ]+  , simple+      "dict_int8_index"+      (dictField "x" False AUtf8 0 (AInt 8 True))+      [dict 0 (ix8 [1, 0, 1]) xy, dict 0 (ix8 [0]) xy]+  , simple+      "dict_int64_values"+      (dictField "x" False (AInt 64 True) 0 i32t)+      [dict 0 (ix [1, 0, 1]) (i64 [100, 200]), dict 0 (ix [0]) (i64 [100, 200])]+  , (simple "dict_replacement" (dictField "x" False AUtf8 0 i32t) [dict 0 (ix [0, 1, 1]) ab, dict 0 (ix [2, 0, 1]) xyz])+      { caseStreamOnly = True+      , caseReplaceDicts = True+      }+  , (simple "dict_delta" (dictField "x" False AUtf8 0 i32t) [dict 0 (ix [0, 1]) ab, dict 0 (ix [2, 0]) abc])+      { caseStreamOnly = True+      , caseReplaceDicts = True+      }+  , simple+      "dict_in_struct"+      (field "x" False AStruct [dictField "d" True AUtf8 0 i32t])+      [ ok (mkStruct 3 Nothing (V.singleton ("d", dict 0 (i32m [Just 2, Just 0, Just 1]) abc)))+      , ok (mkStruct 1 Nothing (V.singleton ("d", dict 0 (i32m [Just 1]) abc)))+      ]+  , -- Dictionaries nested in dictionary values. pyarrow numbers dictionary+    -- ids in schema pre-order, so the outer dictionary is 0 and the inner 1.+    simple+      "dict_struct_of_dict"+      (nestedDict "x" False AStruct (dictField "d" True AUtf8 1 (AInt 16 True)) (AInt 8 True))+      [dict 0 (ix8 [2, 0, 1, 0]) (structOfDict [1, 0, 2] abc), dict 0 (ix8 [1]) (structOfDict [1, 0, 2] abc)]+  , simple+      "dict_list_of_dict"+      (nestedDict "x" True AList (dictField "item" True AUtf8 1 (AInt 8 True)) i32t)+      [ dict 0 (i32m [Just 1, Nothing, Just 0, Just 2]) listOfDict+      , dict 0 (i32m [Just 2]) listOfDict+      ]+  , -- The inner dictionary changes between batches, so both dictionaries are replaced.+    (simple+      "dict_nested_replacement"+      (nestedDict "x" False AStruct (dictField "d" True AUtf8 1 (AInt 16 True)) (AInt 8 True))+      [dict 0 (ix8 [0, 1]) (structOfDict [0, 1] abc), dict 0 (ix8 [1, 0]) (structOfDict [0, 0] xy)])+      { caseStreamOnly = True+      , caseReplaceDicts = True+      }+  , -- A nullable dictionary column whose rows are all null over an empty dictionary.+    simple+      "dict_all_null_empty"+      (dictField "x" True AUtf8 0 i32t)+      [dict 0 (i32m (replicate 3 Nothing)) (fromTexts V.empty), dict 0 (i32m [Nothing]) (fromTexts V.empty)]+  ]+  where+    i32t = AInt 32 True+    ix = primColumn PInt32 . VS.fromList+    ix8 = primColumn PInt8 . VS.fromList+    ix16 = primColumn PInt16 . VS.fromList+    xy = fromTexts (V.fromList ["x", "y"])+    ab = fromTexts (V.fromList ["a", "b"])+    xyz = fromTexts (V.fromList ["x", "y", "z"])+    nestedDict n nullable container inner indexTy =+      Field n nullable container (V.singleton inner) (Just (DictionaryEncoding 0 indexTy False)) V.empty+    structOfDict inner vals = ok (mkStruct (length inner) Nothing (V.singleton ("d", dict 1 (ix16 inner) vals)))+    listOfDict = ok (mkList Nothing (VS.fromList [0, 2, 2, 3]) (dict 1 (ix8 [0, 1, 1]) xy))+++reeCases :: [Case]+reeCases =+  [ ree "ree_int32_int64" False (AInt 32 True) (AInt 64 True) [(i32 [3, 5, 8], int64m [Just 100, Just 200, Just 300]), (i32 [2], int64m [Just 7])]+  , ree "ree_int32_int64_nullable" True (AInt 32 True) (AInt 64 True) [(i32 [3, 5, 8], int64m [Just 100, Nothing, Just 300]), (i32 [2], int64m [Just 7])]+  , ree "ree_int16_utf8" True (AInt 16 True) AUtf8 [(i16 [1, 4], utf8m [Just "a", Just "bb"]), (i16 [3], utf8m [Just "c"])]+  , ree "ree_int64_float64" True (AInt 64 True) (AFloatingPoint DoublePrecision) [(i64 [2, 3], dblm [Just 1.5, Nothing]), (i64 [1], dblm [Just 2.5])]+  ]+  where+    int64m = fromMaybes PInt64 . V.fromList+    dblm = fromMaybes PDouble . V.fromList+    ree n nullable reTy valTy batches =+      simple+        n+        (field "x" nullable ARunEndEncoded [field "run_ends" False reTy [], field "values" True valTy []])+        (map (ok . uncurry mkRunEndEncoded) batches)+++metadataCase :: Case+metadataCase =+  Case+    { caseName = "custom_metadata"+    , caseSchema =+        Schema+          { arrowFields =+              V.fromList+                [ (field "a" False (AInt 32 True) []) {fieldMetadata = V.fromList [("unit", "ms"), ("note", "field level")]}+                , (field "b" True AUtf8 []) {fieldMetadata = V.fromList [("k", "v")]}+                ]+          , arrowEndianness = Little+          , arrowMetadata = V.fromList [("origin", "wireform-arrow interop"), ("k2", "v2")]+          , arrowFeatures = V.empty+          }+    , caseBatches =+        [ V.fromList [i32 [1, 2], utf8m [Just "x", Nothing]]+        , V.fromList [i32 [3], utf8m [Just "y"]]+        ]+    , caseStreamOnly = False+    , caseReplaceDicts = False+    }+++mixedCases :: [Case]+mixedCases =+  [ multi "mixed" [m1, m2]+  , multi "zero_row_batches" [m1, m0, m2]+  , multi "only_zero_row_batch" [m0]+  , multi "no_batches" []+  ]+  where+    multi n bs = Case n mixedSchema bs False False+    mixedSchema =+      schemaOf+        [ field "i" False (AInt 64 True) []+        , field "s" True AUtf8 []+        , field "b" True ABool []+        , field "l" True AList [item (AInt 32 True)]+        , field "st" True AStruct [field "f" True (AFloatingPoint DoublePrecision) []]+        , dictField "d" True AUtf8 0 (AInt 32 True)+        , field "v" True AUtf8View []+        ]+    mk i s b (lv, lo, lc) (sv, sf) d v =+      V.fromList+        [ i64 i+        , utf8m s+        , fromMaybeBools (V.fromList b)+        , ok (mkList (valid lv) (VS.fromList lo) (i32m lc))+        , ok (mkStruct (length sv) (valid sv) (V.singleton ("f", fromMaybes PDouble (V.fromList sf))))+        , dict 0 (i32m d) abc+        , fromMaybeUtf8View (V.fromList v)+        ]+    m1 =+      mk+        [10, 20, 30]+        [Just "hello", Nothing, Just "!"]+        [Just True, Nothing, Just False]+        ([True, False, True], [0, 1, 1, 3], [Just 1, Just 2, Just 3])+        ([True, False, True], [Just 1.5, Nothing, Just 2.5])+        [Just 0, Just 2, Nothing]+        [Just "v", Nothing, Just "a view longer than twelve"]+    m0 = mk [] [] [] ([], [0], []) ([], []) [] []+    m2 =+      mk+        [40]+        [Just "w"]+        [Just True]+        ([True], [0, 0], [])+        ([True], [Just 0.5])+        [Just 1]+        [Just "z"]+++-- ============================================================+-- Logical values: layout-independent row contents, used to compare+-- what wireform-arrow decoded with what the case expects.+-- ============================================================++data LV+  = LNull+  | LInt !Integer+  | LDouble !Double+  | LBool !Bool+  | LText !Text+  | LBytes !ByteString+  | LList [LV]+  | LStruct [(Text, LV)]+  | LMap [(LV, LV)]+  | LUnion !Int32 LV+  deriving stock (Eq, Show)+++type Rows = V.Vector LV+++-- | Converted rows, through a boxed vector ('VG.convert').+opt :: AV.Element a => (a -> LV) -> AV.Vector (Maybe a) -> Rows+opt f = V.map (maybe LNull f) . VG.convert+++childField :: Field -> Int -> Either String Field+childField f i =+  maybe (Left ("field " ++ show (fieldName f) ++ " has no child " ++ show i)) Right (fieldChildren f V.!? i)+++sliceRows :: Rows -> Int -> Int -> Either String [LV]+sliceRows rs o n+  | o < 0 || n < 0 || o + n > V.length rs =+      Left ("slice [" ++ show o ++ ", +" ++ show n ++ ") outside child of length " ++ show (V.length rs))+  | otherwise = Right (V.toList (V.slice o n rs))+++at :: Rows -> Int -> Either String LV+at rs i = maybe (Left ("index " ++ show i ++ " outside child of length " ++ show (V.length rs))) Right (rs V.!? i)+++-- | Rows of a fixed-width column, read through its typed view.+primRows :: PrimType a -> ColumnArray -> Either String Rows+primRows p col = case p of+  PInt8 -> go (LInt . toInteger)+  PInt16 -> go (LInt . toInteger)+  PInt32 -> go (LInt . toInteger)+  PInt64 -> go (LInt . toInteger)+  PUInt8 -> go (LInt . toInteger)+  PUInt16 -> go (LInt . toInteger)+  PUInt32 -> go (LInt . toInteger)+  PUInt64 -> go (LInt . toInteger)+  PFloat16 -> go (\(Float16 w) -> LInt (toInteger w))+  PFloat -> go (LDouble . float2Double)+  PDouble -> go LDouble+  PDate32 -> go (LInt . toInteger)+  PDate64 -> go (LInt . toInteger)+  PTime32 -> go (LInt . toInteger)+  PTime64 -> go (LInt . toInteger)+  PTimestamp -> go (LInt . toInteger)+  PDuration -> go (LInt . toInteger)+  PIntervalYearMonth -> go (LInt . toInteger)+  PIntervalDayTime -> go (\(IntervalDayTime d m) -> LList [LInt (toInteger d), LInt (toInteger m)])+  PIntervalMonthDayNano -> go (\(IntervalMonthDayNano m d n) -> LList [LInt (toInteger m), LInt (toInteger d), LInt (toInteger n)])+  PDecimal32 _ _ -> go (LInt . toInteger)+  PDecimal64 _ _ -> go (LInt . toInteger)+  PDecimal128 _ _ -> go (LInt . decimal128ToInteger)+  PDecimal256 _ _ -> go (LInt . decimal256ToInteger)+  where+    go conv = withPrim p (maybe (Left ("asPrim failed on " ++ show col)) (Right . opt conv . toMaybeVector) (asPrim p col))+++rowsOf :: Field -> ColumnArray -> Either String Rows+rowsOf f col = case col of+  ColNull n -> Right (V.replicate n LNull)+  ColPrim p _ _ -> primRows p col+  ColBool {} -> opt LBool <$> toBoolVector col+  ColUtf8 {} -> texts+  ColLargeUtf8 {} -> texts+  ColUtf8View {} -> texts+  ColBinary {} -> bytes+  ColLargeBinary {} -> bytes+  ColFixedSizeBinary {} -> bytes+  ColBinaryView {} -> bytes+  ColStruct n v cs -> do+    crs <- V.imapM (\i (nm, c) -> (,) nm <$> childRows i c) cs+    V.generateM n $ \j ->+      if isValidAt v j+        then LStruct . V.toList <$> traverse (\(nm, rs) -> (,) nm <$> at rs j) crs+        else Right LNull+  ColList _ _ c -> lists c+  ColLargeList _ _ c -> lists c+  ColListView _ _ _ c -> lists c+  ColLargeListView _ _ _ c -> lists c+  ColFixedSizeList _ _ _ c -> lists c+  ColMap _ _ ks vs -> do+    entries <- childField f 0+    kf <- childField entries 0+    vf <- childField entries 1+    kr <- rowsOf kf ks+    vr <- rowsOf vf vs+    unless (V.length kr == V.length vr) (Left "map keys and values differ in length")+    perRow $ \(ChildRange s l) -> do+      kl <- sliceRows kr s l+      vl <- sliceRows vr s l+      Right (LMap (zip kl vl))+  ColDenseUnion tids offs cs -> do+    crs <- V.imapM childRows cs+    V.fromList+      <$> traverse+        (\(t, o) -> do+            rs <- maybe (Left ("union child index " ++ show t ++ " out of range")) Right (crs V.!? fromIntegral t)+            LUnion (typeCode t) <$> at rs (fromIntegral o)+        )+        (zip (VS.toList tids) (VS.toList offs))+  ColSparseUnion tids cs -> do+    crs <- V.imapM childRows cs+    V.fromList+      <$> traverse+        (\(i, t) -> do+            rs <- maybe (Left ("union child index " ++ show t ++ " out of range")) Right (crs V.!? fromIntegral t)+            LUnion (typeCode t) <$> at rs i+        )+        (zip [0 ..] (VS.toList tids))+  ColDictionary _ _ vals -> do+    vr <- rowsOf (f {fieldDictionary = Nothing}) vals+    V.generateM (columnLength col) (maybe (Right LNull) (at vr) . dictKeyAt col)+  ColRunEndEncoded off len ends vals -> do+    er <- childRows 0 ends+    vr <- childRows 1 vals+    endsI <- traverse (\case LInt e -> Right (fromInteger e); other -> Left ("non-integer run end " ++ show other)) (V.toList er)+    let expand _ [] _ = Right []+        expand prev (e : es) k+          | e < prev = Left "run ends are not ascending"+          | otherwise = do+              v <- at vr k+              rest <- expand e es (k + 1)+              Right (replicate (e - prev) v ++ rest)+    full <- V.fromList <$> expand (0 :: Int) endsI 0+    V.fromList <$> sliceRows full off len+  where+    texts = opt LText <$> toTextVector col+    bytes = opt LBytes <$> toBytesVector col+    childRows i c = do+      cf <- childField f i+      rowsOf cf c+    -- One row per list slot; 'listRange' answers Nothing for a null row.+    perRow row = V.generateM (columnLength col) (maybe (Right LNull) row . listRange col)+    lists c = do+      cr <- childRows 0 c+      perRow (\(ChildRange s l) -> LList <$> sliceRows cr s l)+    typeCode :: Int8 -> Int32+    typeCode t = case fieldType f of+      AUnion _ codes | not (V.null codes) -> fromMaybe (-1) (codes V.!? fromIntegral t)+      _ -> fromIntegral t+++-- | Schemas compare equal up to the order of custom-metadata pairs.+normaliseSchema :: Schema -> Schema+normaliseSchema s =+  s+    { arrowFields = V.map normaliseField (arrowFields s)+    , arrowMetadata = sortMeta (arrowMetadata s)+    }+  where+    normaliseField fl =+      fl+        { fieldChildren = V.map normaliseField (fieldChildren fl)+        , fieldMetadata = sortMeta (fieldMetadata fl)+        }+    sortMeta = V.fromList . sortOn fst . V.toList+++{- | Compare a decoded file with the case's expectation. Top-level+dictionary columns must also expand ('expandDictionary') to the same+rows.+-}+compareDecoded :: Case -> Schema -> [V.Vector ColumnArray] -> Either String ()+compareDecoded c sch batches = do+  unless (normaliseSchema sch == normaliseSchema (caseSchema c)) $+    Left ("schema mismatch: got " ++ show sch ++ " want " ++ show (caseSchema c))+  unless (length batches == length (caseBatches c)) $+    Left ("got " ++ show (length batches) ++ " batches, want " ++ show (length (caseBatches c)))+  forM_ (zip3 [0 :: Int ..] batches (caseBatches c)) $ \(bi, got, want) -> do+    let fields = arrowFields (caseSchema c)+    unless (V.length got == V.length fields) $+      Left ("batch " ++ show bi ++ ": got " ++ show (V.length got) ++ " columns")+    forM_ (zip3 (V.toList fields) (V.toList got) (V.toList want)) $ \(fl, g, w) -> do+      let ctx = "batch " ++ show bi ++ " column " ++ show (fieldName fl) ++ ": "+      gr <- either (\e -> Left (ctx ++ "decoded column is malformed: " ++ e ++ " (" ++ show g ++ ")")) Right (rowsOf fl g)+      wr <- either (\e -> Left (ctx ++ "expected column is malformed: " ++ e)) Right (rowsOf fl w)+      unless (gr == wr) $+        Left (ctx ++ "got " ++ show (V.toList gr) ++ " want " ++ show (V.toList wr) ++ " (decoded " ++ show g ++ ")")+      when (isJust (fieldDictionary fl)) $ do+        expanded <- either (\e -> Left (ctx ++ "expandDictionary failed: " ++ e)) Right (expandDictionary g)+        er <- either (\e -> Left (ctx ++ "expanded column is malformed: " ++ e)) Right (rowsOf fl {fieldDictionary = Nothing} expanded)+        unless (er == gr) $+          Left (ctx ++ "expanded rows " ++ show (V.toList er) ++ " differ from " ++ show (V.toList gr))+++-- ============================================================+-- Driver+-- ============================================================++data Tally = Tally {tPass :: !Int, tFail :: !Int, tSkip :: !Int, tFailures :: [String]}+++report :: IORef Tally -> String -> String -> Either String () -> IO ()+report ref dir name = \case+  Right () -> do+    putStrLn ("[" ++ dir ++ "] PASS " ++ name)+    modifyIORef' ref (\t -> t {tPass = tPass t + 1})+  Left e -> do+    let line = "[" ++ dir ++ "] FAIL " ++ name ++ ": " ++ truncateMsg e+    putStrLn line+    modifyIORef' ref (\t -> t {tFail = tFail t + 1, tFailures = line : tFailures t})+++skip :: IORef Tally -> String -> String -> String -> IO ()+skip ref dir name why = do+  putStrLn ("[" ++ dir ++ "] SKIP " ++ name ++ ": " ++ why)+  modifyIORef' ref (\t -> t {tSkip = tSkip t + 1})+++truncateMsg :: String -> String+truncateMsg s = if length s > 2000 then take 2000 s ++ " ..." else s+++-- | Body-compression variants this build of wireform-arrow supports.+compressionVariants :: [(String, Maybe BodyCompressionCodec)]+compressionVariants =+  [("plain", Nothing)]+#ifdef HAVE_ZSTD+    ++ [("zstd", Just BodyZstd)]+#endif+#ifdef HAVE_LZ4+    ++ [("lz4", Just LZ4Frame)]+#endif+++-- | Whether this build can decode files pyarrow wrote with @variant@.+variantSupported :: String -> Bool+variantSupported v = case lookup v compressionVariants of+  Just _ -> True+  Nothing -> v == "sliced"+++{- | Write every case into @dir@ as @<case>.<variant>.arrows@ (and+@.arrow@ unless stream-only). Returns the file names written plus the+cases whose encoder threw.+-}+writeHaskellFiles :: Endianness -> FilePath -> IO ([String], [(String, String)])+writeHaskellFiles endian dir = do+  createDirectoryIfMissing True dir+  results <- forM (map (\c -> c {caseSchema = (caseSchema c) {arrowEndianness = endian}}) cases) $ \c -> forM compressionVariants $ \(variant, codec) -> do+    let opts =+          defaultWriteOptions+            { writeBodyCompression = codec+            , writeDictHandling = if caseReplaceDicts c then DictReplaceOnChange else DictEmitOnce+            }+        base = caseName c ++ "." ++ variant+        outputs =+          (base ++ ".arrows", encodeArrowStream opts (caseSchema c) (caseBatches c))+            : if caseStreamOnly c then [] else [(base ++ ".arrow", encodeArrowFile opts (caseSchema c) (caseBatches c))]+    forM outputs $ \(name, bytes) -> do+      r <- try (evaluate (force bytes)) :: IO (Either SomeException (Either String ByteString))+      case r of+        Right (Right bs) -> do+          let decode = if ".arrow" `isSuffixOf` name then decodeArrowFile defaultReadOptions else decodeArrowStream defaultReadOptions+          case decode bs >>= uncurry (compareDecoded c) of+            Left e -> pure (Left (name, "native writer/reader roundtrip: " ++ e))+            Right () -> do+              BS.writeFile (dir </> name) bs+              pure (Right name)+        Right (Left e) -> pure (Left (name, "wireform-arrow encoder rejected the case: " ++ e))+        Left e -> pure (Left (name, "wireform-arrow encoder threw: " ++ show e))+  let (failures, written) = partitionEithers (concat (concat results))+  pure (written, failures)+++runPython :: FilePath -> FilePath -> [String] -> IO (Either String [String])+runPython python script args = do+  r <- try (readProcessWithExitCode python (script : args) "") :: IO (Either SomeException (ExitCode, String, String))+  pure $ case r of+    Left e -> Left ("could not run " ++ python ++ ": " ++ show e)+    Right (ExitSuccess, out, _) -> Right (lines out)+    Right (ExitFailure n, out, err) ->+      Left (python ++ " " ++ unwords (script : args) ++ " exited " ++ show n ++ ": " ++ err ++ out)+++-- | Direction A: wireform-arrow writes, pyarrow reads.+directionA :: Endianness -> IORef Tally -> FilePath -> FilePath -> FilePath -> IO ()+directionA endian ref python script tmp = do+  let dir = tmp </> ("wireform-" ++ show endian)+  (written, encodeFailures) <- writeHaskellFiles endian dir+  forM_ encodeFailures $ \(name, e) -> report ref "A" name (Left e)+  out <- runPython python script ["check", dir]+  case out of+    Left e -> report ref "A" "pyarrow checker" (Left e)+    Right ls -> do+      forM_ (filter (\line -> take 8 line == "ADAPTER ") ls) putStrLn+      let verdicts =+            Map.fromList+              ( mapMaybe+                  ( \l -> case stripPrefix "PASS " l of+                      Just n -> Just (n, Right ())+                      Nothing -> case stripPrefix "FAIL " l of+                        Just rest -> let (n, msg) = break (== ':') rest in Just (n, Left (drop 2 msg))+                        Nothing -> Nothing+                  )+                  ls+              )+      forM_ written $ \name ->+        report ref "A" name (fromMaybe (Left ("pyarrow checker printed no verdict; output: " ++ unlines ls)) (Map.lookup name verdicts))+++-- | Direction B: pyarrow writes, wireform-arrow reads.+directionB :: IORef Tally -> FilePath -> FilePath -> FilePath -> IO ()+directionB ref python script tmp = do+  out <- runPython python script ["write", dir]+  case out of+    Left e -> report ref "B" "pyarrow writer" (Left e)+    Right ls -> forM_ ls $ \l -> case words l of+      ("WROTE" : name : _) -> checkFile name+      ("SKIP" : name : _) -> skip ref "B" (takeWhile (/= ':') name) (drop 2 (dropWhile (/= ':') l))+      ("FAIL" : name : _) -> report ref "B" (takeWhile (/= ':') name) (Left (drop 2 (dropWhile (/= ':') l)))+      _ -> pure ()+  where+    byName = Map.fromList (map (\c -> (caseName c, c)) cases)+    checkFile name = case splitName name of+      Nothing -> report ref "B" name (Left "unexpected file name")+      Just (cn, variant, ext)+        | not (variantSupported variant) ->+            skip ref "B" name ("wireform-arrow built without -f+" ++ variant)+        | otherwise -> case Map.lookup cn byName of+            Nothing -> report ref "B" name (Left ("no Haskell-side case named " ++ cn))+            Just c -> do+              bytes <- BS.readFile (dir </> name)+              let decode = if ext == "arrow" then decodeArrowFile defaultReadOptions else decodeArrowStream defaultReadOptions+              r <- try (evaluate (force (decode bytes))) :: IO (Either SomeException (Either String (Schema, [V.Vector ColumnArray])))+              report ref "B" name $ case r of+                Left e -> Left ("wireform-arrow decoder threw: " ++ show e)+                Right (Left e) -> Left ("wireform-arrow failed to decode: " ++ e)+                Right (Right (sch, batches)) -> compareDecoded c sch batches+    dir = tmp </> "pyarrow"+++splitName :: String -> Maybe (String, String, String)+splitName name = case splitOn '.' name of+  [c, v, e] -> Just (c, v, e)+  _ -> Nothing+  where+    splitOn ch s = case break (== ch) s of+      (a, []) -> [a]+      (a, _ : rest) -> a : splitOn ch rest+++-- | Case names must agree between the two definitions of the matrix.+checkCaseNames :: IORef Tally -> FilePath -> FilePath -> IO ()+checkCaseNames ref python script = do+  out <- runPython python script ["cases"]+  case out of+    Left e -> report ref "matrix" "case list" (Left e)+    Right ls -> do+      let py = Set.fromList (filter (not . null) ls)+          hs = Set.fromList (map caseName cases)+      forM_ (Set.toList (Set.difference py hs)) $ \n ->+        report ref "matrix" n (Left "case exists in pyarrow_interop.py but not in Main.hs")+      forM_ (Set.toList (Set.difference hs py)) $ \n ->+        report ref "matrix" n (Left "case exists in Main.hs but not in pyarrow_interop.py")+      when (py == hs) $ putStrLn ("[matrix] " ++ show (Set.size hs) ++ " cases defined on both sides")+++findScript :: IO (Maybe FilePath)+findScript = do+  override <- lookupEnv "WIREFORM_ARROW_INTEROP_SCRIPT"+  let candidates =+        maybe [] pure override+          ++ ["test-interop/pyarrow_interop.py", "wireform-arrow/test-interop/pyarrow_interop.py"]+  firstExisting candidates+  where+    firstExisting [] = pure Nothing+    firstExisting (p : ps) = do+      exists <- doesFileExist p+      if exists then pure (Just p) else firstExisting ps+++main :: IO ()+main = do+  hSetEncoding stdout utf8+  args <- getArgs+  case args of+    [mode, dir] | mode == "--write" || mode == "--write-big" -> do+      (written, failures) <- writeHaskellFiles (if mode == "--write-big" then Big else Little) dir+      putStrLn ("wrote " ++ show (length written) ++ " files to " ++ dir)+      forM_ failures $ \(n, e) -> putStrLn ("FAIL " ++ n ++ ": " ++ e)+      if null failures then exitSuccess else exitFailure+    [] -> runSuite+    _ -> do+      putStrLn "usage: wireform-arrow-pyarrow-interop [--write DIR | --write-big DIR]"+      exitFailure+++runSuite :: IO ()+runSuite = do+  python <- fromMaybe "python3" <$> lookupEnv "WIREFORM_ARROW_PYTHON"+  required <- (== Just "1") <$> lookupEnv "WIREFORM_ARROW_REQUIRE_PYARROW"+  probe <- try (readProcessWithExitCode python ["-c", "import pyarrow; print(pyarrow.__version__)"] "") :: IO (Either SomeException (ExitCode, String, String))+  let unavailable why =+        if required+          then do+            putStrLn ("FAIL: pyarrow is required (WIREFORM_ARROW_REQUIRE_PYARROW=1) but unavailable: " ++ why)+            exitFailure+          else do+            putStrLn ("SKIP: pyarrow interop suite (" ++ why ++ "); set WIREFORM_ARROW_PYTHON to a Python with pyarrow, and WIREFORM_ARROW_REQUIRE_PYARROW=1 to make this a failure")+            exitSuccess+  version <- case probe of+    Left e -> unavailable ("cannot run " ++ python ++ ": " ++ show e)+    Right (ExitSuccess, out, _) -> pure (concat (take 1 (lines out)))+    Right (_, _, err) -> unavailable (python ++ " cannot import pyarrow: " ++ concat (take 1 (reverse (lines err))))+  script <-+    findScript >>= \case+      Just s -> pure s+      Nothing -> do+        putStrLn "FAIL: cannot find test-interop/pyarrow_interop.py (set WIREFORM_ARROW_INTEROP_SCRIPT)"+        exitFailure+  putStrLn ("pyarrow " ++ version ++ " via " ++ python ++ "; compression variants: " ++ unwords (map fst compressionVariants))+  ref <- newIORef (Tally 0 0 0 [])+  checkCaseNames ref python script+  withSystemTempDirectory "wireform-arrow-interop" $ \tmp -> do+    directionA Little ref python script tmp+    directionA Big ref python script tmp+    directionB ref python script tmp+  t <- readIORef ref+  putStrLn ""+  putStrLn ("pyarrow interop: " ++ show (tPass t) ++ " passed, " ++ show (tFail t) ++ " failed, " ++ show (tSkip t) ++ " skipped")+  unless (null (tFailures t)) $ do+    putStrLn "Failures:"+    mapM_ putStrLn (reverse (tFailures t))+    exitFailure
+ test-interop/pyarrow_interop.py view
@@ -0,0 +1,938 @@+#!/usr/bin/env python3+"""pyarrow side of the wireform-arrow interop test-suite.++The Haskell test-suite ``wireform-arrow-pyarrow-interop``+(``test-interop/Main.hs``) drives this script. Both sides define the+same case matrix independently: here every case is a list of+``pyarrow.RecordBatch`` values built with pyarrow's own constructors,+and ``Main.hs`` builds the same data as wireform-arrow ``ColumnArray``+values. Case names tie the two definitions together, and the Haskell+driver fails when the two name sets differ.++Subcommands (all output is line oriented so the Haskell driver can+parse it):++  cases          print every case name, one per line+  check DIR      for every DIR/<case>.<variant>.{arrows,arrow} written+                 by wireform-arrow: open it with ipc.open_stream or+                 ipc.open_file, run validate(full=True) on every batch,+                 and compare schema (including custom metadata) and+                 batch values against the expected batches built here.+                 Prints "PASS <file>" or "FAIL <file>: <reason>".+  write DIR      write every case with pyarrow, once per variant+                 (plain, sliced, zstd, lz4), as both stream and file+                 (stream only for replacement / delta dictionary+                 cases). Prints "WROTE <file>" or+                 "SKIP <case>.<variant>: <reason>".+  regen-goldens  regenerate test/golden/pa_*.arrows, the fixtures the+                 pyarrow-free wireform-arrow-test suite reads.++pyarrow's Python layer cannot materialise INTERVAL(YEAR_MONTH) or+INTERVAL(DAY_TIME) arrays (no Python array class exists for them), so+those cases are built through the Arrow C data interface with ctypes+and handled only through RecordBatch-level C++ operations (validate,+equals, IPC read/write, slice, concat).+"""++from __future__ import annotations++import ctypes+import decimal+import os+import struct+import sys+from pathlib import Path++import pyarrow as pa+import pyarrow.ipc as ipc++ROOT = Path(__file__).resolve().parent.parent+GOLDEN_DIR = ROOT / "test" / "golden"++D = decimal.Decimal+++# ---------------------------------------------------------------------------+# Arrow C data interface (only for the two interval types pyarrow's Python+# layer cannot represent).+# ---------------------------------------------------------------------------+++class _ArrowSchema(ctypes.Structure):+    pass+++class _ArrowArray(ctypes.Structure):+    pass+++_SREL = ctypes.CFUNCTYPE(None, ctypes.POINTER(_ArrowSchema))+_AREL = ctypes.CFUNCTYPE(None, ctypes.POINTER(_ArrowArray))+_ArrowSchema._fields_ = [+    ("format", ctypes.c_char_p),+    ("name", ctypes.c_char_p),+    ("metadata", ctypes.c_char_p),+    ("flags", ctypes.c_int64),+    ("n_children", ctypes.c_int64),+    ("children", ctypes.c_void_p),+    ("dictionary", ctypes.c_void_p),+    ("release", _SREL),+    ("private_data", ctypes.c_void_p),+]+_ArrowArray._fields_ = [+    ("length", ctypes.c_int64),+    ("null_count", ctypes.c_int64),+    ("offset", ctypes.c_int64),+    ("n_buffers", ctypes.c_int64),+    ("n_children", ctypes.c_int64),+    ("buffers", ctypes.POINTER(ctypes.c_void_p)),+    ("children", ctypes.POINTER(ctypes.POINTER(_ArrowArray))),+    ("dictionary", ctypes.c_void_p),+    ("release", _AREL),+    ("private_data", ctypes.c_void_p),+]++# Memory handed to pyarrow through the C interface. Never freed: the+# process exits through os._exit, so teardown order cannot run a+# release callback after its ctypes thunk is gone.+_KEEP: list = []+++@_SREL+def _schema_release(p):+    p.contents.release = _SREL()+++@_AREL+def _array_release(p):+    p.contents.release = _AREL()+++def _c_type(fmt: bytes) -> pa.DataType:+    s = _ArrowSchema(format=fmt, name=b"", flags=2, release=_schema_release)+    _KEEP.append(s)+    return pa.DataType._import_from_c(ctypes.addressof(s))+++MONTH_INTERVAL = _c_type(b"tiM")+DAY_TIME_INTERVAL = _c_type(b"tiD")+++def _raw_batch(schema: pa.Schema, length: int, cols) -> pa.RecordBatch:+    """Import a record batch whose columns are given as raw buffers.++    ``cols`` is a list of ``(null_count, [validity or None, data])``.+    """+    children = []+    for null_count, bufs in cols:+        ptrs = (ctypes.c_void_p * len(bufs))()+        for i, b in enumerate(bufs):+            if b is None:+                ptrs[i] = None+            else:+                cb = ctypes.create_string_buffer(b, len(b) + 8)+                _KEEP.append(cb)+                ptrs[i] = ctypes.addressof(cb)+        arr = _ArrowArray(+            length=length,+            null_count=null_count,+            offset=0,+            n_buffers=len(bufs),+            n_children=0,+            buffers=ptrs,+            release=_array_release,+        )+        _KEEP.extend([ptrs, arr])+        children.append(ctypes.pointer(arr))+    top_bufs = (ctypes.c_void_p * 1)(None)+    kids = (ctypes.POINTER(_ArrowArray) * len(children))(*children)+    top = _ArrowArray(+        length=length,+        null_count=0,+        offset=0,+        n_buffers=1,+        n_children=len(children),+        buffers=top_bufs,+        children=kids,+        release=_array_release,+    )+    _KEEP.extend([top_bufs, kids, top])+    return pa.RecordBatch._import_from_c(ctypes.addressof(top), schema)+++# ---------------------------------------------------------------------------+# Helpers+# ---------------------------------------------------------------------------+++def bitmap(valid: list[bool]) -> bytes:+    out = bytearray((len(valid) + 7) // 8)+    for i, v in enumerate(valid):+        if v:+            out[i // 8] |= 1 << (i % 8)+    return bytes(out)+++def no_nulls(xs: list) -> list:+    return [x for x in xs if x is not None]+++def float16_array(bits: list) -> pa.Array:+    n = len(bits)+    data = struct.pack("<" + "H" * n, *[0 if b is None else b for b in bits])+    nulls = sum(1 for b in bits if b is None)+    validity = pa.py_buffer(bitmap([b is not None for b in bits])) if nulls else None+    return pa.Array.from_buffers(pa.float16(), n, [validity, pa.py_buffer(data)], nulls)+++def mdn(t):+    return None if t is None else pa.MonthDayNano(list(t))+++def schema1(name: str, typ: pa.DataType, nullable: bool, metadata=None) -> pa.Schema:+    return pa.schema([pa.field(name, typ, nullable=nullable)], metadata=metadata)+++def batch1(schema: pa.Schema, arr: pa.Array) -> pa.RecordBatch:+    return pa.record_batch([arr], schema=schema)+++class Case:+    def __init__(self, name, schema, batches, stream_only=False, delta=False):+        self.name = name+        self.schema = schema+        self.batches = batches+        self.stream_only = stream_only+        self.delta = delta+++CASES: dict[str, Case] = {}+++def add(case: Case) -> None:+    if case.name in CASES:+        raise RuntimeError(f"duplicate case {case.name}")+    CASES[case.name] = case+++def scalar(name: str, typ: pa.DataType, b1: list, b2: list, mk=None) -> None:+    """Register ``name`` (non-nullable field, nulls dropped from b1) and+    ``name_nullable`` (nullable field, b1 with nulls, b2 without)."""+    build = mk if mk is not None else (lambda xs: pa.array(xs, type=typ))+    s = schema1("x", typ, False)+    add(Case(name, s, [batch1(s, build(no_nulls(b1))), batch1(s, build(b2))]))+    sn = schema1("x", typ, True)+    add(Case(name + "_nullable", sn, [batch1(sn, build(b1)), batch1(sn, build(b2))]))+++def raw_scalar(name: str, typ: pa.DataType, fmt: str, b1: list, b2: list) -> None:+    """Like ``scalar`` for types that only exist at the C++ layer."""++    def flat(x):+        return list(x) if isinstance(x, tuple) else [x]++    def mk(schema, xs):+        width = struct.calcsize("<" + fmt)+        data = b"".join(+            b"\x00" * width if x is None else struct.pack("<" + fmt, *flat(x)) for x in xs+        )+        nulls = sum(1 for x in xs if x is None)+        validity = bitmap([x is not None for x in xs]) if nulls else None+        return _raw_batch(schema, len(xs), [(nulls, [validity, data])])++    s = schema1("x", typ, False)+    add(Case(name, s, [mk(s, no_nulls(b1)), mk(s, b2)]))+    sn = schema1("x", typ, True)+    add(Case(name + "_nullable", sn, [mk(sn, b1), mk(sn, b2)]))+++# ---------------------------------------------------------------------------+# The case matrix. Keep in sync with test-interop/Main.hs.+# ---------------------------------------------------------------------------+++def build_cases() -> None:+    I64MAX = 9223372036854775807+    scalar("int8", pa.int8(), [0, 1, None, -1, 127, -128], [42, -42])+    scalar("int16", pa.int16(), [0, 1, None, -1, 32767, -32768], [1000, -1000])+    scalar("int32", pa.int32(), [0, 1, None, -1, 2147483647, -2147483648], [7, -7])+    scalar("int64", pa.int64(), [0, 1, None, -1, I64MAX, -I64MAX - 1], [123456789012, -5])+    scalar("uint8", pa.uint8(), [0, 1, None, 255], [128])+    scalar("uint16", pa.uint16(), [0, 1, None, 65535], [256])+    scalar("uint32", pa.uint32(), [0, 1, None, 4294967295], [65536])+    scalar("uint64", pa.uint64(), [0, 1, None, 18446744073709551615], [4294967296])+    scalar(+        "float16",+        pa.float16(),+        [0x0000, 0x3C00, None, 0xC000, 0x7BFF, 0x7C00],+        [0x3800],+        mk=float16_array,+    )+    inf = float("inf")+    scalar("float32", pa.float32(), [0.0, 1.5, None, -2.5, inf, 3.4028234663852886e38], [0.25])+    scalar("float64", pa.float64(), [0.0, 1.5, None, -2.5, -inf, 1e300, 5e-324], [3.141592653589793])+    scalar(+        "bool",+        pa.bool_(),+        [True, False, None, True, True, False, False, True, False],+        [False, True],+    )+    strs = ["", "a", None, "h\u00e9llo w\u00f6rld", "\U0001F600 emoji"]+    scalar("utf8", pa.utf8(), strs, ["second batch"])+    scalar("large_utf8", pa.large_utf8(), strs, ["second batch"])+    bins = [b"", b"\x00\x01\x02", None, b"\xff" * 20]+    scalar("binary", pa.binary(), bins, [b"z"])+    scalar("large_binary", pa.large_binary(), bins, [b"z"])+    scalar(+        "fixed_size_binary4",+        pa.binary(4),+        [b"abcd", b"\x00\x00\x00\x00", None, b"\xff\xfe\xfd\xfc"],+        [b"wxyz"],+    )+    views = [+        "",+        "short",+        None,+        "exactly12byt",+        "thirteen byte",+        "this string is definitely longer than twelve bytes",+    ]+    scalar("utf8_view", pa.string_view(), views, ["inline", "an out of line string value"])+    scalar(+        "binary_view",+        pa.binary_view(),+        [None if v is None else v.encode() for v in views],+        [b"inline", b"\x00" * 13],+    )+    scalar("date32", pa.date32(), [0, 18000, None, -1, 19000], [20000])+    scalar("date64", pa.date64(), [0, 1555200000000, None, -86400000], [86400000])+    scalar("time32_s", pa.time32("s"), [0, 3600, None, 86399], [60])+    scalar("time32_ms", pa.time32("ms"), [0, 1000, None, 86399999], [1])+    scalar("time64_us", pa.time64("us"), [0, 12345000000, None, 86399999999], [1])+    scalar("time64_ns", pa.time64("ns"), [0, 1, None, 86399999999999], [2])+    for unit in ["s", "ms", "us", "ns"]:+        ts = [0, 1700000000, None, -1]+        scalar(f"timestamp_{unit}", pa.timestamp(unit), ts, [42])+        scalar(f"timestamp_{unit}_utc", pa.timestamp(unit, tz="UTC"), ts, [42])+        scalar(f"duration_{unit}", pa.duration(unit), [0, 60, None, -1], [1])+    scalar(+        "timestamp_us_new_york",+        pa.timestamp("us", tz="America/New_York"),+        [0, 1700000000000000, None],+        [5],+    )+    scalar(+        "decimal32_9_2", pa.decimal32(9, 2),+        [D("-9999999.99"), D("1.23"), None, D("0.00"), D("9999999.99")], [D("0.01")],+    )+    scalar(+        "decimal64_18_2", pa.decimal64(18, 2),+        [D("-9999999999999999.99"), D("1.23"), None, D("0.00"), D("9999999999999999.99")], [D("0.01")],+    )+    scalar(+        "decimal128_18_2",+        pa.decimal128(18, 2),+        [D("0.00"), D("1.23"), None, D("-1.23"), D("9999999999999999.99")],+        [D("0.01")],+    )+    scalar(+        "decimal128_38_10",+        pa.decimal128(38, 10),+        [D("12345678901234567890.1234567890"), None, D("-0.0000000001")],+        [D("1.0000000000")],+    )+    scalar(+        "decimal256_76_5",+        pa.decimal256(76, 5),+        [+            D("1.00000"),+            None,+            D("-12345678901234567890123456789012345678901234567890.12345"),+        ],+        [D("0.00001")],+    )+    scalar(+        "interval_month_day_nano",+        pa.month_day_nano_interval(),+        [mdn((1, 2, 3)), mdn((0, 0, 0)), None, mdn((-1, -2, -3))],+        [mdn((12, 30, 1000000000))],+    )+    raw_scalar("interval_year_month", MONTH_INTERVAL, "i", [0, 14, None, -3], [5])+    raw_scalar(+        "interval_day_time",+        DAY_TIME_INTERVAL,+        "ii",+        [(0, 0), (1, 1000), None, (-1, -1)],+        [(2, 2)],+    )++    # Null type: every slot is null, no buffers.+    s = schema1("x", pa.null(), True)+    add(Case("null", s, [batch1(s, pa.nulls(4)), batch1(s, pa.nulls(1))]))++    # Struct<i: int32 not null, s: utf8>.+    st = pa.struct([pa.field("i", pa.int32(), nullable=False), pa.field("s", pa.utf8())])+    rows = [{"i": 1, "s": "a"}, {"i": 2, "s": None}, {"i": 3, "s": "c"}]+    scalar("struct", st, rows[:1] + [None] + rows[1:], [{"i": 4, "s": "d"}])+    # A struct with no fields: only the row count (and validity) is on the wire.+    scalar("struct_empty", pa.struct([]), [{}, None, {}, {}], [{}])++    scalar(+        "list_int32",+        pa.list_(pa.int32()),+        [[1, 2], None, [], [3, None, 5]],+        [[6], [7, 8]],+    )+    scalar(+        "large_list_utf8",+        pa.large_list(pa.utf8()),+        [["a", "b"], None, [], ["c"]],+        [["d"]],+    )+    scalar(+        "fixed_size_list_int16_3",+        pa.list_(pa.int16(), 3),+        [[1, 2, 3], None, [4, None, 6]],+        [[7, 8, 9]],+    )+    # fixed_size_list<int32, 0>: every row is empty, the child has no elements.+    scalar("fixed_size_list_int32_0", pa.list_(pa.int32(), 0), [[], None, [], []], [[], []])+    scalar(+        "map_utf8_int32",+        pa.map_(pa.utf8(), pa.int32()),+        [[("a", 1), ("b", None)], None, [], [("c", 3)]],+        [[("d", 4)]],+    )++    def list_view(large: bool):+        typ = pa.large_list_view(pa.int32()) if large else pa.list_view(pa.int32())+        it = pa.int64() if large else pa.int32()+        child = pa.array([10, 20, 30, None, 50, 60], pa.int32())++        def mk(offs, sizes, valid, values=child):+            mask = None if valid is None else pa.array([not v for v in valid], pa.bool_())+            cls = pa.LargeListViewArray if large else pa.ListViewArray+            return cls.from_arrays(+                pa.array(offs, it), pa.array(sizes, it), values, type=typ, mask=mask+            )++        name = "large_list_view_int32" if large else "list_view_int32"+        b2 = mk([0], [1], None, pa.array([99], pa.int32()))+        s = schema1("x", typ, False)+        add(Case(name, s, [batch1(s, mk([4, 0, 1, 0], [2, 3, 0, 1], None)), batch1(s, b2)]))+        sn = schema1("x", typ, True)+        add(+            Case(+                name + "_nullable",+                sn,+                [+                    batch1(sn, mk([4, 0, 1, 0, 2], [2, 0, 0, 1, 2], [True, False, True, True, True])),+                    batch1(sn, b2),+                ],+            )+        )++    list_view(False)+    list_view(True)++    # Unions. Children are nullable; the union itself has no validity.+    def union(name: str, dense: bool, codes: list[int], nullable_field: bool):+        names = ["i", "s"]+        if dense:+            def mk(types, offsets, ints, strs):+                return pa.UnionArray.from_dense(+                    pa.array(types, pa.int8()),+                    pa.array(offsets, pa.int32()),+                    [pa.array(ints, pa.int32()), pa.array(strs, pa.utf8())],+                    names,+                    codes,+                )++            b1 = mk(+                [codes[0], codes[1], codes[0], codes[0], codes[1]],+                [0, 0, 1, 2, 1],+                [1, 2, None],+                ["a", "b"],+            )+            b2 = mk([codes[1]], [0], [], ["z"])+        else:+            def mk(types, ints, strs):+                return pa.UnionArray.from_sparse(+                    pa.array(types, pa.int8()),+                    [pa.array(ints, pa.int32()), pa.array(strs, pa.utf8())],+                    names,+                    codes,+                )++            b1 = mk(+                [codes[0], codes[1], codes[0], codes[1]],+                [1, 0, None, 0],+                [None, "a", "x", "b"],+            )+            b2 = mk([codes[1]], [0], ["z"])+        s = schema1("x", b1.type, nullable_field)+        add(Case(name, s, [batch1(s, b1), batch1(s, b2)]))++    union("dense_union", True, [0, 1], False)+    union("dense_union_nullable", True, [0, 1], True)+    union("sparse_union", False, [0, 1], False)+    union("sparse_union_nullable", False, [0, 1], True)+    union("dense_union_type_codes", True, [3, 7], True)+    union("sparse_union_type_codes", False, [5, 2], True)++    # Dictionaries.+    def dict_arr(indices, values, index_type=pa.int32()):+        return pa.DictionaryArray.from_arrays(pa.array(indices, index_type), values)++    abc = pa.array(["a", "b", "c"], pa.utf8())+    dt = pa.dictionary(pa.int32(), pa.utf8())+    s = schema1("x", dt, False)+    add(Case("dict_utf8", s, [batch1(s, dict_arr([0, 1, 0, 2, 1], abc)), batch1(s, dict_arr([2, 2, 0], abc))]))+    sn = schema1("x", dt, True)+    add(+        Case(+            "dict_utf8_nullable",+            sn,+            [batch1(sn, dict_arr([0, None, 1, 2], abc)), batch1(sn, dict_arr([1, 0], abc))],+        )+    )+    xy = pa.array(["x", "y"], pa.utf8())+    s8 = schema1("x", pa.dictionary(pa.int8(), pa.utf8()), False)+    add(+        Case(+            "dict_int8_index",+            s8,+            [batch1(s8, dict_arr([1, 0, 1], xy, pa.int8())), batch1(s8, dict_arr([0], xy, pa.int8()))],+        )+    )+    ints = pa.array([100, 200], pa.int64())+    si = schema1("x", pa.dictionary(pa.int32(), pa.int64()), False)+    add(Case("dict_int64_values", si, [batch1(si, dict_arr([1, 0, 1], ints)), batch1(si, dict_arr([0], ints))]))+    sr = schema1("x", dt, False)+    add(+        Case(+            "dict_replacement",+            sr,+            [+                batch1(sr, dict_arr([0, 1, 1], pa.array(["a", "b"]))),+                batch1(sr, dict_arr([2, 0, 1], pa.array(["x", "y", "z"]))),+            ],+            stream_only=True,+        )+    )+    add(+        Case(+            "dict_delta",+            sr,+            [+                batch1(sr, dict_arr([0, 1], pa.array(["a", "b"]))),+                batch1(sr, dict_arr([2, 0], pa.array(["a", "b", "c"]))),+            ],+            stream_only=True,+            delta=True,+        )+    )+    nested_t = pa.struct([pa.field("d", dt)])+    sd = schema1("x", nested_t, False)+    add(+        Case(+            "dict_in_struct",+            sd,+            [+                batch1(sd, pa.StructArray.from_arrays([dict_arr([2, 0, 1], abc)], fields=[pa.field("d", dt)])),+                batch1(sd, pa.StructArray.from_arrays([dict_arr([1], abc)], fields=[pa.field("d", dt)])),+            ],+        )+    )++    # Dictionaries nested in dictionary values. pyarrow numbers dictionary+    # ids in schema pre-order: the outer dictionary is 0, the inner one 1.+    inner_t = pa.dictionary(pa.int16(), pa.utf8())++    def struct_of_dict(indices, inner_indices, inner_values):+        values = pa.StructArray.from_arrays(+            [dict_arr(inner_indices, inner_values, pa.int16())], fields=[pa.field("d", inner_t)]+        )+        return dict_arr(indices, values, pa.int8())++    so = schema1("x", pa.dictionary(pa.int8(), pa.struct([pa.field("d", inner_t)])), False)+    add(+        Case(+            "dict_struct_of_dict",+            so,+            [+                batch1(so, struct_of_dict([2, 0, 1, 0], [1, 0, 2], abc)),+                batch1(so, struct_of_dict([1], [1, 0, 2], abc)),+            ],+        )+    )+    list_values = pa.ListArray.from_arrays(pa.array([0, 2, 2, 3], pa.int32()), dict_arr([0, 1, 1], xy, pa.int8()))+    sl = schema1("x", pa.dictionary(pa.int32(), list_values.type), True)+    add(+        Case(+            "dict_list_of_dict",+            sl,+            [batch1(sl, dict_arr([1, None, 0, 2], list_values)), batch1(sl, dict_arr([2], list_values))],+        )+    )+    # The inner dictionary changes between batches, so both are replaced.+    add(+        Case(+            "dict_nested_replacement",+            so,+            [batch1(so, struct_of_dict([0, 1], [0, 1], abc)), batch1(so, struct_of_dict([1, 0], [0, 0], xy))],+            stream_only=True,+        )+    )+    # A nullable dictionary column whose rows are all null may reference an+    # empty dictionary.+    empty = pa.array([], pa.utf8())+    add(+        Case(+            "dict_all_null_empty",+            sn,+            [batch1(sn, dict_arr([None, None, None], empty)), batch1(sn, dict_arr([None], empty))],+        )+    )++    # Run-end encoded.+    def ree(name, re_type, val_type, b1, b2, nullable):+        typ = pa.run_end_encoded(re_type, val_type)++        def mk(spec):+            ends, vals = spec+            return pa.RunEndEncodedArray.from_arrays(+                pa.array(ends, re_type), pa.array(vals, val_type), type=typ+            )++        s = schema1("x", typ, nullable)+        add(Case(name, s, [batch1(s, mk(b1)), batch1(s, mk(b2))]))++    ree("ree_int32_int64", pa.int32(), pa.int64(), ([3, 5, 8], [100, 200, 300]), ([2], [7]), False)+    ree("ree_int32_int64_nullable", pa.int32(), pa.int64(), ([3, 5, 8], [100, None, 300]), ([2], [7]), True)+    ree("ree_int16_utf8", pa.int16(), pa.utf8(), ([1, 4], ["a", "bb"]), ([3], ["c"]), True)+    ree("ree_int64_float64", pa.int64(), pa.float64(), ([2, 3], [1.5, None]), ([1], [2.5]), True)++    # Custom schema- and field-level metadata.+    ms = pa.schema(+        [+            pa.field("a", pa.int32(), nullable=False, metadata={"unit": "ms", "note": "field level"}),+            pa.field("b", pa.utf8(), metadata={"k": "v"}),+        ],+        metadata={"origin": "wireform-arrow interop", "k2": "v2"},+    )+    add(+        Case(+            "custom_metadata",+            ms,+            [+                pa.record_batch([pa.array([1, 2], pa.int32()), pa.array(["x", None])], schema=ms),+                pa.record_batch([pa.array([3], pa.int32()), pa.array(["y"])], schema=ms),+            ],+        )+    )++    # Multi-column batches, including zero-row batches.+    mixed_s = pa.schema(+        [+            pa.field("i", pa.int64(), nullable=False),+            pa.field("s", pa.utf8()),+            pa.field("b", pa.bool_()),+            pa.field("l", pa.list_(pa.int32())),+            pa.field("st", pa.struct([pa.field("f", pa.float64())])),+            pa.field("d", dt),+            pa.field("v", pa.string_view()),+        ]+    )++    def mixed(i, s_, b, l, f, d, v):+        return pa.record_batch(+            [+                pa.array(i, pa.int64()),+                pa.array(s_, pa.utf8()),+                pa.array(b, pa.bool_()),+                pa.array(l, pa.list_(pa.int32())),+                pa.array([None if x is None else {"f": x} for x in f], mixed_s.field("st").type),+                dict_arr(d, abc),+                pa.array(v, pa.string_view()),+            ],+            schema=mixed_s,+        )++    m1 = mixed(+        [10, 20, 30],+        ["hello", None, "!"],+        [True, None, False],+        [[1], None, [2, 3]],+        [1.5, None, 2.5],+        [0, 2, None],+        ["v", None, "a view longer than twelve"],+    )+    m0 = mixed([], [], [], [], [], [], [])+    m2 = mixed([40], ["w"], [True], [[]], [0.5], [1], ["z"])+    add(Case("mixed", mixed_s, [m1, m2]))+    add(Case("zero_row_batches", mixed_s, [m1, m0, m2]))+    add(Case("only_zero_row_batch", mixed_s, [m0]))+    add(Case("no_batches", mixed_s, []))+++# ---------------------------------------------------------------------------+# check: pyarrow reads what wireform-arrow wrote+# ---------------------------------------------------------------------------+++def one_line(msg) -> str:+    return " ".join(str(msg).split())+++def column_diff(got: pa.RecordBatch, want: pa.RecordBatch) -> str:+    parts = []+    for i in range(want.num_columns):+        name = want.schema.field(i).name+        try:+            g, w = got.column(i), want.column(i)+        except Exception:  # interval types have no Python array class+            parts.append(f"column {name} differs (no Python view of this type)")+            continue+        if not g.equals(w):+            try:+                parts.append(f"column {name}: got {g.to_pylist()!r} want {w.to_pylist()!r}")+            except Exception as e:  # noqa: BLE001+                parts.append(f"column {name}: got {g!r} want {w!r} ({e})")+    return "; ".join(parts) or "batches differ"+++def compare(case: Case, schema: pa.Schema, batches: list[pa.RecordBatch]) -> str | None:+    if not schema.equals(case.schema, check_metadata=True):+        return f"schema mismatch: got {schema!r} (metadata {schema.metadata}) want {case.schema!r} (metadata {case.schema.metadata})"+    if len(batches) != len(case.batches):+        return f"got {len(batches)} batches, want {len(case.batches)}"+    for n, (got, want) in enumerate(zip(batches, case.batches)):+        try:+            got.validate(full=True)+        except Exception as e:  # noqa: BLE001+            return f"batch {n}: validate(full=True) failed: {e}"+        if got.num_rows != want.num_rows:+            return f"batch {n}: got {got.num_rows} rows, want {want.num_rows}"+        if not got.equals(want):+            return f"batch {n}: {column_diff(got, want)}"+    return None+++def _swap_words(buffer, width):+    """Independent fixture adapter: Arrow body words, metadata stays LE."""+    if buffer is None or width == 1:+        return buffer+    data = bytes(buffer)+    if len(data) % width:+        raise ValueError(f"partial {width}-byte numeric word")+    return pa.py_buffer(b"".join(data[i:i + width][::-1] for i in range(0, len(data), width)))+++def _native_from_big(array):+    """Normalize an IPC array read with ensure_native_endian=False.++    Arrow C++ 26 cannot swap views or REE itself. This independent Python+    implementation converts their raw buffers (and nested values), then the+    normal C++ full validator and semantic comparator check the result.+    No expected values or wireform code participate in the conversion.+    The selected fixtures contain views/REE, plain scalars, strings, lists,+    structs and dictionaries. Unsupported adapter types fail explicitly.+    """+    typ = array.type+    if pa.types.is_dictionary(typ):+        return pa.DictionaryArray.from_arrays(_native_from_big(array.indices), _native_from_big(array.dictionary), ordered=typ.ordered)+    if pa.types.is_run_end_encoded(typ):+        return pa.RunEndEncodedArray.from_arrays(_native_from_big(array.run_ends), _native_from_big(array.values), type=typ)+    buffers = array.buffers()[:typ.num_buffers]+    children = []+    if pa.types.is_string_view(typ) or pa.types.is_binary_view(typ):+        buffers = array.buffers()+        data = bytearray(buffers[1])+        for off in range(0, len(data), 16):+            length = int.from_bytes(data[off:off + 4], "big", signed=True)+            data[off:off + 4] = data[off:off + 4][::-1]+            if length > 12:+                data[off + 8:off + 12] = data[off + 8:off + 12][::-1]+                data[off + 12:off + 16] = data[off + 12:off + 16][::-1]+        buffers[1] = pa.py_buffer(data)+    elif pa.types.is_string(typ) or pa.types.is_binary(typ) or pa.types.is_large_string(typ) or pa.types.is_large_binary(typ):+        buffers[1] = _swap_words(buffers[1], 8 if pa.types.is_large_string(typ) or pa.types.is_large_binary(typ) else 4)+    elif pa.types.is_struct(typ):+        children = [_native_from_big(array.field(i)) for i in range(typ.num_fields)]+    elif pa.types.is_list(typ) or pa.types.is_large_list(typ):+        buffers[1] = _swap_words(buffers[1], 8 if pa.types.is_large_list(typ) else 4)+        children = [_native_from_big(array.values)]+    elif pa.types.is_fixed_size_list(typ):+        children = [_native_from_big(array.values)]+    elif pa.types.is_boolean(typ) or pa.types.is_null(typ) or pa.types.is_fixed_size_binary(typ):+        pass+    elif pa.types.is_integer(typ) or pa.types.is_floating(typ) or pa.types.is_decimal(typ) or (pa.types.is_temporal(typ) and not pa.types.is_interval(typ)):+        buffers[1] = _swap_words(buffers[1], typ.bit_width // 8)+    else:+        raise NotImplementedError(f"fixture endian adapter for {typ}")+    return pa.Array.from_buffers(typ, len(array), buffers, null_count=array.null_count, offset=array.offset, children=children)+++def _read_ipc(path, options=None):+    source = pa.OSFile(str(path), "rb")+    if path.suffix == ".arrows":+        reader = ipc.open_stream(source, options=options)+        return reader.schema, list(reader)+    reader = ipc.open_file(source, options=options)+    return reader.schema, [reader.get_batch(i) for i in range(reader.num_record_batches)]+++def check(directory: Path) -> int:+    for path in sorted(directory.iterdir()):+        parts = path.name.split(".")+        if len(parts) != 3 or parts[2] not in ("arrows", "arrow"):+            continue+        case = CASES.get(parts[0])+        if case is None:+            print(f"FAIL {path.name}: no pyarrow-side case named {parts[0]}")+            continue+        try:+            try:+                schema, batches = _read_ipc(path)+            except pa.ArrowNotImplementedError as e:+                # These are explicit reference capability errors, not corrupt+                # input or value mismatches. All other failures remain failures.+                known = ("Swapping endianness of string_view", "Swapping endianness of binary_view", "swapping endianness of run-end encoded array")+                if str(e) not in known:+                    raise+                schema, raw_batches = _read_ipc(path, ipc.IpcReadOptions(ensure_native_endian=False))+                schema = pa.schema(list(schema), metadata=schema.metadata)+                batches = [pa.record_batch([_native_from_big(b.column(i)) for i in range(b.num_columns)], schema=schema) for b in raw_batches]+                print(f"ADAPTER {path.name}: PyArrow native limitation: {e}; independent raw-buffer byte swap + full C++ validation")+            err = compare(case, schema, batches)+        except Exception as e:  # noqa: BLE001+            err = f"pyarrow failed to read: {type(e).__name__}: {e}"+        if err is None:+            print(f"PASS {path.name}")+        else:+            print(f"FAIL {path.name}: {one_line(err)}")+    return 0+++# ---------------------------------------------------------------------------+# write: wireform-arrow reads what pyarrow wrote+# ---------------------------------------------------------------------------+++def sliced(b: pa.RecordBatch) -> pa.RecordBatch:+    """The same rows as ``b``, but as a slice at a non-zero offset of a+    larger batch, so pyarrow has to truncate / rebase buffers when it+    writes them."""+    big = pa.concat_batches([b, b, b])+    return big.slice(b.num_rows, b.num_rows)+++VARIANTS = {+    "plain": None,+    "sliced": None,+    "zstd": "zstd",+    "lz4": "lz4",+}+++def write(directory: Path) -> int:+    directory.mkdir(parents=True, exist_ok=True)+    for case in CASES.values():+        for variant, codec in VARIANTS.items():+            try:+                batches = case.batches+                if variant == "sliced":+                    batches = [sliced(b) for b in batches]+                opts = ipc.IpcWriteOptions(+                    compression=codec, emit_dictionary_deltas=case.delta+                )+            except Exception as e:  # noqa: BLE001+                print(f"SKIP {case.name}.{variant}: {one_line(e)}")+                continue+            kinds = [("arrows", ipc.new_stream)]+            if not case.stream_only:+                kinds.append(("arrow", ipc.new_file))+            for ext, new in kinds:+                path = directory / f"{case.name}.{variant}.{ext}"+                try:+                    with pa.OSFile(str(path), "wb") as sink:+                        with new(sink, case.schema, options=opts) as w:+                            for b in batches:+                                w.write_batch(b)+                except Exception as e:  # noqa: BLE001+                    print(f"FAIL {path.name}: pyarrow failed to write: {one_line(e)}")+                    continue+                print(f"WROTE {path.name}")+    return 0+++# ---------------------------------------------------------------------------+# regen-goldens: fixtures for the pyarrow-free wireform-arrow-test suite+# ---------------------------------------------------------------------------+++def regen_goldens() -> int:+    GOLDEN_DIR.mkdir(parents=True, exist_ok=True)++    sch = pa.schema([pa.field("a", pa.int32(), nullable=False)])+    batch = pa.record_batch([pa.array([1, 2, 3, 4, 5], pa.int32())], schema=sch)+    with pa.OSFile(str(GOLDEN_DIR / "pa_int32.arrows"), "wb") as f:+        with ipc.new_stream(f, sch) as w:+            w.write_batch(batch)++    sch2 = pa.schema(+        [+            pa.field("i", pa.int64(), nullable=False),+            pa.field("s", pa.string(), nullable=True),+            pa.field("b", pa.bool_(), nullable=True),+        ]+    )+    batch2 = pa.record_batch(+        [+            pa.array([10, 20, 30], pa.int64()),+            pa.array(["alpha", None, "gamma"], pa.string()),+            pa.array([True, False, None], pa.bool_()),+        ],+        schema=sch2,+    )+    with pa.OSFile(str(GOLDEN_DIR / "pa_mixed.arrows"), "wb") as f:+        with ipc.new_stream(f, sch2) as w:+            w.write_batch(batch2)++    dict_arr = pa.DictionaryArray.from_arrays(+        pa.array([0, 1, 0, 2, 1], pa.int32()), pa.array(["a", "b", "c"], pa.string())+    )+    sch3 = pa.schema([pa.field("d", dict_arr.type)])+    batch3 = pa.record_batch([dict_arr], schema=sch3)+    with pa.OSFile(str(GOLDEN_DIR / "pa_dict.arrows"), "wb") as f:+        with ipc.new_stream(f, sch3) as w:+            w.write_batch(batch3)++    print(f"Regenerated goldens under {GOLDEN_DIR.relative_to(ROOT)}/")+    return 0+++def main(argv: list[str]) -> int:+    cmd = argv[1] if len(argv) > 1 else ""+    if cmd == "regen-goldens":+        return regen_goldens()+    build_cases()+    if cmd == "cases":+        for name in CASES:+            print(name)+        return 0+    if cmd == "check" and len(argv) == 3:+        return check(Path(argv[2]))+    if cmd == "write" and len(argv) == 3:+        return write(Path(argv[2]))+    print(__doc__, file=sys.stderr)+    return 2+++if __name__ == "__main__":+    rc = main(sys.argv)+    sys.stdout.flush()+    sys.stderr.flush()+    # Skip interpreter teardown: see _KEEP.+    os._exit(rc)
+ test/Main.hs view
@@ -0,0 +1,1634 @@+{-# LANGUAGE OverloadedStrings #-}++{- | Round-trip tests for the Arrow IPC writer + reader.++Two layers of coverage:++  1. Internal round-trips: every 'ColumnArray' constructor+     the writer emits must be parsed back correctly by+     'decodeArrowStream' + 'decodeRecordBatch'.+  2. Golden pyarrow interop: the bytes in+     @test/golden/pa_*.arrows@ were produced by pyarrow+     ('pa.ipc.new_stream') against the reference spec.+     'decodeArrowStream' must accept them verbatim.++Previously (PR #16 deferred list) the wireform IPC framing+was a simplified encoding that pyarrow couldn't read; after+'Arrow.FlatBufferIPC' landed with a real FlatBuffers layout+the bidirectional interop works and these tests pin it.+-}+module Main (main) where++import Arrow.Column+import Arrow.File (asBatches, asSchema, readArrowStream)+import Arrow.Read.Columns (defaultReadOptions)+import Arrow.FlatBufferIPC (+  SparseTensor (..),+  Tensor (..),+  TensorDim (..),+  buildSchemaMessage,+  decodeSchemaMessage,+  decodeSparseTensorFrame,+  decodeTensorFrame,+  encodeSparseTensorFrame,+  encodeTensorFrame,+  decodeRecordBatch,+ )+import Arrow.Record qualified as AR+import Arrow.Stream (+  DictHandling (..),+  WriteOptions (..),+  bodyCompressionAvailable,+  decodeArrowStream,+  defaultWriteOptions,+  encodeArrowStream,+  openStreamReader,+  streamReaderIter,+  streamReaderNext,+  streamReaderProjected,+  streamReaderSchema,+  streamReaderToList,+ )+import Arrow.Types+import Arrow.Write (writeArrowStream)+import Columnar.Stream qualified as IS+import Control.Monad (unless, when)+import Data.ByteString qualified as BS+import Data.Int (Int32, Int64)+import Data.Maybe (listToMaybe)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Vector qualified as V+import Data.Vector.Generic qualified as VG+import Data.Vector.Storable qualified as VS+import Data.Word (Word16, Word32, Word64, Word8)+import Test.Arrow.Core qualified as Core+import Test.Arrow.Decimal qualified as Decimal+import Test.Arrow.Features qualified as Features+import Test.Arrow.Malformed qualified as Malformed+import Test.Arrow.Props qualified as Props+import Test.Arrow.Sections qualified as Sections+import Test.Arrow.Vector qualified as Vector+import System.Exit (exitFailure)+++-- Test records used by 'nestedStructRoundTrip'.+data Address = Address+  { cityF :: Text+  , zipF :: Text+  }+  deriving (Show, Eq)+++data Customer = Customer+  { nameF :: Text+  , addrF :: Address+  , ageF :: Int32+  }+  deriving (Show, Eq)+++{- | Customer with an optional address; exercises the+'encoderFromRowEncoder' / 'decoderFromRowDecoder' nullable+nested-struct path.+-}+data CustomerOpt = CustomerOpt+  { coNameF :: Text+  , maybeAddrF :: Maybe Address+  , coAgeF :: Int32+  }+  deriving (Show, Eq)+++main :: IO ()+main = do+  putStrLn "wireform-arrow writer/reader round-trip suite"+  Decimal.tests+  Features.tests++  -- Every column: build a single-batch Arrow stream containing just+  -- that column, serialise it, parse it, materialise, and assert+  -- the recovered ColumnArray equals the one we put in.++  roundTripPrim "Int8" (primColumn PInt8 (VS.fromList [0, 1, -1, 100, -128, 127]))+  roundTripPrim "Int16" (primColumn PInt16 (VS.fromList [0, 1, -1, 32767, -32768]))+  roundTripPrim "Int32" (primColumn PInt32 (VS.fromList [0, 1, -1, maxBound, minBound]))+  roundTripPrim "Int64" (primColumn PInt64 (VS.fromList [0, 1, -1, maxBound, minBound]))+  roundTripPrim "UInt8" (primColumn PUInt8 (VS.fromList ([0, 255] :: [Word8])))+  roundTripPrim "UInt16" (primColumn PUInt16 (VS.fromList ([0, 65535] :: [Word16])))+  roundTripPrim "UInt32" (primColumn PUInt32 (VS.fromList ([0, maxBound] :: [Word32])))+  roundTripPrim "UInt64" (primColumn PUInt64 (VS.fromList ([0, maxBound] :: [Word64])))+  roundTripPrim+    "Float16"+    (primColumn PFloat16 (VS.fromList [Float16 0, Float16 0x3C00, Float16 0xBC00]))+  roundTripPrim "Float" (primColumn PFloat (VS.fromList [0.0, 1.5, -2.25, 3.14 :: Float]))+  roundTripPrim "Double" (primColumn PDouble (VS.fromList [0.0, 1.5, -2.25, 3.14159265 :: Double]))+  roundTripPrim "Bool" (fromBools (V.fromList [True, False, True, False, True]))++  roundTripPrim "Date32" (primColumn PDate32 (VS.fromList [0 :: Int32, 18000, -1]))+  roundTripPrim "Date64" (primColumn PDate64 (VS.fromList [0 :: Int64, 1700000000000]))+  roundTripPrim "Time32" (primColumn PTime32 (VS.fromList [0 :: Int32, 12345]))+  roundTripPrim "Time64" (primColumn PTime64 (VS.fromList [0 :: Int64, 12345000000]))+  roundTripPrim+    "Timestamp"+    (primColumn PTimestamp (VS.fromList [0 :: Int64, 1700000000_000_000_000]))+  roundTripPrim "Duration" (primColumn PDuration (VS.fromList [0 :: Int64, 60_000_000_000]))++  roundTripPrim+    "IntervalYearMonth"+    (primColumn PIntervalYearMonth (VS.fromList [0 :: Int32, 12, -6, 100]))+  roundTripPrim+    "IntervalDayTime"+    ( primColumn+        PIntervalDayTime+        ( VS.fromList+            [IntervalDayTime 1 500, IntervalDayTime 2 (-1), IntervalDayTime 30 0]+        )+    )+  roundTripPrim+    "IntervalMonthDayNano"+    ( primColumn+        PIntervalMonthDayNano+        ( VS.fromList+            [IntervalMonthDayNano 1 3 1000, IntervalMonthDayNano 2 4 (-500)]+        )+    )++  roundTripPrim "Decimal32" (primColumn (PDecimal32 9 2) (VS.fromList [-999999999,0,123,999999999]))+  roundTripPrim "Decimal64" (primColumn (PDecimal64 18 4) (VS.fromList [-999999999999999999,0,123,999999999999999999]))++  roundTripPrim+    "Decimal128"+    ( primColumn+        (PDecimal128 18 2)+        (VS.fromList [decimal128FromInteger 0, decimal128FromInteger 100])+    )+  roundTripPrim+    "Decimal256"+    ( primColumn+        (PDecimal256 38 4)+        (VS.fromList [decimal256FromInteger 0, decimal256FromInteger 42])+    )++  roundTripPrim+    "FixedSizeBinary"+    ( fixture $+        fromMaybeFixedSizeBinary+          4+          ( V.fromList+              [Just (BS.pack [0, 1, 2, 3]), Just (BS.pack [0xFF, 0xFE, 0xFD, 0xFC])]+          )+    )++  roundTripPrim "Utf8" (fromTexts (V.fromList ["alpha", "beta", "", "\xe2\x9a\xa1"]))+  roundTripPrim "Binary" (fromByteStrings (V.fromList [BS.pack [1, 2, 3], BS.empty, BS.pack [0xFF]]))+  roundTripPrim "LargeUtf8" (fromMaybeLargeTexts (V.fromList [Just "alpha", Just "beta"]))+  roundTripPrim+    "LargeBinary"+    (fromMaybeLargeByteStrings (V.fromList [Just (BS.pack [0, 1, 2]), Just (BS.pack [0xFF])]))++  -- Nullable variants (every fixture holds at least one null row).+  roundTripPrim+    "Int8Maybe"+    (fromMaybes PInt8 (V.fromList [Just 1, Nothing, Just (-1), Just 42]))+  roundTripPrim+    "Int16Maybe"+    (fromMaybes PInt16 (V.fromList [Just 100, Nothing, Just (-200)]))+  roundTripPrim+    "Int32Maybe"+    (fromMaybes PInt32 (V.fromList [Nothing, Just 0, Just maxBound]))+  roundTripPrim+    "Int64Maybe"+    (fromMaybes PInt64 (V.fromList [Just 0, Nothing, Just (-1)]))+  roundTripPrim+    "UInt8Maybe"+    (fromMaybes PUInt8 (V.fromList [Just 0, Just 255, Nothing]))+  roundTripPrim+    "UInt16Maybe"+    (fromMaybes PUInt16 (V.fromList [Just 0, Nothing]))+  roundTripPrim+    "UInt32Maybe"+    (fromMaybes PUInt32 (V.fromList [Just 0, Nothing, Just maxBound]))+  roundTripPrim+    "UInt64Maybe"+    (fromMaybes PUInt64 (V.fromList [Just 0, Nothing]))+  roundTripPrim+    "Float16Maybe"+    (fromMaybes PFloat16 (V.fromList [Just (Float16 0), Just (Float16 0x3C00), Nothing]))+  roundTripPrim+    "FloatMaybe"+    (fromMaybes PFloat (V.fromList [Just 1.5, Nothing, Just (-2.25)]))+  roundTripPrim+    "DoubleMaybe"+    (fromMaybes PDouble (V.fromList [Just 1.5, Nothing]))+  roundTripPrim+    "BoolMaybe"+    (fromMaybeBools (V.fromList [Just True, Nothing, Just False, Just True]))++  roundTripPrim+    "Utf8Maybe"+    (fromMaybeTexts (V.fromList [Just "alpha", Nothing, Just ""]))+  roundTripPrim+    "BinaryMaybe"+    (fromMaybeByteStrings (V.fromList [Just (BS.pack [1, 2]), Nothing]))+  roundTripPrim+    "LargeUtf8Maybe"+    (fromMaybeLargeTexts (V.fromList [Nothing, Just "beta"]))+  roundTripPrim+    "LargeBinaryMaybe"+    (fromMaybeLargeByteStrings (V.fromList [Just (BS.pack [0xFF]), Nothing]))+  roundTripPrim+    "FixedSizeBinaryMaybe"+    ( fixture $+        fromMaybeFixedSizeBinary+          3+          (V.fromList [Just (BS.pack [1, 2, 3]), Nothing, Just (BS.pack [4, 5, 6])])+    )++  roundTripPrim "Date32Maybe" (fromMaybes PDate32 (V.fromList [Just 0, Nothing]))+  roundTripPrim "Date64Maybe" (fromMaybes PDate64 (V.fromList [Just 0, Nothing]))+  roundTripPrim "Time32Maybe" (fromMaybes PTime32 (V.fromList [Just 0, Nothing]))+  roundTripPrim "Time64Maybe" (fromMaybes PTime64 (V.fromList [Just 0, Nothing]))+  roundTripPrim "TimestampMaybe" (fromMaybes PTimestamp (V.fromList [Just 0, Nothing]))+  roundTripPrim "DurationMaybe" (fromMaybes PDuration (V.fromList [Just 0, Nothing]))++  -- ============================================================+  -- Nested columns+  -- ============================================================++  -- Struct with two primitive children.+  roundTripNested+    "Struct"+    ( nestedField "s" False AStruct $+        V.fromList+          [ plainField "id" False (AInt 64 True)+          , plainField "name" False AUtf8+          ]+    )+    ( fixture . mkStruct 3 Nothing $+        V.fromList+          [ ("id", primColumn PInt64 (VS.fromList [1, 2, 3]))+          , ("name", fromTexts (V.fromList ["a", "b", "c"]))+          ]+    )++  roundTripNested+    "StructMaybe"+    ( nestedField "s" True AStruct $+        V.fromList+          [ plainField "id" False (AInt 32 True)+          , plainField "flag" False ABool+          ]+    )+    ( fixture $+        mkStruct+          3+          (validityFromBools (V.fromList [True, False, True]))+          ( V.fromList+              [ ("id", primColumn PInt32 (VS.fromList [1, 2, 3]))+              , ("flag", fromBools (V.fromList [True, False, True]))+              ]+          )+    )++  roundTripNested+    "List<int32>"+    ( nestedField "l" False AList $+        V.fromList+          [plainField "item" False (AInt 32 True)]+    )+    ( fixture $+        mkList+          Nothing+          (VS.fromList [0, 2, 2, 5])+          (primColumn PInt32 (VS.fromList [10, 20, 30, 40, 50]))+    )++  roundTripNested+    "ListMaybe<int32>"+    ( nestedField "l" True AList $+        V.fromList+          [plainField "item" False (AInt 32 True)]+    )+    ( fixture $+        mkList+          (validityFromBools (V.fromList [True, False, True]))+          (VS.fromList [0, 2, 2, 5])+          (primColumn PInt32 (VS.fromList [10, 20, 30, 40, 50]))+    )++  roundTripNested+    "LargeList<int32>"+    ( nestedField "l" False ALargeList $+        V.fromList+          [plainField "item" False (AInt 32 True)]+    )+    ( fixture $+        mkLargeList+          Nothing+          (VS.fromList [0, 2, 2, 5])+          (primColumn PInt32 (VS.fromList [1, 2, 3, 4, 5]))+    )++  roundTripNested+    "LargeListMaybe<int32>"+    ( nestedField "l" True ALargeList $+        V.fromList+          [plainField "item" False (AInt 32 True)]+    )+    ( fixture $+        mkLargeList+          (validityFromBools (V.fromList [True, False, True]))+          (VS.fromList [0, 2, 2, 5])+          (primColumn PInt32 (VS.fromList [1, 2, 3, 4, 5]))+    )++  roundTripNested+    "FixedSizeList<3 of int32>"+    ( nestedField "l" False (AFixedSizeList 3) $+        V.fromList+          [plainField "item" False (AInt 32 True)]+    )+    ( fixture $+        mkFixedSizeList+          3+          2+          Nothing+          (primColumn PInt32 (VS.fromList [1, 2, 3, 4, 5, 6]))+    )++  roundTripNested+    "FixedSizeListMaybe<2 of int32>"+    ( nestedField "l" True (AFixedSizeList 2) $+        V.fromList+          [plainField "item" False (AInt 32 True)]+    )+    ( fixture $+        mkFixedSizeList+          2+          3+          (validityFromBools (V.fromList [True, False, True]))+          (primColumn PInt32 (VS.fromList [1, 2, 3, 4, 5, 6]))+    )++  -- Map<string, int32>. Arrow encodes maps as a list of struct+  -- <key, value> pairs; the map field has one child (the struct).+  roundTripNested+    "Map<string, int32>"+    ( nestedField "m" False (AMap False) $+        V.fromList+          [ nestedField "entries" False AStruct $+              V.fromList+                [ plainField "key" False AUtf8+                , plainField "value" False (AInt 32 True)+                ]+          ]+    )+    ( fixture $+        mkMap+          Nothing+          (VS.fromList [0, 2, 2, 3])+          (fromTexts (V.fromList ["a", "b", "c"]))+          (primColumn PInt32 (VS.fromList [1, 2, 3]))+    )++  -- Dense union over (int32, utf8).+  roundTripNested+    "DenseUnion<int32, utf8>"+    ( nestedField "u" False (AUnion Dense (V.fromList [0, 1])) $+        V.fromList+          [ plainField "v_int" False (AInt 32 True)+          , plainField "v_text" False AUtf8+          ]+    )+    ( fixture $+        mkDenseUnion+          (VS.fromList [0, 1, 0])+          (VS.fromList [0, 0, 1])+          ( V.fromList+              [ primColumn PInt32 (VS.fromList [100, 200])+              , fromTexts (V.fromList ["hello"])+              ]+          )+    )++  -- Sparse union over (bool, int32).+  roundTripNested+    "SparseUnion<bool, int32>"+    ( nestedField "u" False (AUnion Sparse (V.fromList [0, 1])) $+        V.fromList+          [ plainField "flag" False ABool+          , plainField "value" False (AInt 32 True)+          ]+    )+    ( fixture $+        mkSparseUnion+          (VS.fromList [0, 1, 0])+          ( V.fromList+              [ fromBools (V.fromList [True, False, False])+              , primColumn PInt32 (VS.fromList [0, 42, 0])+              ]+          )+    )++  -- FlatBuffers reader / writer round-trip: build a typical+  -- multi-column batch with the FB writer, parse back with the FB+  -- reader, assert schema equality.+  flatBufRoundTrip+  flatBufSchemaSelfCheck++  -- Golden pyarrow interop: reference .arrows files produced+  -- by pyarrow's ipc.new_stream, checked into test/golden.+  -- Decoding them proves the FlatBuffers reader handles the+  -- shapes arrow-cpp emits (vs only what wireform's own+  -- writer produces).+  pyarrowGoldenRoundTrip++  -- Hedgehog property suites: column core, generated round-trips,+  -- malformed-input robustness and the Arrow-layout vectors. Each+  -- returns False on failure.+  coreOk <- Core.tests+  propsOk <- Props.tests+  malformedOk <- Malformed.tests+  sectionsOk <- Sections.tests+  vectorOk <- Vector.tests+  unless (coreOk && propsOk && malformedOk && sectionsOk && vectorOk) $+    failTest "FAIL: wireform-arrow property suites"++  putStrLn "All wireform-arrow round-trip tests passed."+++flatBufSchemaSelfCheck :: IO ()+flatBufSchemaSelfCheck = do+  let cases =+        [ Schema (V.fromList [plainField "a" False (AInt 32 True)]) Little V.empty V.empty+        , Schema+            ( V.fromList+                [ plainField "id" False (AInt 64 True)+                , plainField "name" True AUtf8+                , plainField "amount" False (ADecimal 12 4)+                , plainField "ts" True (ATimestamp Nanosecond (Just "UTC"))+                , plainField "blob" True ABinary+                , plainField "tag" False (AFixedSizeBinary 16)+                ]+            )+            Little+            V.empty+            V.empty+        , -- Post-V5 type tags (Utf8View / BinaryView / RunEndEncoded /+          -- ListView / LargeListView). Arrow.Column doesn't materialise+          -- their data buffers, but the schema flatbuffer round-trips+          -- so wireform can interoperate with newer Arrow producers.+          Schema+            ( V.fromList+                [ plainField "v" True AUtf8View+                , plainField "b" True ABinaryView+                , plainField "ree" True ARunEndEncoded+                , plainField "lv" True AListView+                , plainField "llv" True ALargeListView+                ]+            )+            Little+            V.empty+            V.empty+        ]+  mapM_+    ( \sch -> do+        let bs = buildSchemaMessage sch+        case decodeSchemaMessage bs of+          Right got+            | got == sch ->+                putStrLn $ "OK: FlatBuffers schema self-roundtrip: " ++ describe sch+          Right got ->+            failTest $+              "FB schema roundtrip mismatch:\n got: "+                ++ show got+                ++ "\n exp: "+                ++ show sch+          Left e ->+            failTest $ "FB schema roundtrip decode failed: " ++ e+    )+    cases+  where+    describe sch =+      "("+        ++ show (V.length (arrowFields sch))+        ++ " field"+        ++ (if V.length (arrowFields sch) == 1 then "" else "s")+        ++ ")"+++flatBufRoundTrip :: IO ()+flatBufRoundTrip = do+  -- Multi-column round-trip exercising the high-level API in+  -- "Arrow.Stream": Schema + batches go in, bytes come out, and+  -- the inverse recovers the same shape.+  highLevelRoundTrip+    "Multi-column"+    ( Schema+        { arrowFields =+            V.fromList+              [ plainField "i" False (AInt 32 True)+              , plainField "s" True AUtf8+              ]+        , arrowEndianness = Little+        , arrowMetadata = V.empty+        , arrowFeatures = V.empty+        }+    )+    ( V.fromList+        [ primColumn PInt32 (VS.fromList ([1, 2, 3] :: [Int32]))+        , fromMaybeTexts (V.fromList [Just "x", Nothing, Just "z"])+        ]+    )++  -- Post-V5 columns: writer + reader byte-compatible end to end.+  highLevelRoundTrip+    "Utf8View"+    (Schema (V.singleton (plainField "v" True AUtf8View)) Little V.empty V.empty)+    ( V.singleton+        ( fromMaybeUtf8View+            ( V.fromList+                [ Just "short"+                , Nothing+                , Just "this string is definitely longer than twelve bytes"+                ]+            )+        )+    )++  highLevelRoundTrip+    "ListView<int32>"+    ( Schema+        ( V.singleton+            ( nestedField+                "lv"+                False+                AListView+                ( V.singleton+                    (plainField "item" False (AInt 32 True))+                )+            )+        )+        Little+        V.empty+        V.empty+    )+    ( V.singleton+        ( fixture $+            mkListView+              Nothing+              (VS.fromList ([0, 2, 5] :: [Int32]))+              (VS.fromList ([2, 3, 1] :: [Int32]))+              (primColumn PInt32 (VS.fromList ([10, 20, 30, 40, 50, 60] :: [Int32])))+        )+    )++  highLevelRoundTrip+    "RunEndEncoded(int32, int64?)"+    ( Schema+        ( V.singleton+            ( nestedField "ree" True ARunEndEncoded $+                V.fromList+                  [ plainField "run_ends" False (AInt 32 True)+                  , plainField "values" True (AInt 64 True)+                  ]+            )+        )+        Little+        V.empty+        V.empty+    )+    ( V.singleton+        ( fixture $+            mkRunEndEncoded+              (primColumn PInt32 (VS.fromList ([3, 5, 8] :: [Int32])))+              (fromMaybes PInt64 (V.fromList [Just 100, Nothing, Just 300]))+        )+    )++  -- Dictionary-encoded utf8: the high-level API auto-extracts+  -- the dictionary batch and auto-resolves on read.+  let dictField =+        Field+          "d"+          False+          AUtf8+          V.empty+          (Just (DictionaryEncoding 0 (AInt 32 True) False))+          V.empty+  highLevelRoundTrip+    "Dictionary<utf8>"+    (Schema (V.singleton dictField) Little V.empty V.empty)+    ( V.singleton+        ( fixture $+            mkDictionary+              0+              (primColumn PInt32 (VS.fromList ([0, 1, 0, 2, 1] :: [Int32])))+              (fromTexts (V.fromList ["a", "b", "c"]))+        )+    )++  -- ANull column: schema metadata round-trip + ColNull row count+  highLevelRoundTrip+    "Null"+    (Schema (V.singleton (plainField "n" False ANull)) Little V.empty V.empty)+    (V.singleton (ColNull 5))++  -- Custom metadata round-trip on schema + field+  customMetadataRoundTrip++  -- Nested struct via Arrow.Record.structE / structD+  nestedStructRoundTrip++  -- Nullable nested struct via encoderFromRowEncoder /+  -- decoderFromRowDecoder + Arrow.Record.nullable / nullableD.+  nullableNestedStructRoundTrip++  -- Schema fingerprint: determinism + structural equivalence.+  schemaFingerprintTests++  -- Record helpers: subsetTable / projectTable /+  -- columnDWithDefault / NameStrategy / validateMapKeysSorted.+  recordHelperTests++  -- Streaming reader: pull batches one at a time, then drain.+  streamingRoundTrip+    (Schema (V.fromList [plainField "n" False (AInt 32 True)]) Little V.empty V.empty)+    [ V.singleton (primColumn PInt32 (VS.fromList ([1, 2] :: [Int32])))+    , V.singleton (primColumn PInt32 (VS.fromList ([3] :: [Int32])))+    , V.singleton (primColumn PInt32 (VS.fromList ([4, 5, 6, 7] :: [Int32])))+    ]++  -- Column projection on a multi-column stream: a 3-column+  -- batch should narrow to exactly the requested columns in+  -- the requested order.+  projectionRoundTrip++  -- ZSTD body compression (writer + reader): exercises+  -- BodyCompression on a multi-column batch, asserting the+  -- decoded values match the source.+  bodyCompressionRoundTrip+    BodyZstd+    ( Schema+        ( V.fromList+            [ plainField "n" False (AInt 64 True)+            , plainField "s" False AUtf8+            ]+        )+        Little+        V.empty+        V.empty+    )+    ( V.fromList+        [ primColumn+            PInt64+            ( VS.fromList+                ([1 .. 1000] :: [Int64]) -- enough bytes that ZSTD shrinks+            )+        , fromTexts (V.replicate 1000 "highly-compressible-payload")+        ]+    )++  -- LZ4_FRAME body compression: same shape / sizing as the ZSTD+  -- case. Verifies the lz4-hs Codec.Lz4 frame compressor ++  -- decompressor round-trip through the full BodyCompression+  -- pipeline (per-buffer envelope, length prefix, offsets+  -- rewritten on decode, etc.).+  bodyCompressionRoundTrip+    LZ4Frame+    ( Schema+        ( V.fromList+            [ plainField "n" False (AInt 64 True)+            , plainField "s" False AUtf8+            ]+        )+        Little+        V.empty+        V.empty+    )+    ( V.fromList+        [ primColumn PInt64 (VS.fromList ([1 .. 1000] :: [Int64]))+        , fromTexts (V.replicate 1000 "highly-compressible-payload")+        ]+    )++  -- DictReplaceOnChange: two batches with the SAME dict id but+  -- different values. The writer should emit two dict batches;+  -- the reader should resolve each record batch against the+  -- most-recently-emitted dict for that id.+  dictReplacementRoundTrip++  -- Tensor message round-trip.+  tensorRoundTrip++  -- SparseTensor (COO) round-trip.+  sparseTensorRoundTrip+++{- | Consume the golden .arrows files in @test/golden/@ and+assert 'decodeArrowStream' returns the ColumnArray values we+expect from the pyarrow-side generator.++The fixtures:+  pa_int32.arrows   : int32 column [1,2,3,4,5]+  pa_mixed.arrows   : (int64, nullable utf8, nullable bool), 3 rows+  pa_dict.arrows    : dictionary<utf8, int32> with values ["a","b","c"]+                       and indices [0,1,0,2,1]+-}+pyarrowGoldenRoundTrip :: IO ()+pyarrowGoldenRoundTrip = do+  goldenCheck+    "pa_int32.arrows"+    (V.singleton (primColumn PInt32 (VS.fromList [1, 2, 3, 4, 5 :: Int32])))++  goldenCheck+    "pa_mixed.arrows"+    ( V.fromList+        [ primColumn PInt64 (VS.fromList [10, 20, 30 :: Int64])+        , fromMaybeTexts (V.fromList [Just "alpha", Nothing, Just "gamma"])+        , fromMaybeBools (V.fromList [Just True, Just False, Nothing])+        ]+    )++  -- Dictionary-encoded batch: the decoder resolves+  -- ColDictionary's values against the dict batches pyarrow+  -- emitted ahead of the record batch.+  goldenDictCheck+    "pa_dict.arrows"+    [0, 1, 0, 2, 1]+    (fromTexts (V.fromList ["a", "b", "c"]))+++goldenCheck :: FilePath -> V.Vector ColumnArray -> IO ()+goldenCheck name expected = do+  bs <- BS.readFile ("test/golden/" <> name)+  case decodeArrowStream defaultReadOptions bs of+    Left e -> failTest $ "golden " <> name <> ": decode: " <> e+    Right (_sch, batches)+      | [b] <- batches+      , b == expected ->+          putStrLn $ "OK: pyarrow golden " <> name+      | otherwise ->+          failTest $+            "golden "+              <> name+              <> " mismatch:\n got: "+              <> show batches+              <> "\n exp: "+              <> show [expected]+++-- Dictionary-encoded batches need a bespoke matcher: the keys and+-- the resolved values are compared piecewise. pyarrow fields are+-- nullable, so the column decodes as a 'ColDictionary' whose keys+-- carry a validity slot (here without null rows).+goldenDictCheck+  :: FilePath -> [Int] -> ColumnArray -> IO ()+goldenDictCheck name expectedIndices expectedValues = do+  bs <- BS.readFile ("test/golden/" <> name)+  case decodeArrowStream defaultReadOptions bs of+    Left e -> failTest $ "golden " <> name <> ": decode: " <> e+    Right (_sch, batches) -> case batches of+      [b] | V.length b == 1 -> case V.head b of+        col@(ColDictionary _ _ vals)+          | hasValiditySlot col+          , nullCount col == 0+          , keysOf col == map Just expectedIndices+          , vals == expectedValues ->+              putStrLn $ "OK: pyarrow golden " <> name+          | otherwise ->+              failTest $+                "golden "+                  <> name+                  <> " dict mismatch:\n idx="+                  <> show (keysOf col)+                  <> " vals="+                  <> show vals+        other ->+          failTest $ "golden " <> name <> " expected ColDictionary, got " <> columnTag other+      _ -> failTest $ "golden " <> name <> " expected 1 batch with 1 column"+  where+    keysOf col = map (dictKeyAt col) [0 .. columnLength col - 1]+++sparseTensorRoundTrip :: IO ()+sparseTensorRoundTrip = do+  -- Tiny 3x3 sparse int32 tensor with 2 non-zeros at (0,1) and+  -- (2,0). COO indices are Int64 pairs.+  let !idx =+        BS.pack+          [ 0+          , 0+          , 0+          , 0+          , 0+          , 0+          , 0+          , 0 -- row 0+          , 1+          , 0+          , 0+          , 0+          , 0+          , 0+          , 0+          , 0 -- col 1+          , 2+          , 0+          , 0+          , 0+          , 0+          , 0+          , 0+          , 0 -- row 2+          , 0+          , 0+          , 0+          , 0+          , 0+          , 0+          , 0+          , 0 -- col 0+          ]+      !vals =+        BS.pack+          [ 7+          , 0+          , 0+          , 0 -- value 7+          , 9+          , 0+          , 0+          , 0 -- value 9+          ]+      !st =+        SparseTensor+          { sparseTensorType = AInt 32 True+          , sparseTensorShape =+              V.fromList+                [TensorDim 3 "rows", TensorDim 3 "cols"]+          , sparseNonZeroLength = 2+          , sparseIndicesType = AInt 64 True+          , sparseIndicesBody = idx+          , sparseIndicesCanonical = True+          , sparseTensorBody = vals+          }+      !frame = encodeSparseTensorFrame st+  case decodeSparseTensorFrame frame of+    Left e -> failTest $ "decodeSparseTensorFrame: " ++ e+    Right (sout, _)+      | sparseTensorType sout == AInt 32 True+      , sparseNonZeroLength sout == 2+      , sparseIndicesBody sout == idx+      , sparseTensorBody sout == vals ->+          putStrLn "OK: SparseTensor message round-trip (COO, 3x3 int32, nnz=2)"+      | otherwise ->+          failTest $ "sparse tensor mismatch:\n got " ++ show sout+++tensorRoundTrip :: IO ()+tensorRoundTrip = do+  -- 2x3 tensor of Int32: raw little-endian body, row-major.+  let !body =+        BS.pack+          [ 0x01+          , 0+          , 0+          , 0+          , 0x02+          , 0+          , 0+          , 0+          , 0x03+          , 0+          , 0+          , 0+          , 0x04+          , 0+          , 0+          , 0+          , 0x05+          , 0+          , 0+          , 0+          , 0x06+          , 0+          , 0+          , 0+          ]+      !tin =+        Tensor+          { tensorType = AInt 32 True+          , tensorShape =+              V.fromList+                [TensorDim 2 "rows", TensorDim 3 "cols"]+          , tensorStrides = V.empty+          , tensorBody = body+          }+      !frame = encodeTensorFrame tin+  case decodeTensorFrame frame of+    Left e -> failTest $ "decodeTensorFrame: " ++ e+    Right (tout, rest)+      | BS.null rest+      , tensorType tout == AInt 32 True+      , V.toList (tensorShape tout)+          == [TensorDim 2 "rows", TensorDim 3 "cols"]+      , tensorBody tout == body ->+          putStrLn "OK: Tensor message round-trip (2x3 int32)"+      | otherwise ->+          failTest $+            "tensor round-trip mismatch:\n got "+              ++ show tout+              ++ " rest="+              ++ show (BS.length rest)+              ++ "B"+++dictReplacementRoundTrip :: IO ()+dictReplacementRoundTrip = do+  let !sch =+        Schema+          ( V.singleton+              ( Field+                  "d"+                  False+                  AUtf8+                  V.empty+                  (Just (DictionaryEncoding 0 (AInt 32 True) False))+                  V.empty+              )+          )+          Little+          V.empty+          V.empty+      !batch1 =+        V.singleton . fixture $+          mkDictionary+            0+            (primColumn PInt32 (VS.fromList [0, 1, 0]))+            (fromTexts (V.fromList ["a", "b"]))+      !batch2 =+        V.singleton . fixture $+          mkDictionary+            0+            (primColumn PInt32 (VS.fromList [0, 1, 0]))+            (fromTexts (V.fromList ["x", "y"]))+      !opts = defaultWriteOptions {writeDictHandling = DictReplaceOnChange}+      !bytes = encodeArrowStream opts sch [batch1, batch2]+  case bytes >>= decodeArrowStream defaultReadOptions of+    Left e -> failTest $ "dict-replace round-trip: " ++ e+    Right (_, batches)+      | [b1, b2] <- batches+      , V.length b1 == 1+      , V.length b2 == 1 ->+          let !c1 = V.unsafeIndex b1 0+              !c2 = V.unsafeIndex b2 0+          in case (c1, c2) of+               (ColDictionary _ ix1 v1, ColDictionary _ ix2 v2)+                 | valuesToList v1 == ["a", "b"]+                 , valuesToList v2 == ["x", "y"]+                 , keysToList ix1 == [Just 0, Just 1, Just 0]+                 , keysToList ix2 == [Just 0, Just 1, Just 0] ->+                     putStrLn "OK: dictionary replacement across batches"+               _ ->+                 failTest $+                   "dict-replace mismatch:\n got "+                     ++ show batches+      | otherwise ->+          failTest $+            "dict-replace expected 2 batches, got "+              ++ show (length batches)+  where+    valuesToList :: ColumnArray -> [Text]+    valuesToList c = either (const []) (VG.toList . VG.mapMaybe id) (toTextVector c)+    keysToList :: ColumnArray -> [Maybe Int32]+    keysToList c = maybe [] (VG.toList . toMaybeVector) (asPrim PInt32 c)+++-- | Without the codec's cabal flag the writer silently falls back to an+-- uncompressed body and the round trip would pass without exercising+-- compression, so an unavailable codec is skipped outright.+bodyCompressionRoundTrip :: BodyCompressionCodec -> Schema -> V.Vector ColumnArray -> IO ()+bodyCompressionRoundTrip codec sch cols+  | not (bodyCompressionAvailable codec) =+      putStrLn $ "SKIP: body compression " ++ show codec ++ " (not built in; enable the zstd / lz4 cabal flags)"+  | otherwise = do+    let !opts = defaultWriteOptions {writeBodyCompression = Just codec}+        !encoded = encodeArrowStream opts sch [cols]+    case encoded >>= \bytes -> (,) bytes <$> decodeArrowStream defaultReadOptions bytes of+      Left e -> failTest $ "body-compression round-trip: " ++ e+      Right (bytes, (_sch', batches))+        | [got] <- batches+        , got == cols ->+            putStrLn $+              "OK: body compression "+                ++ show codec+                ++ " ("+                ++ show (BS.length bytes)+                ++ " bytes)"+        | otherwise ->+            failTest $+              "body-compression mismatch: got "+                ++ show batches+++streamingRoundTrip :: Schema -> [V.Vector ColumnArray] -> IO ()+streamingRoundTrip sch batches = do+  let bytes = encodeArrowStream defaultWriteOptions sch batches+  case bytes >>= openStreamReader defaultReadOptions of+    Left e -> failTest $ "openStreamReader: " ++ e+    Right rd0 -> do+      when (streamReaderSchema rd0 /= sch) $+        failTest "streamReaderSchema mismatch"+      -- Pull first batch via streamReaderNext; drain rest via toList.+      case streamReaderNext rd0 of+        Left e -> failTest $ "streamReaderNext (first): " ++ e+        Right Nothing -> failTest "streamReaderNext: stream empty"+        Right (Just (cols0, rd1)) -> do+          when (Just cols0 /= listToMaybe batches) $+            failTest $+              "streamReaderNext (first) mismatch:\n got "+                ++ show cols0+                ++ "\n exp "+                ++ show (take 1 batches)+          case streamReaderToList rd1 of+            Left e -> failTest $ "streamReaderToList: " ++ e+            Right rest+              | rest == drop 1 batches ->+                  putStrLn "OK: streaming reader iterates all batches"+              | otherwise ->+                  failTest $+                    "streamReaderToList: tail mismatch\n got "+                      ++ show rest+                      ++ "\n exp "+                      ++ show (drop 1 batches)+  -- Iter-shaped variant: the same drain via Columnar.Stream.+  case bytes >>= openStreamReader defaultReadOptions of+    Left e -> failTest $ "openStreamReader defaultReadOptions (iter): " ++ e+    Right rd0 ->+      case IS.iterToList (streamReaderIter rd0) of+        Left e -> failTest $ "streamReaderIter drain: " ++ e+        Right got+          | got == batches ->+              putStrLn "OK: streamReaderIter drains all batches"+          | otherwise ->+              failTest $+                "streamReaderIter mismatch:\n got "+                  ++ show got+                  ++ "\n exp "+                  ++ show batches+++{- | Schema-level + field-level @custom_metadata@ pairs survive+a full encode to decode round-trip via the FlatBuffers schema+writer + reader.+-}+customMetadataRoundTrip :: IO ()+customMetadataRoundTrip = do+  let !field =+        (plainField "n" False (AInt 32 True))+          { fieldMetadata = V.fromList [("description", "row id"), ("unit", "count")]+          }+      !sch =+        Schema+          { arrowFields = V.singleton field+          , arrowEndianness = Little+          , arrowMetadata =+              V.fromList+                [ ("pandas", "{}")+                , ("creator", "wireform-test")+                ]+          , arrowFeatures = V.empty+          }+      !batch = V.singleton (primColumn PInt32 (VS.fromList ([1, 2, 3] :: [Int32])))+      !bytes = encodeArrowStream defaultWriteOptions sch [batch]+  case bytes >>= decodeArrowStream defaultReadOptions of+    Left e -> failTest $ "customMetadata roundtrip: " ++ e+    Right (sch', _batches) -> do+      expect+        "schema custom_metadata roundtrips"+        (arrowMetadata sch' == arrowMetadata sch)+      let recoveredField = V.unsafeIndex (arrowFields sch') 0+      expect+        "field custom_metadata roundtrips"+        (fieldMetadata recoveredField == fieldMetadata field)+++{- | Nested record via 'structE' + 'structD'. The inner record+(Address) becomes a 'ColStruct' column inside the outer+record (Customer); a round-trip through 'encodeTable' /+'decodeTable' must recover the exact value.+-}+nestedStructRoundTrip :: IO ()+nestedStructRoundTrip = do+  let !addrEnc =+        AR.fieldE "city" cityF AR.utf8E+          <> AR.fieldE "zip" zipF AR.utf8E+      !addrDec =+        Address+          <$> AR.columnD "city" AR.utf8D+          <*> AR.columnD "zip" AR.utf8D+      !custEnc =+        AR.fieldE "name" nameF AR.utf8E+          <> AR.structE "addr" addrF addrEnc+          <> AR.fieldE "age" ageF AR.int32E+      !custDec =+        Customer+          <$> AR.columnD "name" AR.utf8D+          <*> AR.structD "addr" addrDec+          <*> AR.columnD "age" AR.int32D+      !tbl = AR.table custEnc custDec+      !rows =+        V.fromList+          [ Customer "Alice" (Address "Atlantis" "00001") 30+          , Customer "Bob" (Address "Brisbane" "4000") 45+          , Customer "Carol" (Address "Calcutta" "700001") 28+          ]+      (!sch, !cols) = AR.encodeTable tbl rows+      !bytes = encodeArrowStream defaultWriteOptions sch [cols]+  case bytes >>= decodeArrowStream defaultReadOptions of+    Left e -> failTest $ "nested struct decode: " ++ e+    Right (sch', batches) -> case batches of+      [batch] -> case AR.decodeTable tbl sch' batch of+        Left e -> failTest $ "nested struct decodeTable: " ++ e+        Right got+          | got == rows ->+              putStrLn "OK: nested struct via structE / structD"+          | otherwise ->+              failTest $+                "nested struct mismatch:\n got "+                  ++ show (V.toList got)+                  ++ "\n exp "+                  ++ show (V.toList rows)+      _ -> failTest "nested struct: expected 1 batch"+++{- | Nullable nested record. Same shape as 'nestedStructRoundTrip'+but the @addr@ column is @Maybe Address@; the encoder builds+a nullable 'ColStruct' with a top-level validity mask, and the+decoder reconstructs the @Just@/@Nothing@ pattern.+-}+nullableNestedStructRoundTrip :: IO ()+nullableNestedStructRoundTrip = do+  let !addrEnc =+        AR.fieldE "city" cityF AR.utf8E+          <> AR.fieldE "zip" zipF AR.utf8E+      !addrDec =+        Address+          <$> AR.columnD "city" AR.utf8D+          <*> AR.columnD "zip" AR.utf8D+      !custEnc =+        AR.fieldE "name" coNameF AR.utf8E+          <> AR.structEMaybe "addr_opt" maybeAddrF addrEnc+          <> AR.fieldE "age" coAgeF AR.int32E+      !custDec =+        CustomerOpt+          <$> AR.columnD "name" AR.utf8D+          <*> AR.structDMaybe "addr_opt" addrDec+          <*> AR.columnD "age" AR.int32D+      !tbl = AR.table custEnc custDec :: AR.Table CustomerOpt+      !rows =+        V.fromList+          [ CustomerOpt "Alice" (Just (Address "Atlantis" "00001")) 30+          , CustomerOpt "Bob" Nothing 45+          , CustomerOpt "Carol" (Just (Address "Calcutta" "700001")) 28+          , CustomerOpt "Dave" Nothing 50+          ]+      (!sch, !cols) = AR.encodeTable tbl rows+      !bytes = encodeArrowStream defaultWriteOptions sch [cols]+  case bytes >>= decodeArrowStream defaultReadOptions of+    Left e -> failTest $ "nullable nested struct decode: " ++ e+    Right (sch', batches) -> case batches of+      [batch] -> case AR.decodeTable tbl sch' batch of+        Left e -> failTest $ "nullable nested decodeTable: " ++ e+        Right got+          | got == rows ->+              putStrLn "OK: nullable nested struct via structEMaybe / structDMaybe"+          | otherwise ->+              failTest $+                "nullable nested mismatch:\n got "+                  ++ show (V.toList got)+                  ++ "\n exp "+                  ++ show (V.toList rows)+      _ -> failTest "nullable nested struct: expected 1 batch"+++{- | 'schemaFingerprint' tests: determinism, equivalence-class+equality, and difference detection.+-}+schemaFingerprintTests :: IO ()+schemaFingerprintTests = do+  let !sch1 =+        Schema+          ( V.fromList+              [ plainField "id" False (AInt 64 True)+              , plainField "name" True AUtf8+              ]+          )+          Little+          V.empty+          V.empty+      !sch2 =+        Schema+          ( V.fromList+              [ plainField "id" False (AInt 64 True)+              , plainField "name" True AUtf8+              ]+          )+          Little+          (V.fromList [("creator", "wireform")]) -- different annotation+          (V.fromList [FeatureDictionaryReplacement]) -- different feature flag+      !sch3 =+        Schema+          ( V.fromList+              [ plainField "id" False (AInt 64 True)+              , plainField "name2" True AUtf8 -- different field name+              ]+          )+          Little+          V.empty+          V.empty+      !fp1 = schemaFingerprint sch1+      !fp2 = schemaFingerprint sch2+      !fp3 = schemaFingerprint sch3+  expect+    "fingerprint is deterministic across calls"+    (fp1 == schemaFingerprint sch1)+  expect+    "fingerprint ignores annotation fields"+    (fp1 == fp2)+  expect+    "fingerprint distinguishes different field names"+    (fp1 /= fp3)+  expect+    "schemaEquivalent matches fingerprint equality (1==2)"+    (schemaEquivalent sch1 sch2 == (fp1 == fp2))+  expect+    "schemaEquivalent matches fingerprint equality (1==3)"+    (schemaEquivalent sch1 sch3 == (fp1 == fp3))+++{- | 'NameStrategy', 'columnDWithDefault', 'projectTable',+'subsetTable', and 'validateMapKeysSorted' tests.+-}+fst3 :: (a, b, c) -> a+fst3 (a, _, _) = a+++snd3 :: (a, b, c) -> b+snd3 (_, b, _) = b+++recordHelperTests :: IO ()+recordHelperTests = do+  -- NameStrategy+  expect+    "NameAsIs is identity"+    (AR.applyNameStrategy AR.NameAsIs "userId" == "userId")+  expect+    "NameSnakeCase userId -> user_id"+    (AR.applyNameStrategy AR.NameSnakeCase "userId" == "user_id")+  expect+    "NameSnakeCase userIDValue -> user_id_value (acronym boundary)"+    (AR.applyNameStrategy AR.NameSnakeCase "userIDValue" == "user_id_value")+  expect+    "NameSnakeCase XMLHttpRequest -> xml_http_request"+    (AR.applyNameStrategy AR.NameSnakeCase "XMLHttpRequest" == "xml_http_request")+  expect+    "NameCamelCase user_id -> userId"+    (AR.applyNameStrategy AR.NameCamelCase "user_id" == "userId")+  expect+    "NameUpperSnakeCase userId -> USER_ID"+    (AR.applyNameStrategy AR.NameUpperSnakeCase "userId" == "USER_ID")++  -- validateMapKeysSorted+  -- Build a ColMap with sorted keys vs unsorted keys.+  let !sortedKeys = fromTexts (V.fromList ["a", "b", "c"])+      !unsortedKeys = fromTexts (V.fromList ["b", "a", "c"])+      !vals = primColumn PInt32 (VS.fromList [1, 2, 3 :: Int32])+      !offsets = VS.fromList [0, 3 :: Int32]+      !sortedMap = fixture (mkMap Nothing offsets sortedKeys vals)+      !unsortedMap = fixture (mkMap Nothing offsets unsortedKeys vals)+  case validateMapKeysSorted sortedMap of+    Right () -> putStrLn "OK: validateMapKeysSorted accepts sorted keys"+    Left e -> failTest $ "expected sorted accept, got " ++ e+  case validateMapKeysSorted unsortedMap of+    Left _ -> putStrLn "OK: validateMapKeysSorted rejects unsorted keys"+    Right () -> failTest "validateMapKeysSorted should have rejected unsorted"++  -- columnDWithDefault: missing column substitutes the default.+  -- Build a writer that emits only (name, age); the reader+  -- expects (name, age, opt) and falls back on the default+  -- for the missing 'opt' column.+  let !partialEnc =+        AR.fieldE "name" (fst3 :: (Text, Int32, Text) -> Text) AR.utf8E+          <> AR.fieldE "age" (snd3 :: (Text, Int32, Text) -> Int32) AR.int32E+      !partialDec =+        (\n a -> (n, a, "" :: Text))+          <$> AR.columnD "name" AR.utf8D+          <*> AR.columnD "age" AR.int32D+      !partialTbl =+        AR.table partialEnc partialDec+          :: AR.Table (Text, Int32, Text)+      !partialRows =+        V.fromList+          [ ("Alice" :: Text, 30 :: Int32, "ignored" :: Text)+          , ("Bob", 45, "ignored")+          ]+      (!partialSch, !partialCols) = AR.encodeTable partialTbl partialRows+      !fullDec =+        (\n a o -> (n, a, o))+          <$> AR.columnD "name" AR.utf8D+          <*> AR.columnD "age" AR.int32D+          <*> AR.columnDWithDefault "opt" ("default" :: Text) AR.utf8D+  case AR.runRowDecoder fullDec (arrowFields partialSch) partialCols of+    Right got+      | V.toList got == [("Alice", 30, "default"), ("Bob", 45, "default")] ->+          putStrLn "OK: columnDWithDefault substitutes for missing column"+      | otherwise ->+          failTest $ "columnDWithDefault wrong values: " ++ show (V.toList got)+    Left e -> failTest $ "columnDWithDefault decode: " ++ e++  -- projectTable: pick a subset of columns by name+  let (!schWide, !colsWide) =+        let !enc =+              AR.fieldE "a" (\(x, _, _) -> x :: Int32) AR.int32E+                <> AR.fieldE "b" (\(_, y, _) -> y :: Int32) AR.int32E+                <> AR.fieldE "c" (\(_, _, z) -> z :: Int32) AR.int32E+            !dec =+              (,,)+                <$> AR.columnD "a" AR.int32D+                <*> AR.columnD "b" AR.int32D+                <*> AR.columnD "c" AR.int32D+            !tbl = AR.table enc dec :: AR.Table (Int32, Int32, Int32)+            !rs = V.fromList [(1, 10, 100), (2, 20, 200)]+        in AR.encodeTable tbl rs+  case AR.projectTable ["c", "a"] schWide colsWide of+    Just (sch', cols') -> do+      let !names = V.toList (V.map fieldName (arrowFields sch'))+      expect+        ("projectTable preserves order: got " ++ show names)+        (names == ["c", "a"])+      expect+        "projectTable yields matching column count"+        (V.length cols' == 2)+    Nothing -> failTest "projectTable returned Nothing for present cols"+  case AR.projectTable ["c", "missing"] schWide colsWide of+    Nothing -> putStrLn "OK: projectTable returns Nothing for missing column"+    Just _ -> failTest "projectTable should have returned Nothing"++  -- subsetTable: build a Table whose encoder emits only some+  -- columns+  let !custTbl =+        AR.table+          ( AR.fieldE "name" (fst :: (Text, Int32) -> Text) AR.utf8E+              <> AR.fieldE "age" (snd :: (Text, Int32) -> Int32) AR.int32E+          )+          ( (,)+              <$> AR.columnD "name" AR.utf8D+              <*> AR.columnD "age" AR.int32D+          )+          :: AR.Table (Text, Int32)+  case AR.subsetTable ["name"] custTbl of+    Just sub -> do+      let !rsSub = V.fromList [("Alice" :: Text, 30 :: Int32), ("Bob", 45)]+          (!schSub, !colsSub) = AR.encodeTable sub rsSub+      expect+        "subsetTable schema has 1 field"+        (V.length (arrowFields schSub) == 1)+      expect+        "subsetTable schema field is 'name'"+        (V.toList (V.map fieldName (arrowFields schSub)) == ["name"])+      expect+        "subsetTable encoded 1 column"+        (V.length colsSub == 1)+    Nothing -> failTest "subsetTable returned Nothing for ['name']"+  case AR.subsetTable ["nope"] custTbl of+    Nothing -> putStrLn "OK: subsetTable returns Nothing for missing column"+    Just _ -> failTest "subsetTable should have returned Nothing"+++projectionRoundTrip :: IO ()+projectionRoundTrip = do+  let !sch =+        Schema+          ( V.fromList+              [ plainField "a" False (AInt 32 True)+              , plainField "b" False (AInt 64 True)+              , plainField "c" False AUtf8+              ]+          )+          Little+          V.empty+          V.empty+      !batch =+        V.fromList+          [ primColumn PInt32 (VS.fromList ([1, 2, 3] :: [Int32]))+          , primColumn PInt64 (VS.fromList ([10, 20, 30] :: [Int64]))+          , fromTexts (V.fromList ["x", "y", "z"])+          ]+      !bytes = encodeArrowStream defaultWriteOptions sch [batch]+  case bytes >>= openStreamReader defaultReadOptions of+    Left e -> failTest $ "projection openStreamReader: " ++ e+    Right rd0 ->+      -- Ask for c then a, in that order: should drop b and reorder.+      case streamReaderProjected ["c", "a"] rd0 of+        Left e -> failTest $ "streamReaderProjected: " ++ e+        Right (projSch, batches')+          | length batches' == 1+          , [proj] <- batches'+          , V.length proj == 2+          , V.length (arrowFields projSch) == 2+          , fieldName (V.unsafeIndex (arrowFields projSch) 0) == "c"+          , fieldName (V.unsafeIndex (arrowFields projSch) 1) == "a"+          , V.unsafeIndex proj 0 == V.unsafeIndex batch 2+          , V.unsafeIndex proj 1 == V.unsafeIndex batch 0 ->+              putStrLn "OK: streamReaderProjected narrows + reorders"+          | otherwise ->+              failTest $+                "streamReaderProjected unexpected: "+                  ++ show batches'+++{- | Generic single-batch round-trip helper for the high-level+'encodeArrowStream' / 'decodeArrowStream' API.+-}+highLevelRoundTrip :: String -> Schema -> V.Vector ColumnArray -> IO ()+highLevelRoundTrip label sch cols = do+  let bytes = encodeArrowStream defaultWriteOptions sch [cols]+  case bytes >>= decodeArrowStream defaultReadOptions of+    Left e -> failTest $ label ++ ": decodeArrowStream: " ++ e+    Right (sch', batches)+      | sch' /= sch ->+          failTest $ label ++ ": schema mismatch"+      | [got] <- batches ->+          if got == cols+            then putStrLn $ "OK: high-level round-trip " ++ label+            else+              failTest $+                label+                  ++ ": column mismatch\n got: "+                  ++ show (V.toList got)+                  ++ "\n exp: "+                  ++ show (V.toList cols)+      | otherwise ->+          failTest $+            label+              ++ ": expected 1 batch, got "+              ++ show (length batches)+++-- | Build a simple leaf field with no children.+plainField :: Text -> Bool -> ArrowType -> Field+plainField nm nullable ty = Field nm nullable ty V.empty Nothing V.empty+++-- | Field with explicit children, no dictionary.+nestedField :: Text -> Bool -> ArrowType -> V.Vector Field -> Field+nestedField nm nullable ty children = Field nm nullable ty children Nothing V.empty+++-- | Round-trip a pre-built Field/ColumnArray pair.+roundTripNested :: String -> Field -> ColumnArray -> IO ()+roundTripNested label field col = do+  let !schema =+        Schema+          { arrowEndianness = Little+          , arrowFields = V.singleton field+          , arrowMetadata = V.empty+          , arrowFeatures = V.empty+          }+      !stream = writeArrowStream schema (V.singleton (V.singleton col))+  case stream >>= readArrowStream of+    Left e -> failTest (label ++ ": readArrowStream: " ++ e)+    Right as -> do+      expect+        (label ++ ": batch count == 1")+        (V.length (asBatches as) == 1)+      let (rb, body) = V.unsafeIndex (asBatches as) 0+      case decodeRecordBatch defaultReadOptions (asSchema as) rb body of+        Left e -> failTest (label ++ ": materialize: " ++ e)+        Right cols -> do+          expect (label ++ ": column count == 1") (V.length cols == 1)+          let !got = V.unsafeIndex cols 0+          expect+            (label ++ ": null count matches")+            (nullCount got == nullCount col)+          when (got /= col) $+            failTest+              ( label+                  ++ ": got "+                  ++ show got+                  ++ ", expected "+                  ++ show col+              )+      expect (label ++ ": nested round-trip preserves column") True+++{- | Round-trip a single flat (non-nested) column through a single-field+single-batch Arrow stream.+-}+roundTripPrim :: String -> ColumnArray -> IO ()+roundTripPrim label col = do+  let !ty = inferArrowType col+      !nullable = nullCount col > 0+      !schema =+        Schema+          { arrowEndianness = Little+          , arrowFields =+              V.singleton+                Field+                  { fieldName = T.pack label+                  , fieldNullable = nullable+                  , fieldType = ty+                  , fieldChildren = V.empty+                  , fieldDictionary = Nothing+                  , fieldMetadata = V.empty+                  }+          , arrowMetadata = V.empty+          , arrowFeatures = V.empty+          }+      !stream = writeArrowStream schema (V.singleton (V.singleton col))+  case stream >>= readArrowStream of+    Left e -> failTest (label ++ ": readArrowStream: " ++ e)+    Right as -> do+      expect+        (label ++ ": schema endianness")+        (arrowEndianness (asSchema as) == Little)+      expect+        (label ++ ": batch count == 1")+        (V.length (asBatches as) == 1)+      let (rb, body) = V.unsafeIndex (asBatches as) 0+      case decodeRecordBatch defaultReadOptions (asSchema as) rb body of+        Left e -> failTest (label ++ ": materialize: " ++ e)+        Right cols -> do+          expect (label ++ ": column count == 1") (V.length cols == 1)+          let !got = V.unsafeIndex cols 0+          expect+            (label ++ ": row count matches")+            (columnLength got == columnLength col)+          expect+            (label ++ ": null count and validity slot match")+            ( nullCount got == nullCount col+                && hasValiditySlot got == hasValiditySlot col+            )+          when (got /= col) $+            failTest (label ++ ": got " ++ show got ++ ", expected " ++ show col)+      expect (label ++ ": round-trip preserves column") True+++{- | Derive the appropriate 'ArrowType' for a 'ColumnArray' variant+the test driver feeds in. Used only to build per-test schemas.+-}+inferArrowType :: ColumnArray -> ArrowType+inferArrowType = \case+  ColInt8 _ _ -> AInt 8 True+  ColInt16 _ _ -> AInt 16 True+  ColInt32 _ _ -> AInt 32 True+  ColInt64 _ _ -> AInt 64 True+  ColUInt8 _ _ -> AInt 8 False+  ColUInt16 _ _ -> AInt 16 False+  ColUInt32 _ _ -> AInt 32 False+  ColUInt64 _ _ -> AInt 64 False+  ColFloat16 _ _ -> AFloatingPoint Half+  ColFloat _ _ -> AFloatingPoint Single+  ColDouble _ _ -> AFloatingPoint DoublePrecision+  ColBool _ _ -> ABool+  ColUtf8 {} -> AUtf8+  ColBinary {} -> ABinary+  ColLargeUtf8 {} -> ALargeUtf8+  ColLargeBinary {} -> ALargeBinary+  ColFixedSizeBinary w _ _ _ -> AFixedSizeBinary w+  ColDate32 _ _ -> ADate DateDay+  ColDate64 _ _ -> ADate DateMillisecond+  ColTime32 _ _ -> ATime Second 32+  ColTime64 _ _ -> ATime Microsecond 64+  ColTimestamp _ _ -> ATimestamp Nanosecond Nothing+  ColDuration _ _ -> ADuration Nanosecond+  ColDecimal32 p s _ _ -> ADecimal32 p s+  ColDecimal64 p s _ _ -> ADecimal64 p s+  ColDecimal128 p s _ _ -> ADecimal p s+  ColDecimal256 p s _ _ -> ADecimal256 p s+  ColIntervalYearMonth _ _ -> AInterval YearMonth+  ColIntervalDayTime _ _ -> AInterval DayTime+  ColIntervalMonthDayNano _ _ -> AInterval MonthDayNano+  -- The test driver doesn't invoke inferArrowType for nested columns;+  -- those are fed through a dedicated roundTripNested helper below.+  other -> error ("inferArrowType: unsupported: " ++ columnTag other)+++-- | Unwrap a validated fixture constructor, failing loudly on 'Left'.+fixture :: Either String ColumnArray -> ColumnArray+fixture = either (error . ("bad fixture: " ++)) id+++expect :: String -> Bool -> IO ()+expect label True = putStrLn ("OK: " ++ label)+expect label False = failTest ("FAIL: " ++ label)+++failTest :: String -> IO ()+failTest msg = do+  putStrLn msg+  exitFailure
+ test/Test/Arrow/Core.hs view
@@ -0,0 +1,944 @@+{-# LANGUAGE GADTs #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PatternSynonyms #-}++{- | Properties of the column core ("Arrow.Column", its buffers and+builders) against a boxed row model.++Columns are generated together with their model rows through the+public construction API (the @from*@ conversions, the builders and the+@mk*@ constructors), usually as a slice of a longer column so bitmaps+start at non-zero bit offsets and offsets at non-zero bases. Every+row operation is then checked against the same operation on the model.+-}+module Test.Arrow.Core (tests) where++import Arrow.Column+import Arrow.Column.Internal (Bitmap (..), Validity (..), checkValidity, columnBuffers, concatBitmaps, copyBitmap, integralPrim, IntegralPrim (..), sliceBitmap, takeBitmap)+import Arrow.Record qualified as AR+import Control.Exception (ErrorCall, displayException, evaluate, try)+import Control.Monad (forM, replicateM)+import Control.Monad.ST (runST)+import Data.Bits (setBit)+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Internal qualified as BSI+import Data.Either (isLeft, isRight)+import Data.Foldable (foldl')+import Data.Int (Int32, Int64)+import Data.List (isInfixOf, sortOn)+import Data.Maybe (catMaybes, isJust, isNothing)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Text.Encoding qualified as TE+import Data.Vector qualified as V+import Data.Vector.Generic qualified as VG+import Data.Vector.Storable qualified as VS+import Data.Word (Word8)+import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)+import Foreign.Ptr (minusPtr)+import GHC.Float (castDoubleToWord64, castFloatToWord32, castWord32ToFloat, castWord64ToDouble)+import Hedgehog+import Hedgehog.Gen qualified as Gen+import Hedgehog.Range qualified as Range+++tests :: IO Bool+tests =+  checkParallel $+    Group+      "Arrow.Core"+      [ ("bitmaps at arbitrary bit offsets match a bool-list model", prop_bitmaps)+      , ("validity is normalised: no nulls means Nothing", prop_validityNormalised)+      , ("generated columns observe as their model", prop_observeModel)+      , ("typed accessors agree with the model and reject out-of-range rows", prop_accessors)+      , ("boxed conversions and Arrow.Record utf8D agree with the model on every kind they accept", prop_conversions)+      , ("dense binary construction copies sliced rows and empty values", prop_denseBinary)+      , ("binary and text construction check offset overflow before allocating", prop_denseBinaryOffsetLimit)+      , ("builders build what the from* conversions build, and reset on freeze", prop_builders)+      , ("Eq is reflexive and agrees with the model", prop_eq)+      , ("slice keeps the clamped window", prop_slice)+      , ("slice of slice composes", prop_sliceOfSlice)+      , ("concat of adjacent slices is the column", prop_concatAdjacent)+      , ("concat appends rows", prop_concatAppends)+      , ("take gathers rows and rejects bad indices", prop_take)+      , ("copyColumn is equal and shares no memory", prop_copyColumn)+      , ("expandDictionary keeps rows", prop_expandDictionary)+      , ("maskValidity nulls masked rows", prop_maskValidity)+      , ("map key order agrees with the model", prop_mapKeys)+      , ("mk* constructors reject invalid buffers", prop_mkRejects)+      ]+++-- ============================================================+-- Model+-- ============================================================++data Value+  = VNull+  | VInt Integer+  | VBits Integer+  | VBool Bool+  | VBytes ByteString+  | VText Text+  | VPair Int32 Int32+  | VTriple Int32 Int32 Int64+  | VList [Value]+  | VStruct [Value]+  | VMap [(Value, Value)]+  | VUnion Int Value+  deriving stock (Eq, Show)+++data Ty+  = TPrim SomePrimType+  | TBool+  | TUtf8+  | TBinary+  | TLargeUtf8+  | TLargeBinary+  | TFixed Int+  | TUtf8View+  | TBinaryView+  | TNull+  | TStruct [Ty]+  | TList Ty+  | TLargeList Ty+  | -- | Large (64-bit) offsets when True.+    TListView Bool Ty+  | TFixedList Int Ty+  | TMap Ty+  | TDense [Ty]+  | TSparse [Ty]+  | TDict SomePrimType Ty+  | TRee SomePrimType Ty+  deriving stock (Show)+++primTypes :: [SomePrimType]+primTypes =+  [ SomePrimType PInt8+  , SomePrimType PInt16+  , SomePrimType PInt32+  , SomePrimType PInt64+  , SomePrimType PUInt8+  , SomePrimType PUInt16+  , SomePrimType PUInt32+  , SomePrimType PUInt64+  , SomePrimType PFloat16+  , SomePrimType PFloat+  , SomePrimType PDouble+  , SomePrimType PDate32+  , SomePrimType PDate64+  , SomePrimType PTime32+  , SomePrimType PTime64+  , SomePrimType PTimestamp+  , SomePrimType PDuration+  , SomePrimType PIntervalYearMonth+  , SomePrimType PIntervalDayTime+  , SomePrimType PIntervalMonthDayNano+  , SomePrimType (PDecimal32 9 2)+  , SomePrimType (PDecimal64 18 4)+  , SomePrimType (PDecimal128 38 4)+  , SomePrimType (PDecimal256 76 10)+  ]+++keyTypes :: [SomePrimType]+keyTypes = [SomePrimType PInt8, SomePrimType PInt16, SomePrimType PInt32, SomePrimType PInt64, SomePrimType PUInt8, SomePrimType PUInt16, SomePrimType PUInt32, SomePrimType PUInt64]+++genTy :: Int -> Gen Ty+genTy d =+  Gen.frequency $+    [ (6, TPrim <$> Gen.element primTypes)+    , (2, pure TBool)+    , (2, pure TUtf8)+    , (1, pure TBinary)+    , (1, pure TLargeUtf8)+    , (1, pure TLargeBinary)+    , (1, TFixed <$> Gen.int (Range.linear 0 5))+    , (1, pure TUtf8View)+    , (1, pure TBinaryView)+    , (1, pure TNull)+    ]+      ++ if d <= 0+        then []+        else+          [ (1, TStruct <$> Gen.list (Range.linear 0 3) sub)+          , (1, TList <$> sub)+          , (1, TLargeList <$> sub)+          , (1, TListView <$> Gen.bool <*> sub)+          , (1, TFixedList <$> Gen.int (Range.linear 0 3) <*> sub)+          , (1, TMap <$> sub)+          , (1, TDense <$> Gen.list (Range.linear 1 3) sub)+          , (1, TSparse <$> Gen.list (Range.linear 1 3) sub)+          , (1, TDict <$> Gen.element keyTypes <*> genTy 0)+          , (1, TRee <$> Gen.element [SomePrimType PInt16, SomePrimType PInt32, SomePrimType PInt64] <*> genTy 0)+          ]+  where+    sub = genTy (d - 1)+++genPrim :: PrimType a -> Gen (a, Value)+genPrim = \case+  PInt8 -> int (Gen.int8 Range.linearBounded)+  PInt16 -> int (Gen.int16 Range.linearBounded)+  PInt32 -> int (Gen.int32 Range.linearBounded)+  PInt64 -> int (Gen.int64 Range.linearBounded)+  PUInt8 -> int (Gen.word8 Range.linearBounded)+  PUInt16 -> int (Gen.word16 Range.linearBounded)+  PUInt32 -> int (Gen.word32 Range.linearBounded)+  PUInt64 -> int (Gen.word64 Range.linearBounded)+  PFloat16 -> (\w -> (Float16 w, VBits (toInteger w))) <$> Gen.word16 Range.linearBounded+  PFloat -> (\w -> (castWord32ToFloat w, VBits (toInteger w))) <$> Gen.word32 Range.linearBounded+  PDouble -> (\w -> (castWord64ToDouble w, VBits (toInteger w))) <$> Gen.word64 Range.linearBounded+  PDate32 -> int (Gen.int32 Range.linearBounded)+  PDate64 -> int (Gen.int64 Range.linearBounded)+  PTime32 -> int (Gen.int32 Range.linearBounded)+  PTime64 -> int (Gen.int64 Range.linearBounded)+  PTimestamp -> int (Gen.int64 Range.linearBounded)+  PDuration -> int (Gen.int64 Range.linearBounded)+  PIntervalYearMonth -> int (Gen.int32 Range.linearBounded)+  PIntervalDayTime -> (\a b -> (IntervalDayTime a b, VPair a b)) <$> Gen.int32 Range.linearBounded <*> Gen.int32 Range.linearBounded+  PIntervalMonthDayNano ->+    (\a b c -> (IntervalMonthDayNano a b c, VTriple a b c))+      <$> Gen.int32 Range.linearBounded+      <*> Gen.int32 Range.linearBounded+      <*> Gen.int64 Range.linearBounded+  PDecimal32 _ _ -> int (Gen.int32 Range.linearBounded)+  PDecimal64 _ _ -> int (Gen.int64 Range.linearBounded)+  PDecimal128 _ _ -> (\x -> (decimal128FromInteger x, VInt x)) <$> Gen.integral (Range.linearFrom 0 (negate (2 ^ (127 :: Int))) (2 ^ (127 :: Int) - 1))+  PDecimal256 _ _ -> (\x -> (decimal256FromInteger x, VInt x)) <$> Gen.integral (Range.linearFrom 0 (negate (2 ^ (255 :: Int))) (2 ^ (255 :: Int) - 1))+  where+    int :: Integral b => Gen b -> Gen (b, Value)+    int = fmap (\x -> (x, VInt (toInteger x)))+++primValue :: PrimType a -> a -> Value+primValue t x = case t of+  PInt8 -> VInt (toInteger x)+  PInt16 -> VInt (toInteger x)+  PInt32 -> VInt (toInteger x)+  PInt64 -> VInt (toInteger x)+  PUInt8 -> VInt (toInteger x)+  PUInt16 -> VInt (toInteger x)+  PUInt32 -> VInt (toInteger x)+  PUInt64 -> VInt (toInteger x)+  PFloat16 -> let Float16 w = x in VBits (toInteger w)+  PFloat -> VBits (toInteger (castFloatToWord32 x))+  PDouble -> VBits (toInteger (castDoubleToWord64 x))+  PDate32 -> VInt (toInteger x)+  PDate64 -> VInt (toInteger x)+  PTime32 -> VInt (toInteger x)+  PTime64 -> VInt (toInteger x)+  PTimestamp -> VInt (toInteger x)+  PDuration -> VInt (toInteger x)+  PIntervalYearMonth -> VInt (toInteger x)+  PIntervalDayTime -> let IntervalDayTime a b = x in VPair a b+  PIntervalMonthDayNano -> let IntervalMonthDayNano a b c = x in VTriple a b c+  PDecimal32 _ _ -> VInt (toInteger x)+  PDecimal64 _ _ -> VInt (toInteger x)+  PDecimal128 _ _ -> VInt (decimal128ToInteger x)+  PDecimal256 _ _ -> VInt (decimal256ToInteger x)+++maybeOf :: Gen a -> Gen (Maybe a)+maybeOf g = Gen.frequency [(1, pure Nothing), (4, Just <$> g)]+++genValid :: Int -> Gen [Bool]+genValid n = Gen.frequency [(1, pure (replicate n True)), (3, replicateM n (Gen.frequency [(1, pure False), (3, pure True)]))]+++validityOf :: [Bool] -> Maybe Validity+validityOf = validityFromBools . V.fromList+++genText :: Gen Text+genText = Gen.text (Range.linear 0 20) Gen.unicode+++genBytes :: Gen ByteString+genBytes = Gen.bytes (Range.linear 0 20)+++right :: Either String a -> a+right = either error id+++-- | A column of @n@ rows and its model, usually a slice of a longer column.+genColumnOf :: Ty -> Int -> Gen ([Value], ColumnArray)+genColumnOf ty n = do+  pre <- Gen.frequency [(1, pure 0), (3, Gen.int (Range.linear 0 11))]+  post <- Gen.int (Range.linear 0 3)+  (m, c) <- genRaw ty (pre + n + post)+  pure (take n (drop pre m), sliceColumnArray pre n c)+++genRaw :: Ty -> Int -> Gen ([Value], ColumnArray)+genRaw ty n = case ty of+  TPrim (SomePrimType t) -> withPrim t $ do+    rows <- replicateM n (maybeOf (genPrim t))+    let xs = map (fmap fst) rows+    viaBuilder <- Gen.bool+    let c =+          if viaBuilder+            then runST $ do+              b <- newPrimBuilder t 0+              mapM_ (appendPrimMaybe b) xs+              freezeBuilder b+            else fromMaybes t (V.fromList xs)+    pure (map (maybe VNull snd) rows, c)+  TBool -> do+    rows <- replicateM n (maybeOf Gen.bool)+    viaBuilder <- Gen.bool+    let c =+          if viaBuilder+            then runST $ do+              b <- newBoolBuilder 0+              mapM_ (appendBoolMaybe b) rows+              freezeBuilder b+            else fromMaybeBools (V.fromList rows)+    pure (map (maybe VNull VBool) rows, c)+  TUtf8 -> do+    rows <- replicateM n (maybeOf genText)+    how <- Gen.int (Range.constant 0 2)+    c <- case how of+      0 -> pure (fromMaybeTexts (V.fromList rows))+      1 -> pure $ runST $ do+        b <- newUtf8Builder 0+        mapM_ (appendTextMaybe b) rows+        freezeBuilder b+      _ -> do+        prefix <- genText+        let base = BS.length (TE.encodeUtf8 prefix)+            lens = map (maybe 0 (BS.length . TE.encodeUtf8)) rows+            dat = TE.encodeUtf8 (prefix <> T.concat (map (maybe "" id) rows))+            offs = VS.fromList (map fromIntegral (scanl (+) base lens)) :: VS.Vector Int32+        pure (right (mkUtf8 (validityOf (map isJust rows)) offs dat))+    pure (map (maybe VNull VText) rows, c)+  TLargeUtf8 -> do+    rows <- replicateM n (maybeOf genText)+    viaBuilder <- Gen.bool+    let c =+          if viaBuilder+            then runST $ do+              b <- newLargeUtf8Builder 0+              mapM_ (appendTextMaybe b) rows+              freezeBuilder b+            else fromMaybeLargeTexts (V.fromList rows)+    pure (map (maybe VNull VText) rows, c)+  TBinary -> do+    rows <- replicateM n (maybeOf genBytes)+    viaBuilder <- Gen.bool+    let c =+          if viaBuilder+            then runST $ do+              b <- newBinaryBuilder 0+              mapM_ (appendBytesMaybe b) rows+              freezeBuilder b+            else fromMaybeByteStrings (V.fromList rows)+    pure (map (maybe VNull VBytes) rows, c)+  TLargeBinary -> do+    rows <- replicateM n (maybeOf genBytes)+    viaBuilder <- Gen.bool+    let c =+          if viaBuilder+            then runST $ do+              b <- newLargeBinaryBuilder 0+              mapM_ (appendBytesMaybe b) rows+              freezeBuilder b+            else fromMaybeLargeByteStrings (V.fromList rows)+    pure (map (maybe VNull VBytes) rows, c)+  TFixed w -> do+    rows <- replicateM n (maybeOf (Gen.bytes (Range.singleton w)))+    pure (map (maybe VNull VBytes) rows, right (fromMaybeFixedSizeBinary w (V.fromList rows)))+  TUtf8View -> do+    rows <- replicateM n (maybeOf (Gen.text (Range.linear 0 30) Gen.unicode))+    pure (map (maybe VNull VText) rows, fromMaybeUtf8View (V.fromList rows))+  TBinaryView -> do+    rows <- replicateM n (maybeOf (Gen.bytes (Range.linear 0 30)))+    pure (map (maybe VNull VBytes) rows, fromMaybeBinaryView (V.fromList rows))+  TNull -> pure (replicate n VNull, ColNull n)+  TStruct tys -> do+    valid <- genValid n+    children <- forM tys $ \t -> do+      extra <- Gen.int (Range.linear 0 2)+      genColumnOf t (n + extra)+    let names = map (\i -> T.pack ("f" ++ show i)) [0 .. length tys - 1]+        c = right (mkStruct n (validityOf valid) (V.fromList (zip names (map snd children))))+        row i ok = if ok then VStruct (map (\(cm, _) -> cm !! i) children) else VNull+    pure (zipWith row [0 ..] valid, c)+  TList ct -> do+    (m, offs, child) <- genListParts ct n+    valid <- genValid n+    pure (listModel valid m offs, right (mkList (validityOf valid) (VS.fromList (map fromIntegral offs)) child))+  TLargeList ct -> do+    (m, offs, child) <- genListParts ct n+    valid <- genValid n+    pure (listModel valid m offs, right (mkLargeList (validityOf valid) (VS.fromList (map fromIntegral offs)) child))+  TListView large ct -> do+    len <- Gen.int (Range.linear 0 8)+    (cm, child) <- genColumnOf ct len+    views <- replicateM n $ do+      off <- Gen.int (Range.linear 0 len)+      size <- Gen.int (Range.linear 0 (len - off))+      pure (off, size)+    valid <- genValid n+    let row (off, size) ok = if ok then VList (take size (drop off cm)) else VNull+        offs = VS.fromList (map (fromIntegral . fst) views) :: VS.Vector Int32+        sizes = VS.fromList (map (fromIntegral . snd) views) :: VS.Vector Int32+    let c =+          if large+            then right (mkLargeListView (validityOf valid) (VS.map fromIntegral offs) (VS.map fromIntegral sizes) child)+            else right (mkListView (validityOf valid) offs sizes child)+    pure (zipWith row views valid, c)+  TFixedList w ct -> do+    extra <- Gen.int (Range.linear 0 2)+    (cm, child) <- genColumnOf ct (n * w + extra)+    valid <- genValid n+    let row i ok = if ok then VList (take w (drop (i * w) cm)) else VNull+    pure (zipWith row [0 ..] valid, right (mkFixedSizeList w n (validityOf valid) child))+  TMap vt -> do+    lens <- replicateM n (Gen.int (Range.linear 0 3))+    base <- Gen.int (Range.linear 0 3)+    let total = base + sum lens+    (km, keys) <- genColumnOf (TPrim (SomePrimType PInt32)) total+    (vm, vals) <- genColumnOf vt total+    valid <- genValid n+    let offs = scanl (+) base lens+        row (s, e) ok = if ok then VMap (take (e - s) (drop s (zip km vm))) else VNull+    pure (zipWith row (zip offs (drop 1 offs)) valid, right (mkMap (validityOf valid) (VS.fromList (map fromIntegral offs)) keys vals))+  TDense tys -> do+    children <- forM tys $ \t -> do+      len <- Gen.int (Range.linear 1 5)+      genColumnOf t len+    picks <- replicateM n $ do+      ci <- Gen.int (Range.linear 0 (length tys - 1))+      off <- Gen.int (Range.linear 0 (length (fst (children !! ci)) - 1))+      pure (ci, off)+    let c =+          right+            ( mkDenseUnion+                (VS.fromList (map (fromIntegral . fst) picks))+                (VS.fromList (map (fromIntegral . snd) picks))+                (V.fromList (map snd children))+            )+    pure (map (\(ci, off) -> VUnion ci (fst (children !! ci) !! off)) picks, c)+  TSparse tys -> do+    children <- forM tys $ \t -> genColumnOf t n+    picks <- replicateM n (Gen.int (Range.linear 0 (length tys - 1)))+    let c = right (mkSparseUnion (VS.fromList (map fromIntegral picks)) (V.fromList (map snd children)))+    pure (zipWith (\i ci -> VUnion ci (fst (children !! ci) !! i)) [0 ..] picks, c)+  TDict (SomePrimType kt) vt -> case integralPrim kt of+    Nothing -> error "key type"+    Just IntegralPrim -> do+      k <- Gen.int (Range.linear 0 6)+      (vm, vals) <- genColumnOf vt k+      ks <- replicateM n (if k == 0 then pure Nothing else maybeOf (Gen.int (Range.linear 0 (k - 1))))+      let keys = fromMaybes kt (V.fromList (map (fmap fromIntegral) ks))+      pure (map (maybe VNull (vm !!)) ks, right (mkDictionary 7 keys vals))+  TRee (SomePrimType rt) vt -> case integralPrim rt of+    Nothing -> error "run end type"+    Just IntegralPrim -> do+      lens <- genParts n+      (vm, vals) <- genColumnOf vt (length lens)+      let re = primColumn rt (VS.fromList (map fromIntegral (scanl1 (+) lens)))+      pure (concat (zipWith replicate lens vm), right (mkRunEndEncoded re vals))+++genParts :: Int -> Gen [Int]+genParts n+  | n <= 0 = pure []+  | otherwise = do+      k <- Gen.int (Range.linear 1 (min n 4))+      (k :) <$> genParts (n - k)+++-- | Row lengths, a non-zero base, and a child long enough to cover them.+genListParts :: Ty -> Int -> Gen ([Value], [Int], ColumnArray)+genListParts ct n = do+  lens <- replicateM n (Gen.int (Range.linear 0 3))+  base <- Gen.int (Range.linear 0 3)+  extra <- Gen.int (Range.linear 0 2)+  (cm, child) <- genColumnOf ct (base + sum lens + extra)+  pure (cm, scanl (+) base lens, child)+++listModel :: [Bool] -> [Value] -> [Int] -> [Value]+listModel valid cm offs = zipWith (\(s, e) ok -> if ok then VList (take (e - s) (drop s cm)) else VNull) (zip offs (drop 1 offs)) valid+++-- ============================================================+-- Observation through the public accessors+-- ============================================================++observe :: ColumnArray -> [Value]+observe c = map (rowValue c) [0 .. columnLength c - 1]+++rowValue :: ColumnArray -> Int -> Value+rowValue c i+  | not (isValidAt (validity c) i) = VNull+  | otherwise = case c of+      ColNull _ -> VNull+      ColPrim t _ _ -> case asPrim t c of+        Just arr -> withPrim t (maybe VNull (primValue t) (primAt arr i))+        Nothing -> error "asPrim"+      ColBool _ _ -> maybe VNull VBool (boolAt c i)+      ColUtf8 {} -> maybe VNull VText (anyTextAt c i)+      ColLargeUtf8 {} -> maybe VNull VText (anyTextAt c i)+      ColUtf8View {} -> maybe VNull VText (anyTextAt c i)+      ColBinary {} -> maybe VNull VBytes (anyBytesAt c i)+      ColLargeBinary {} -> maybe VNull VBytes (anyBytesAt c i)+      ColBinaryView {} -> maybe VNull VBytes (anyBytesAt c i)+      ColFixedSizeBinary {} -> maybe VNull VBytes (anyBytesAt c i)+      ColStruct _ _ cs -> VStruct (map (\(_, x) -> rowValue x i) (V.toList cs))+      ColList _ _ x -> listOf x+      ColLargeList _ _ x -> listOf x+      ColListView _ _ _ x -> listOf x+      ColLargeListView _ _ _ x -> listOf x+      ColFixedSizeList _ _ _ x -> listOf x+      ColMap _ _ k x -> case listRange c i of+        Just (ChildRange s l) -> VMap (map (\q -> (rowValue k (s + q), rowValue x (s + q))) [0 .. l - 1])+        Nothing -> VNull+      ColDenseUnion t o cs ->+        let ci = fromIntegral (t VS.! i)+        in VUnion ci (rowValue (cs V.! ci) (fromIntegral (o VS.! i)))+      ColSparseUnion t cs ->+        let ci = fromIntegral (t VS.! i)+        in VUnion ci (rowValue (cs V.! ci) i)+      ColDictionary _ _ vals -> maybe VNull (rowValue vals) (dictKeyAt c i)+      ColRunEndEncoded off _ re vals ->+        let ends = runEnds re+        in rowValue vals (length (takeWhile (<= off + i) ends))+  where+    listOf x = case listRange c i of+      Just (ChildRange s l) -> VList (map (rowValue x) [s .. s + l - 1])+      Nothing -> VNull+++runEnds :: ColumnArray -> [Int]+runEnds = \case+  ColInt16 _ xs -> map fromIntegral (VS.toList xs)+  ColInt32 _ xs -> map fromIntegral (VS.toList xs)+  ColInt64 _ xs -> map fromIntegral (VS.toList xs)+  _ -> []+++-- ============================================================+-- Properties+-- ============================================================++genColumn :: Gen (Ty, [Value], ColumnArray)+genColumn = do+  ty <- genTy 2+  n <- Gen.int (Range.linear 0 12)+  (m, c) <- genColumnOf ty n+  pure (ty, m, c)+++-- | Bytes holding @bools@ at bit offset @k@, with random bits around them.+bitmapAt :: Int -> [Bool] -> [Bool] -> [Bool] -> Bitmap+bitmapAt k pre bools post =+  let bits = take k (pre ++ repeat False) ++ bools ++ post+      nbytes = (length bits + 7) `div` 8+      byte j = foldl' (\acc q -> if bitIx (j * 8 + q) then setBit acc q else acc) (0 :: Word8) [0 .. 7]+      bitIx p = p < length bits && bits !! p+  in Bitmap (BS.pack (map byte [0 .. nbytes - 1])) k (length bools)+++prop_bitmaps :: Property+prop_bitmaps = withTests 300 . property $ do+  k <- forAll (Gen.int (Range.linear 0 70))+  pre <- forAll (Gen.list (Range.singleton k) Gen.bool)+  bools <- forAll (Gen.list (Range.linear 0 200) Gen.bool)+  post <- forAll (Gen.list (Range.linear 0 9) Gen.bool)+  let b = bitmapAt k pre bools post+      n = length bools+  bitmapToBools b === bools+  map (bitAt b) [0 .. n - 1] === bools+  bitmapSetCount b === length (filter id bools)+  s <- forAll (Gen.int (Range.linear 0 n))+  l <- forAll (Gen.int (Range.linear 0 (n - s)))+  bitmapToBools (sliceBitmap s l b) === take l (drop s bools)+  bitmapSetCount (sliceBitmap s l b) === length (filter id (take l (drop s bools)))+  let cp = copyBitmap b+  bitmapOffset cp === 0+  bitmapToBools cp === bools+  bools2 <- forAll (Gen.list (Range.linear 0 70) Gen.bool)+  k2 <- forAll (Gen.int (Range.linear 0 13))+  let b2 = bitmapAt k2 (replicate k2 True) bools2 []+  bitmapToBools (concatBitmaps [sliceBitmap s l b, b2, b]) === take l (drop s bools) ++ bools2 ++ bools+  ix <- forAll (if n == 0 then pure [] else Gen.list (Range.linear 0 40) (Gen.int (Range.linear 0 (n - 1))))+  let (tb, set) = takeBitmap (VS.fromList ix) b+  bitmapToBools tb === map (bools !!) ix+  set === length (filter (bools !!) ix)+  bitmapToBools (bitmapGenerate n (bools !!)) === bools+++prop_validityNormalised :: Property+prop_validityNormalised = withTests 300 . property $ do+  t <- forAll (Gen.element primTypes)+  case t of+    SomePrimType pt -> withPrim pt $ do+      rows <- forAll (Gen.list (Range.linear 0 40) (Gen.frequency [(1, pure Nothing), (6, pure (Just ()))]))+      xs <- forAll (traverse (traverse (const (fst <$> genPrim pt))) rows)+      let c = fromMaybes pt (V.fromList xs)+          nulls = length (filter isNothing xs)+      nullCount c === nulls+      assert (isNothing (validity c) == (nulls == 0))+      -- A window without nulls loses its validity.+      let firstNull = length (takeWhile isJust xs)+      assert (isNothing (validity (sliceColumnArray 0 firstNull c)))+      case validity c of+        Nothing -> success+        Just v -> do+          validityNullCount v === nulls+          assert (isLeft (checkValidity "t" (length xs) (Just (Validity (validityBits v) (nulls + 1)))))+          assert (isLeft (checkValidity "t" (length xs + 1) (Just v)))+          checkValidity "t" (length xs) (Just v) === Right (Just v)+  bools <- forAll (Gen.list (Range.linear 0 50) Gen.bool)+  let b = bitmapFromBools (V.fromList bools)+  fmap validityNullCount (mkValidity b) === (if and bools then Nothing else Just (length (filter not bools)))+++prop_observeModel :: Property+prop_observeModel = withTests 500 . property $ do+  (_, m, c) <- forAll genColumn+  columnLength c === length m+  observe c === m+++prop_accessors :: Property+prop_accessors = withTests 300 . property $ do+  t <- forAll (Gen.element primTypes)+  n <- forAll (Gen.int (Range.linear 0 30))+  case t of+    SomePrimType pt -> withPrim pt $ do+      (m, c) <- forAll (genColumnOf (TPrim t) n)+      arr <- maybe (failure) pure (asPrim pt c)+      map (fmap (primValue pt) . primAt arr) [0 .. n - 1] === map (\v -> if v == VNull then Nothing else Just v) m+      primAt arr (-1) === Nothing+      primAt arr n === Nothing+      map (fmap (primValue pt)) (VG.toList (toMaybeVector arr)) === map (\v -> if v == VNull then Nothing else Just v) m+      VS.length (toStorable arr) === n+  (tm, tc) <- forAll (genColumnOf TUtf8 n)+  u <- maybe failure pure (asUtf8 tc)+  map (fmap VText . textAt u) [0 .. n - 1] === map (\v -> if v == VNull then Nothing else Just v) tm+  map (fmap (VText . TE.decodeUtf8) . bytesAt (utf8Bytes u)) [0 .. n - 1] === map (\v -> if v == VNull then Nothing else Just v) tm+  textAt u n === Nothing+  fmap VG.toList (toTextVector tc) === Right (map (\case VText x -> Just x; _ -> Nothing) tm)+  (bm, bc) <- forAll (genColumnOf TBool n)+  map (boolAt bc) [-1 .. n] === Nothing : map (\case VBool x -> Just x; _ -> Nothing) bm ++ [Nothing]+  where+    utf8Bytes (Utf8Array b) = b+++prop_conversions :: Property+prop_conversions = withTests 500 . property $ do+  n <- forAll (Gen.int (Range.linear 0 30))+  key <- forAll (Gen.element keyTypes)+  w <- forAll (Gen.int (Range.linear 0 3))+  large <- forAll Gen.bool+  let int32 = TPrim (SomePrimType PInt32)+  ty <-+    forAll . Gen.element $+      [ TUtf8, TLargeUtf8, TUtf8View, TDict key TUtf8, TDict key TLargeUtf8+      , TBinary, TLargeBinary, TFixed w, TBinaryView, TDict key TBinary+      , TBool+      , TList int32, TLargeList int32, TListView large int32, TFixedList w int32+      ]+  (m, c) <- forAll (genColumnOf ty n)+  let texts = map (\case VText x -> Just x; _ -> Nothing) m+      bytes = map (\case VText x -> Just (TE.encodeUtf8 x); VBytes x -> Just x; _ -> Nothing) m+      lists = map (\case VList xs -> Just xs; _ -> Nothing) m+      leaf = \case+        TDict _ t -> leaf t+        t -> t+      textKind = case leaf ty of TUtf8 -> True; TLargeUtf8 -> True; TUtf8View -> True; _ -> False+      listKind = case ty of TList _ -> True; TLargeList _ -> True; TListView _ _ -> True; TFixedList _ _ -> True; _ -> False+  case ty of+    TBool -> fmap VG.toList (toBoolVector c) === Right (map (\case VBool x -> Just x; _ -> Nothing) m)+    _+      | listKind -> do+          fmap (map (fmap (map (maybe VNull (VInt . toInteger)) . VG.toList)) . VG.toList) (toListVector int32s c) === Right lists+          -- A child conversion that returns too few rows is an error, not a short slice.+          case ty of+            TList _ | any (maybe False (not . null)) lists -> assert (isLeft (toListVector (fmap (VG.drop 1) . int32s) c))+            _ -> success+      | otherwise -> do+          fmap VG.toList (toBytesVector c) === Right bytes+          if textKind+            then do+              fmap VG.toList (toTextVector c) === Right texts+              -- Arrow.Record's utf8D reads the same rows, as slices of+              -- one copy of the column's referenced bytes.+              fmap V.toList (AR.runDecoder (AR.nullableD AR.utf8D) c) === Right texts+              if all isJust texts+                then fmap V.toList (AR.runDecoder AR.utf8D c) === Right (catMaybes texts)+                else assert (isLeft (AR.runDecoder AR.utf8D c))+            else assert (isLeft (toTextVector c))+  where+    int32s ch = maybe (Left ("not an int32 child: " ++ columnTag ch)) (Right . toMaybeVector) (asPrim PInt32 ch)+++prop_denseBinary :: Property+prop_denseBinary = withTests 300 . property $ do+  rows <- forAll (Gen.list (Range.linear 0 200) (Gen.bytes (Range.linear 0 100)))+  padding <- forAll (Gen.bytes (Range.linear 1 16))+  let sliced bytes = BS.take (BS.length bytes) (BS.drop (BS.length padding) (padding <> bytes <> padding))+      checkRows values = do+        let column = fromByteStrings (V.fromList (sliced <$> values))+        fmap VG.toList (toBytesVector column) === Right (Just <$> values)+        nullCount column === 0+  checkRows rows+  checkRows ((\n -> BS.replicate n (fromIntegral n)) <$> [0 .. 64])+++prop_denseBinaryOffsetLimit :: Property+prop_denseBinaryOffsetLimit = withTests 1 . property $ do+  -- The same 1 MiB allocation is shared by every input row: exercise a 2 GiB+  -- total without allocating a 2 GiB input or output.+  let bytes = BS.replicate (1024 * 1024) 0+      rows = V.replicate 2048 bytes+      nullable values = V.singleton Nothing <> (Just <$> values) <> V.singleton Nothing+      checkLimit column = do+        result <- evalIO (try (evaluate (columnLength column)) :: IO (Either ErrorCall Int))+        case result of+          Left err -> assert ("offset range" `isInfixOf` displayException err)+          Right _ -> failure+  checkLimit (fromByteStrings rows)+  checkLimit (fromMaybeByteStrings (nullable rows))+  checkLimit (fromMaybeTexts (nullable (V.replicate 2048 (TE.decodeUtf8 bytes))))+++prop_builders :: Property+prop_builders = withTests 300 . property $ do+  ints <- forAll (Gen.list (Range.linear 0 100) (maybeOf (Gen.int32 Range.linearBounded)))+  ints2 <- forAll (Gen.list (Range.linear 0 20) (maybeOf (Gen.int32 Range.linearBounded)))+  let (c1, c2) = runST $ do+        b <- newPrimBuilder PInt32 3+        mapM_ (appendPrimMaybe b) ints+        x <- freezeBuilder b+        mapM_ (appendPrimMaybe b) ints2+        y <- freezeBuilder b+        pure (x, y)+  c1 === fromMaybes PInt32 (V.fromList ints)+  c2 === fromMaybes PInt32 (V.fromList ints2)+  texts <- forAll (Gen.list (Range.linear 0 60) (maybeOf genText))+  let tc = runST $ do+        b <- newUtf8Builder 1+        mapM_ (appendTextMaybe b) texts+        freezeBuilder b+  tc === fromMaybeTexts (V.fromList texts)+  bytes <- forAll (Gen.list (Range.linear 0 60) (maybeOf genBytes))+  let bc = runST $ do+        b <- newLargeBinaryBuilder 0+        mapM_ (appendBytesMaybe b) bytes+        freezeBuilder b+  bc === fromMaybeLargeByteStrings (V.fromList bytes)+  bools <- forAll (Gen.list (Range.linear 0 100) (maybeOf Gen.bool))+  let boc = runST $ do+        b <- newBoolBuilder 0+        mapM_ (appendBoolMaybe b) bools+        freezeBuilder b+  boc === fromMaybeBools (V.fromList bools)+  nullCount boc === length (filter isNothing bools)+++prop_eq :: Property+prop_eq = withTests 400 . property $ do+  ty <- forAll (genTy 2)+  n <- forAll (Gen.int (Range.linear 0 8))+  (m1, c1) <- forAll (genColumnOf ty n)+  (m2, c2) <- forAll (genColumnOf ty n)+  assert (c1 == c1)+  (c1 == c2) === (m1 == m2)+  k <- forAll (Gen.int (Range.linear 0 n))+  rejoined <- evalEither (concatColumnArray (sliceColumnArray 0 k c1) (sliceColumnArray k (n - k) c1))+  assert (rejoined == c1)+++prop_slice :: Property+prop_slice = withTests 500 . property $ do+  (_, m, c) <- forAll genColumn+  let n = length m+  s <- forAll (Gen.int (Range.linear (-2) (n + 2)))+  l <- forAll (Gen.int (Range.linear (-2) (n + 2)))+  let s' = min n (max 0 s)+      l' = max 0 (min l (n - s'))+      sl = sliceColumnArray s l c+  columnLength sl === l'+  observe sl === take l' (drop s' m)+  columnTag sl === columnTag c+  -- The IPC layout of a sliced run-end column reads the same rows.+  rebased <- evalEither (rebaseRunEnds sl)+  observe rebased === observe sl+  case rebased of+    ColRunEndEncoded off len re _ -> do+      off === 0+      lastOr 0 (runEnds re) === len+    _ -> success+  where+    lastOr d xs = if null xs then d else last xs+++prop_sliceOfSlice :: Property+prop_sliceOfSlice = withTests 300 . property $ do+  (_, m, c) <- forAll genColumn+  let n = length m+  s1 <- forAll (Gen.int (Range.linear 0 n))+  l1 <- forAll (Gen.int (Range.linear 0 (n - s1)))+  s2 <- forAll (Gen.int (Range.linear 0 l1))+  l2 <- forAll (Gen.int (Range.linear 0 (l1 - s2)))+  let a = sliceColumnArray s2 l2 (sliceColumnArray s1 l1 c)+      b = sliceColumnArray (s1 + s2) l2 c+  observe a === observe b+  assert (a == b)+++prop_concatAdjacent :: Property+prop_concatAdjacent = withTests 500 . property $ do+  (_, m, c) <- forAll genColumn+  let n = length m+  ks <- forAll (fmap (sortOn id) (Gen.list (Range.linear 0 4) (Gen.int (Range.linear 0 n))))+  let cuts = zip (0 : ks) (ks ++ [n])+  joined <- evalEither (concatColumnArrays (map (\(a, b) -> sliceColumnArray a (b - a) c) cuts))+  columnLength joined === n+  observe joined === m+  assert (joined == c)+++prop_concatAppends :: Property+prop_concatAppends = withTests 400 . property $ do+  ty <- forAll (genTy 2)+  ns <- forAll (Gen.list (Range.linear 1 4) (Gen.int (Range.linear 0 8)))+  parts <- forAll (traverse (genColumnOf ty) ns)+  joined <- evalEither (concatColumnArrays (map snd parts))+  columnLength joined === sum ns+  observe joined === concatMap fst parts+  -- A type mismatch is rejected.+  (_, other) <- forAll (genColumnOf TNull 1)+  case ty of+    TNull -> success+    _ -> case parts of+      (p : _) -> assert (isLeft (concatColumnArray (snd p) other))+      [] -> success+++prop_take :: Property+prop_take = withTests 400 . property $ do+  (_, m, c) <- forAll genColumn+  let n = length m+  ix <- forAll (if n == 0 then pure [] else Gen.list (Range.linear 0 15) (Gen.int (Range.linear 0 (n - 1))))+  taken <- evalEither (takeColumnArray (VS.fromList ix) c)+  observe taken === map (m !!) ix+  columnTag taken === columnTag c+  assert (isLeft (takeColumnArray (VS.singleton n) c))+  assert (isLeft (takeColumnArray (VS.singleton (-1)) c))+++prop_copyColumn :: Property+prop_copyColumn = withTests 300 . property $ do+  (_, m, c) <- forAll genColumn+  let cp = copyColumn c+  observe cp === m+  assert (cp == c)+  let regions = filter (\b -> BS.length b > 0) . columnBuffers+      overlaps x y =+        let (BSI.BS fx lx) = x+            (BSI.BS fy ly) = y+            d = unsafeForeignPtrToPtr fy `minusPtr` unsafeForeignPtrToPtr fx+        in d < lx && negate d < ly+  assert (not (any (\a -> any (overlaps a) (regions c)) (regions cp)))+++prop_expandDictionary :: Property+prop_expandDictionary = withTests 300 . property $ do+  kt <- forAll (Gen.element keyTypes)+  vt <- forAll (genTy 0)+  n <- forAll (Gen.int (Range.linear 0 10))+  (m, c) <- forAll (genColumnOf (TDict kt vt) n)+  e <- evalEither (expandDictionary c)+  columnLength e === n+  observe e === m+  assert (columnTag e /= "ColDictionary")+++prop_maskValidity :: Property+prop_maskValidity = withTests 300 . property $ do+  (ty, m, c) <- forAll genColumn+  mask <- forAll (Gen.list (Range.singleton (length m)) Gen.bool)+  case maskValidity (validityOf mask) c of+    Left _ -> assert (not (hasValiditySlot c) && not (and mask) && not (isNullTy ty))+    Right masked -> observe masked === zipWith (\ok v -> if ok then v else VNull) mask m+  where+    isNullTy = \case+      TNull -> True+      _ -> False+++prop_mapKeys :: Property+prop_mapKeys = withTests 300 . property $ do+  n <- forAll (Gen.int (Range.linear 0 6))+  (m, c) <- forAll (genColumnOf (TMap TBool) n)+  let entrySorted = \case+        VMap kvs ->+          let ks = map fst kvs+          in notElem VNull ks && and (zipWith (\x y -> keyInt x <= keyInt y) ks (drop 1 ks))+        _ -> True+      keyInt = \case+        VInt x -> x+        _ -> 0+  isRight (validateMapKeysSorted c) === all entrySorted m+++prop_mkRejects :: Property+prop_mkRejects = withTests 300 . property $ do+  texts <- forAll (Gen.list (Range.linear 1 10) genText)+  let dat = TE.encodeUtf8 (T.concat texts)+      offs = VS.fromList (map fromIntegral (scanl (+) 0 (map (BS.length . TE.encodeUtf8) texts))) :: VS.Vector Int32+      n = length texts+  good <- evalEither (mkUtf8 Nothing offs dat)+  good === fromTexts (V.fromList texts)+  -- Offsets past the data, decreasing, or starting below zero.+  assert (isLeft (mkUtf8 Nothing (VS.snoc offs (fromIntegral (BS.length dat) + 1)) dat))+  assert (isLeft (mkBinary Nothing (VS.fromList [0, 2, 1]) "abc"))+  assert (isLeft (mkBinary Nothing (VS.fromList [-1, 0]) "abc"))+  assert (isLeft (mkBinary Nothing VS.empty "abc"))+  -- Invalid UTF-8, and an offset inside a character.+  assert (isLeft (mkUtf8 Nothing (VS.fromList [0, 2]) (BS.pack [0xC0, 0x80])))+  assert (isLeft (mkUtf8 Nothing (VS.fromList [0, 3]) (BS.pack [0xED, 0xA0, 0x80])))+  assert (isLeft (mkUtf8 Nothing (VS.fromList [0, 1, 2]) (BS.pack [0xC3, 0xA9])))+  assert (isRight (mkUtf8 Nothing (VS.fromList [0, 2]) (BS.pack [0xC3, 0xA9])))+  -- Validity of the wrong length.+  assert (isLeft (mkUtf8 (validityOf (replicate (n + 1) False)) offs dat))+  -- Lists, fixed-size lists, structs, unions, dictionaries, run ends.+  let child = fromTexts (V.fromList texts)+  assert (isLeft (mkList Nothing (VS.fromList [0, fromIntegral n + 1]) child))+  assert (isRight (mkList Nothing (VS.fromList [0, fromIntegral n]) child))+  assert (isLeft (mkFixedSizeList 2 n Nothing child))+  assert (isLeft (mkStruct (n + 1) Nothing (V.singleton ("a", child))))+  assert (isLeft (mkDenseUnion (VS.fromList [0]) (VS.fromList [fromIntegral n]) (V.singleton child)))+  assert (isLeft (mkDenseUnion (VS.fromList [1]) (VS.fromList [0]) (V.singleton child)))+  assert (isLeft (mkSparseUnion (VS.fromList (replicate (n + 1) 0)) (V.singleton child)))+  let keys ks = fromMaybes PInt16 (V.fromList ks)+  assert (isLeft (mkDictionary 1 (keys [Just (fromIntegral n)]) child))+  assert (isLeft (mkDictionary 1 (keys [Just (-1)]) child))+  assert (isRight (mkDictionary 1 (keys [Just 0, Nothing]) child))+  assert (isLeft (mkDictionary 1 child child))+  assert (isLeft (mkRunEndEncoded (primColumn PInt32 (VS.fromList [2, 2])) child))+  assert (isLeft (mkRunEndEncoded (primColumn PInt32 (VS.fromList [0])) child))+  assert (isLeft (mkRunEndEncoded (primColumn PDouble (VS.fromList [1])) child))+  -- Views: a long view pointing past its buffer.+  let longText = T.replicate 4 "abcd"+      ok = fromMaybeUtf8View (V.singleton (Just longText))+  case ok of+    ColUtf8View v views bufs -> do+      assert (isRight (mkUtf8View v views bufs))+      assert (isLeft (mkUtf8View v views (V.map (BS.take 3) bufs)))+      assert (isLeft (mkUtf8View v views V.empty))+      assert (isLeft (mkUtf8View v (BS.take 15 views) bufs))+    _ -> failure
+ test/Test/Arrow/Decimal.hs view
@@ -0,0 +1,122 @@+{-# LANGUAGE OverloadedStrings #-}++module Test.Arrow.Decimal (tests) where++import Arrow.Column qualified as C+import Arrow.File qualified as F+import Arrow.FlatBufferIPC qualified as IPC+import Arrow.Read.Columns (decodeRecordBatch, defaultReadOptions)+import Arrow.Stream qualified as S+import Arrow.Types+import Arrow.Write.Columns (validateColumns)+import Control.Monad (forM_, unless)+import Data.ByteString qualified as BS+import Data.ByteString.Lazy qualified as BL+import Data.Either (isLeft)+import Data.Int (Int32)+import Data.Vector qualified as V+import Data.Vector.Storable qualified as VS+++data DecimalCase = DecimalCase+  { decimalType :: !ArrowType+  , decimalColumn :: !C.ColumnArray+  , decimalWidth :: !Int+  }+++cases :: [DecimalCase]+cases =+  [ DecimalCase (ADecimal32 9 2) (C.fromMaybes (C.PDecimal32 9 2) (V.fromList [Just (-999999999), Nothing, Just 0, Just 123, Just 999999999])) 4+  , DecimalCase (ADecimal64 18 2) (C.fromMaybes (C.PDecimal64 18 2) (V.fromList [Just (-999999999999999999), Nothing, Just 0, Just 123, Just 999999999999999999])) 8+  , DecimalCase (ADecimal 38 2) (C.fromMaybes (C.PDecimal128 38 2) (V.fromList (map (fmap C.decimal128FromInteger) [Just (negate (10 ^ (38 :: Int) - 1)), Nothing, Just 0, Just 123, Just (10 ^ (38 :: Int) - 1)]))) 16+  , DecimalCase (ADecimal256 76 2) (C.fromMaybes (C.PDecimal256 76 2) (V.fromList (map (fmap C.decimal256FromInteger) [Just (negate (10 ^ (76 :: Int) - 1)), Nothing, Just 0, Just 123, Just (10 ^ (76 :: Int) - 1)]))) 32+  ]+++load :: Either String a -> IO a+load = either fail pure+++check :: String -> Bool -> IO ()+check name condition = unless condition (fail ("Arrow decimal regression: " ++ name))+++tests :: IO ()+tests = do+  forM_ cases $ \sample -> do+    let column = decimalColumn sample+        field = defaultLeafField "d" True (decimalType sample)+        schema = defaultSchema (V.singleton field)+    roundtrip field column+    roundtrip field (C.sliceColumnArray 1 3 column)+    structure <- load (C.mkStruct 5 Nothing (V.singleton ("d", column)))+    roundtrip (defaultField "s" True AStruct (V.singleton field)) structure+    list <- load (C.mkList (C.validityFromBools (V.fromList [True, False, True])) (VS.fromList [0, 2, 2, 5]) column)+    roundtrip (defaultField "l" True AList (V.singleton field)) list+    dictionary <- load (C.mkDictionary 42 (C.fromMaybes C.PInt8 (V.fromList [Just 0, Nothing, Just 4, Just 1])) column)+    let dictField = field {fieldDictionary = Just (DictionaryEncoding 42 (AInt 8 True) False)}+    encodedDictionary <- load (S.encodeArrowStream S.defaultWriteOptions (defaultSchema (V.singleton dictField)) [V.singleton dictionary])+    (_, dictBatches) <- load (S.decodeArrowStream defaultReadOptions encodedDictionary)+    expanded <- load (C.expandDictionary dictionary)+    actualExpanded <- load (C.expandDictionary (V.head (head dictBatches)))+    check "dictionary values retain decimal width/parameters/nulls" (actualExpanded == expanded)+    bytes <- load (S.encodeArrowStream S.defaultWriteOptions schema [V.singleton column])+    stream <- load (F.readArrowStream bytes)+    let (batch, body) = V.head (F.asBatches stream)+        Buffer off size = rbBuffers batch VS.! 1+        width = decimalWidth sample+        begin = fromIntegral off+        sizeInt = fromIntegral size+        swapped = BS.concat (map (\i -> BS.reverse (BS.take width (BS.drop (begin + i * width) body))) [0 .. C.columnLength column - 1])+        bigBody = BS.take begin body <> swapped <> BS.drop (begin + sizeInt) body+    check "wire buffer uses its declared physical width" (sizeInt == C.columnLength column * width)+    big <- load (decodeRecordBatch defaultReadOptions schema {arrowEndianness = Big} batch bigBody)+    check "big-endian signed decimals retain every bit" (V.head big == column)+    let short = batch {rbBuffers = rbBuffers batch VS.// [(1, Buffer off (size - 1))]}+    check "short data buffers reject using the actual decimal width" (isLeft (decodeRecordBatch defaultReadOptions schema short body))+  let widths = [(ADecimal32, 9), (ADecimal64, 18), (ADecimal, 38), (ADecimal256, 76)]+  forM_ widths $ \(make, maximumPrecision) -> do+    forM_ [-1, 0, maximumPrecision + 1, maxBound] $ \precision -> do+      let field = defaultLeafField "bad" True (make precision 0)+          schema = defaultSchema (V.singleton field)+          nested = defaultSchema (V.singleton (defaultField "parent" True AStruct (V.singleton field)))+      check "reader rejects invalid decimal precision before materialization" (isLeft (IPC.decodeSchemaMessage (IPC.buildSchemaMessage schema)))+      check "file/footer reader rejects invalid decimal precision" (isLeft (IPC.readFileIndex (IPC.writeArrowFileFBWithDicts schema [] [])))+      forM_ [schema, nested] $ \invalid -> do+        check "schema-only stream writer rejects invalid precision" (isLeft (S.encodeArrowStream S.defaultWriteOptions invalid []))+        check "schema-only lazy stream writer rejects invalid precision" (isLeft (S.encodeArrowStreamLazy S.defaultWriteOptions invalid []))+        check "schema-only file writer rejects invalid precision" (isLeft (S.encodeArrowFile S.defaultWriteOptions invalid []))+        check "schema-only lazy file writer rejects invalid precision" (isLeft (S.encodeArrowFileLazy S.defaultWriteOptions invalid []))+    forM_ [-2, maximumPrecision + 3, fromIntegral (minBound :: Int32), fromIntegral (maxBound :: Int32)] $ \scale -> do+      let schema = defaultSchema (V.singleton (defaultLeafField "valid" True (make maximumPrecision scale)))+      metadata <- load (IPC.decodeSchemaMessage (IPC.buildSchemaMessage schema))+      check "negative/large signed32 scales are preserved" (metadata == schema)+      _ <- load (S.encodeArrowStream S.defaultWriteOptions schema [])+      pure ()+    forM_ [toInteger (minBound :: Int32) - 1, toInteger (maxBound :: Int32) + 1, 4294967298] $ \scale -> do+      let schema = defaultSchema (V.singleton (defaultLeafField "bad" True (make maximumPrecision (fromInteger scale))))+      check "scale cannot silently wrap when serialized" (isLeft (S.encodeArrowStream S.defaultWriteOptions schema []) && isLeft (S.encodeArrowFile S.defaultWriteOptions schema []))+  check "physical integers cannot masquerade as decimal columns" (isLeft (validateColumns (V.singleton (defaultLeafField "d" False (ADecimal32 9 2))) (V.singleton (C.primColumn C.PInt32 (VS.singleton 1)))))+  forM_ [("pa_decimal32", ADecimal32 9 2, decimalColumn (head cases)), ("pa_decimal64", ADecimal64 18 2, decimalColumn (cases !! 1))] $ \(name, ty, expected) -> do+    forM_ [".arrows", ".arrow"] $ \extension -> do+      bytes <- BS.readFile ("test/golden/" ++ name ++ extension)+      (schema, batches) <- load ((if extension == ".arrow" then S.decodeArrowFile else S.decodeArrowStream) defaultReadOptions bytes)+      check "independent PyArrow golden retains schema and values" (fieldType (V.head (arrowFields schema)) == ty && V.head (head batches) == expected)+  putStrLn "Arrow decimal width/precision/scale/nesting/endianness regressions passed"+++roundtrip :: Field -> C.ColumnArray -> IO ()+roundtrip field column = do+  let schema = defaultSchema (V.singleton field)+      batch = [V.singleton column]+      encoders =+        [ S.encodeArrowStream S.defaultWriteOptions schema batch+        , BL.toStrict <$> S.encodeArrowStreamLazy S.defaultWriteOptions schema batch+        , S.encodeArrowFile S.defaultWriteOptions schema batch+        , BL.toStrict <$> S.encodeArrowFileLazy S.defaultWriteOptions schema batch+        ]+  forM_ (zip [False, False, True, True] encoders) $ \(file, encoded) -> do+    bytes <- load encoded+    (actualSchema, batches) <- load ((if file then S.decodeArrowFile else S.decodeArrowStream) defaultReadOptions bytes)+    check "schema and decimal values roundtrip without widening" (actualSchema == schema && length batches == 1 && V.head (head batches) == column)
+ test/Test/Arrow/Features.hs view
@@ -0,0 +1,117 @@+{-# LANGUAGE OverloadedStrings #-}++module Test.Arrow.Features (tests) where++import Arrow.Column qualified as C+import Arrow.FlatBufferIPC qualified as IPC+import Arrow.Stream qualified as S+import Arrow.Types+import Control.Monad (forM_, unless)+import Data.ByteString qualified as BS+import Data.ByteString.Builder qualified as BB+import Data.ByteString.Lazy qualified as BL+import Data.Either (isLeft)+import Data.Int (Int64)+import Data.Vector qualified as V+import Data.Vector.Storable qualified as VS+import FlatBuffers.Builder qualified as B+import FlatBuffers.Reader qualified as R+++load :: Either String a -> IO a+load = either fail pure+++check :: String -> Bool -> IO ()+check name condition = unless condition (fail ("Arrow feature regression: " ++ name))+++tests :: IO ()+tests = do+  let field = defaultLeafField "n" False (AInt 32 True)+      column = C.primColumn C.PInt32 (VS.fromList [-123, 0, 456])+      allFlags = [FeatureUnused, FeatureDictionaryReplacement, FeatureCompressedBody]+      combinations = map V.fromList [[], [FeatureUnused], [FeatureDictionaryReplacement], [FeatureCompressedBody], allFlags, reverse allFlags, [FeatureCompressedBody, FeatureDictionaryReplacement, FeatureCompressedBody]]+      options = S.defaultWriteOptions : map (\codec -> S.defaultWriteOptions {S.writeBodyCompression = Just codec}) (filter S.bodyCompressionAvailable [BodyZstd, LZ4Frame])+  forM_ combinations $ \flags -> do+    let schema = (defaultSchema (V.singleton field)) {arrowFeatures = flags}+        metadata = IPC.buildSchemaMessage schema+    decoded <- load (IPC.decodeSchemaMessage metadata)+    check "metadata retains flags, order and duplicates" (decoded == schema)+    codes <- featureCodes metadata+    check "wire values are signed64 enum constants in Schema slot3" (codes == V.map code flags)+    forM_ options $ \opts -> do+      let batch = [V.singleton column]+          encoders =+            [ S.encodeArrowStream opts schema batch+            , BL.toStrict <$> S.encodeArrowStreamLazy opts schema batch+            , S.encodeArrowFile opts schema batch+            , BL.toStrict <$> S.encodeArrowFileLazy opts schema batch+            ]+      forM_ (zip [False, False, True, True] encoders) $ \(file, encoded) -> do+        bytes <- load encoded+        (actualSchema, actualBatches) <- load ((if file then S.decodeArrowFile else S.decodeArrowStream) S.defaultReadOptions bytes)+        check "flags and values survive strict/lazy stream/file and compression" (actualSchema == schema && actualBatches == batch)+        if file+          then do+            index <- load (IPC.readFileIndex bytes)+            check "footer schema retains feature flags" (IPC.fiSchema index == schema)+            reader <- load (S.openFileReader S.defaultReadOptions bytes)+            check "file iterator schema retains feature flags" (S.fileReaderSchema reader == schema)+          else do+            reader <- load (S.openStreamReader S.defaultReadOptions bytes)+            check "stream iterator schema retains feature flags" (S.streamReaderSchema reader == schema)+  forM_ [([], V.empty), ([0, 1, 2], V.fromList allFlags), ([2, 1, 2], V.fromList [FeatureCompressedBody, FeatureDictionaryReplacement, FeatureCompressedBody])] $ \(codes, flags) -> do+    raw <- rawFeatureMessage codes+    schema <- load (IPC.decodeSchemaMessage raw)+    check "independent raw int64 vector is interpreted using feature wire values" (arrowFeatures schema == flags)+  forM_ [-1, 3, 4, 256, minBound, maxBound] $ \unknown -> do+    raw <- rawFeatureMessage [1, unknown, 2]+    check "unknown feature cannot silently lose semantics" (isLeft (IPC.decodeSchemaMessage raw))+  raw <- rawFeatureMessage [0, 1, 2]+  vec <- featureVector raw+  forM_ [0x7fffffff, 0xffffffff] $ \count -> do+    let malformed = BS.take vec raw <> BL.toStrict (BB.toLazyByteString (BB.word32LE count)) <> BS.drop (vec + 4) raw+    check "feature count is bounded before allocation" (isLeft (IPC.decodeSchemaMessage malformed))+  check "truncated feature payload is rejected" (isLeft (IPC.decodeSchemaMessage (BS.take (vec + 4 + 23) raw)))+  putStrLn "Arrow schema feature wire/round-trip/unknown/bounds regressions passed"+  where+    code :: Feature -> Int64+    code FeatureUnused = 0+    code FeatureDictionaryReplacement = 1+    code FeatureCompressedBody = 2+++-- Independent table construction keeps the test's expected wire values out+-- of the production feature encoder, including values it cannot represent.+rawFeatureMessage :: [Int64] -> IO BS.ByteString+rawFeatureMessage values = do+  b <- B.newBuilder+  features <- B.writeVectorInt64 b values+  schema <- B.writeTable b [Nothing, Nothing, Nothing, Just (B.voff features)]+  message <- B.writeTable b [Just (B.scalarI16 4), Just (B.scalarU8 1), Just (B.voff schema)]+  B.finish b message+++featureVector :: BS.ByteString -> IO Int+featureVector bytes = load $ do+  root <- fromIntegral <$> R.peekU32 bytes 0+  message <- R.resolveTable bytes root+  header <- maybe (Left "missing schema header") (R.followUOffset bytes) (message 2)+  schema <- R.resolveTable bytes header+  maybe (Left "missing feature vector") (R.followUOffset bytes) (schema 3)+++featureCodes :: BS.ByteString -> IO (V.Vector Int64)+featureCodes bytes = do+  -- Absence is the default empty vector; non-empty vectors must use slot3.+  root <- load (fromIntegral <$> R.peekU32 bytes 0)+  message <- load (R.resolveTable bytes root)+  header <- load (maybe (Left "missing schema header") (R.followUOffset bytes) (message 2))+  schema <- load (R.resolveTable bytes header)+  case schema 3 of+    Nothing -> pure V.empty+    Just p -> load $ do+      vec <- R.followUOffset bytes p+      (_, elems) <- R.readVectorOfStructs bytes vec 8+      V.mapM (R.peekI64 bytes) elems
+ test/Test/Arrow/Gen.hs view
@@ -0,0 +1,1109 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PatternSynonyms #-}++{- | Hedgehog generators for Arrow schemas with matching column+batches, plus a logical row model ('Value' / 'columnValues') that+compares columns by meaning rather than by physical layout (offsets,+dictionary encoding, run-end encoding, list views and float bit+patterns all reduce to the same row values).++Generated schemas cover every 'ArrowType' the codec supports, nullable+and not, nested to depth 3, with dictionary-encoded fields over every+index type whose value types may themselves be nested and contain+dictionary-encoded fields. Structs may have no fields and fixed-size+lists may have size 0.++Columns are built through the public construction API only: the+@from*@ conversions, the growable builders and the validating @mk*@+constructors (with arbitrary bytes in null slots and around the+referenced ranges). On top of that the generated physical layouts are+deliberately untidy, so the writer's rebase and realign paths run:++* columns are often a window of a longer column ('sliceColumnArray'),+  so bitmaps start at non-zero bit offsets, offsets at non-zero bases,+  run-end-encoded columns carry a logical offset and struct children+  are longer than the struct;+* var-length data carries junk (even invalid UTF-8) outside the+  referenced range, struct, fixed-size-list and sparse-union children+  carry extra rows, list offsets start past zero;+* with 'layoutDenormalised', a nullable column without nulls sometimes+  carries an all-valid validity bitmap with a zero null count.++Unions and run-end-encoded fields have no validity, map keys are+non-null, and a nullable dictionary column whose rows are all null may+reference an empty dictionary.+-}+module Test.Arrow.Gen (+  -- * Schemas and tables+  genField,+  genFieldType,+  genSchema,+  genColumnFor,+  genColumnForWith,+  genBatchFor,+  genTable,+  genOrderableKeyField,+  genDictionaryField,+  keyColumn,+  Layout (..),+  tidyLayout,+  untidyLayout,++  -- * Logical row model+  Value (..),+  columnValues,+  batchValues,+  compareKeyValues,+  hasDictionaries,+  childColumns,+) where++import Arrow.Column (+  ColumnArray,+  Decimal128 (..),+  Decimal256 (..),+  Float16 (..),+  IntervalDayTime (..),+  IntervalMonthDayNano (..),+  PrimType (..),+  SomePrimType (..),+  appendBoolMaybe,+  appendBytesMaybe,+  appendPrimMaybe,+  appendTextMaybe,+  bitAt,+  bitmapFromBools,+  bitmapGenerate,+  bitmapLength,+  decimal128ToInteger,+  decimal256ToInteger,+  freezeBuilder,+  fromBools,+  fromByteStrings,+  fromMaybeBinaryView,+  fromMaybeBools,+  fromMaybeByteStrings,+  fromMaybeFixedSizeBinary,+  fromMaybeLargeByteStrings,+  fromMaybeLargeTexts,+  fromMaybeTexts,+  fromMaybeUtf8View,+  fromMaybes,+  fromTexts,+  isValidAt,+  mkBinary,+  mkBool,+  mkDenseUnion,+  mkDictionary,+  mkFixedSizeBinary,+  mkFixedSizeList,+  mkLargeBinary,+  mkLargeList,+  mkLargeListView,+  mkLargeUtf8,+  mkList,+  mkListView,+  mkMap,+  mkPrim,+  mkRunEndEncoded,+  mkSparseUnion,+  mkStruct,+  mkUtf8,+  newBinaryBuilder,+  newBoolBuilder,+  newLargeBinaryBuilder,+  newLargeUtf8Builder,+  newPrimBuilder,+  newUtf8Builder,+  primColumn,+  primTypeFor,+  sliceColumnArray,+  takeColumnArray,+  validityFromBools,+  withPrim,+  pattern ColBinary,+  pattern ColBinaryView,+  pattern ColBool,+  pattern ColDenseUnion,+  pattern ColDictionary,+  pattern ColFixedSizeBinary,+  pattern ColFixedSizeList,+  pattern ColLargeBinary,+  pattern ColLargeList,+  pattern ColLargeListView,+  pattern ColLargeUtf8,+  pattern ColList,+  pattern ColListView,+  pattern ColMap,+  pattern ColNull,+  pattern ColPrim,+  pattern ColRunEndEncoded,+  pattern ColSparseUnion,+  pattern ColStruct,+  pattern ColUtf8,+  pattern ColUtf8View,+ )+import Arrow.Column.Internal qualified as I+import Arrow.Types+import Control.Monad (forM, replicateM)+import Control.Monad.ST (runST)+import Data.Bits (shiftL, shiftR, (.&.), (.|.))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.Int (Int16, Int32, Int64, Int8)+import Data.List (sortBy)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Text.Encoding qualified as TE+import Data.Vector qualified as V+import Data.Vector.Storable qualified as VS+import Data.Word (Word16, Word32, Word64, Word8)+import Foreign.Storable (Storable)+import GHC.Float (castDoubleToWord64, castFloatToWord32, castWord32ToFloat, castWord64ToDouble)+import Hedgehog (Gen)+import Hedgehog.Gen qualified as Gen+import Hedgehog.Range qualified as Range+++-- ============================================================+-- Schemas+-- ============================================================++{- | A schema of 1 to 4 uniquely named top-level fields, nested up to+depth 3, plus (in a third of the schemas) one field of a shape that+random nesting reaches rarely (see 'genEdgeField').+-}+genSchema :: Gen Schema+genSchema = do+  n <- Gen.int (Range.linear 1 4)+  fs <- forM [0 .. n - 1] $ \i -> genField 3 (T.pack ("c" ++ show i))+  edge <- Gen.frequency [(2, pure []), (1, (: []) <$> genEdgeField (T.pack ("c" ++ show n)))]+  pure (numberDictionaries (Schema (V.fromList (fs ++ edge)) Little V.empty V.empty))+++{- | A struct with no fields, a fixed-size list of size 0, a dictionary+whose struct or list values contain a dictionary field, or a nullable+dictionary (whose columns are sometimes all null over an empty+dictionary).+-}+genEdgeField :: Text -> Gen Field+genEdgeField name = do+  nullable <- Gen.bool+  Gen.choice+    [ pure (Field name nullable AStruct V.empty Nothing V.empty)+    , do+        c <- genField 1 "item"+        pure (Field name nullable (AFixedSizeList 0) (V.singleton c) Nothing V.empty)+    , do+        inner <- genDictionaryField 0 "e"+        container <- Gen.element [AStruct, AList, ALargeList]+        idx <- genIndexType+        pure (Field name nullable container (V.singleton inner) (Just (DictionaryEncoding 0 idx False)) V.empty)+    , (\f -> f {fieldNullable = True}) <$> genDictionaryField 0 name+    ]+++genIndexType :: Gen ArrowType+genIndexType = AInt <$> Gen.element [8, 16, 32, 64] <*> Gen.bool+++{- | Give every dictionary-encoded field a distinct dictionary id (in+pre-order), so each id names exactly one value type.+-}+numberDictionaries :: Schema -> Schema+numberDictionaries sch =+  let go !next f =+        let (next1, dict') = case fieldDictionary f of+              Just de -> (next + 1, Just de {deId = next})+              Nothing -> (next, Nothing)+            (next2, kids) = goList next1 (V.toList (fieldChildren f))+        in (next2, f {fieldDictionary = dict', fieldChildren = V.fromList kids})+      goList next = \case+        [] -> (next, [])+        (f : fs) ->+          let (n1, f') = go next f+              (n2, fs') = goList n1 fs+          in (n2, f' : fs')+      (_, fields) = goList 0 (V.toList (arrowFields sch))+  in sch {arrowFields = V.fromList fields}+++-- | One field named @name@; @depth@ bounds the remaining nesting.+genField :: Int -> Text -> Gen Field+genField depth name =+  Gen.frequency+    [ (8, genLeafField name)+    , (if depth > 0 then 6 else 0, genNestedField depth name)+    , (2, genDictionaryField depth name)+    ]+++leaf :: Text -> Bool -> ArrowType -> Field+leaf name nullable ty = Field name nullable ty V.empty Nothing V.empty+++genLeafField :: Text -> Gen Field+genLeafField name = do+  ty <- genFieldType+  nullable <- case ty of+    ANull -> pure True+    _ -> Gen.bool+  pure (leaf name nullable ty)+++-- | Any non-nested Arrow type.+genFieldType :: Gen ArrowType+genFieldType =+  Gen.choice+    [ pure ANull+    , AInt <$> Gen.element [8, 16, 32, 64] <*> Gen.bool+    , AFloatingPoint <$> Gen.element [Half, Single, DoublePrecision]+    , pure ABinary+    , pure AUtf8+    , pure ABool+    , ADecimal32 <$> Gen.int (Range.linear 1 9) <*> Gen.int (Range.linear 0 9)+    , ADecimal64 <$> Gen.int (Range.linear 1 18) <*> Gen.int (Range.linear 0 18)+    , ADecimal <$> Gen.int (Range.linear 1 38) <*> Gen.int (Range.linear 0 10)+    , ADecimal256 <$> Gen.int (Range.linear 1 76) <*> Gen.int (Range.linear 0 10)+    , ADate <$> Gen.element [DateDay, DateMillisecond]+    , Gen.element [ATime Second 32, ATime Millisecond 32, ATime Microsecond 64, ATime Nanosecond 64]+    , ATimestamp <$> Gen.element [Second, Millisecond, Microsecond, Nanosecond] <*> Gen.element [Nothing, Just "UTC"]+    , AInterval <$> Gen.element [YearMonth, DayTime, MonthDayNano]+    , AFixedSizeBinary <$> Gen.int (Range.linear 1 20)+    , ADuration <$> Gen.element [Second, Millisecond, Microsecond, Nanosecond]+    , pure ALargeBinary+    , pure ALargeUtf8+    , pure AUtf8View+    , pure ABinaryView+    ]+++genNestedField :: Int -> Text -> Gen Field+genNestedField depth name = do+  let child = genField (depth - 1)+  Gen.choice+    [ do+        n <- Gen.frequency [(1, pure 0), (3, Gen.int (Range.linear 1 3))]+        kids <- forM [0 .. n - 1] $ \i -> child (T.pack ("f" ++ show i))+        nullable <- Gen.bool+        pure (Field name nullable AStruct (V.fromList kids) Nothing V.empty)+    , listLike AList+    , listLike ALargeList+    , listLike AListView+    , listLike ALargeListView+    , do+        w <- Gen.frequency [(1, pure 0), (3, Gen.int (Range.linear 1 3))]+        listLike (AFixedSizeList w)+    , do+        sorted <- Gen.bool+        k <- if sorted then genOrderableKeyField "key" else (\f -> f {fieldNullable = False}) <$> genLeafKey+        v <- child "value"+        nullable <- Gen.bool+        let entries = Field "entries" False AStruct (V.fromList [k, v]) Nothing V.empty+        pure (Field name nullable (AMap sorted) (V.singleton entries) Nothing V.empty)+    , do+        mode <- Gen.element [Dense, Sparse]+        n <- Gen.int (Range.linear 1 3)+        kids <- forM [0 .. n - 1] $ \i -> child (T.pack ("u" ++ show i))+        ids <-+          Gen.choice+            [ pure V.empty+            , V.fromList . take n <$> Gen.shuffle [3, 7, 11, 0, 5]+            ]+        pure (Field name False (AUnion mode ids) (V.fromList kids) Nothing V.empty)+    , do+        reTy <- Gen.element [AInt 16 True, AInt 32 True, AInt 64 True]+        v <- child "values"+        let v' = case fieldType v of+              ARunEndEncoded -> v {fieldType = AInt 32 True, fieldChildren = V.empty, fieldDictionary = Nothing}+              _ -> v+        pure (Field name False ARunEndEncoded (V.fromList [leaf "run_ends" False reTy, v']) Nothing V.empty)+    ]+  where+    listLike ty = do+      c <- genField (depth - 1) "item"+      nullable <- Gen.bool+      pure (Field name nullable ty (V.singleton c) Nothing V.empty)+    genLeafKey = do+      ty <- Gen.filter (/= ANull) genFieldType+      pure (leaf "key" False ty)+++{- | A non-null key field whose 'Value' order agrees with the map-key+order of 'Arrow.Column.validateMapKeysSorted', so sorted keys can be+generated by sorting values.+-}+genOrderableKeyField :: Text -> Gen Field+genOrderableKeyField name =+  leaf name False+    <$> Gen.element+      [ AInt 8 True+      , AInt 16 False+      , AInt 32 True+      , AInt 64 False+      , AUtf8+      , ALargeUtf8+      , ABinary+      , ABool+      , ADate DateDay+      , ATimestamp Microsecond Nothing+      , AFixedSizeBinary 3+      , AFloatingPoint DoublePrecision+      , ADecimal 10 2+      , AUtf8View+      ]+++{- | A dictionary-encoded field. The value type is flat or, while+@depth@ allows, nested (struct, list, map, union, ...), so dictionary+fields can sit inside dictionary values.+-}+genDictionaryField :: Int -> Text -> Gen Field+genDictionaryField depth name = do+  value <-+    Gen.frequency+      [ (1, leaf name False <$> Gen.filter (/= ANull) genFieldType)+      , (if depth > 0 then 1 else 0, genNestedField depth name)+      ]+  idx <- genIndexType+  nullable <- Gen.bool+  ordered <- Gen.bool+  pure value {fieldNullable = nullable, fieldDictionary = Just (DictionaryEncoding 0 idx ordered)}+++-- ============================================================+-- Columns+-- ============================================================++-- | How untidy generated physical layouts may be.+data Layout = Layout+  { layoutSlices :: !Bool+  -- ^ build longer columns and slice a window out of them+  , layoutDenormalised :: !Bool+  -- ^ sometimes give a nullable column without nulls an all-valid validity+  }+  deriving stock (Show, Eq)+++-- | Every column is built exactly as long as it needs to be.+tidyLayout :: Layout+tidyLayout = Layout {layoutSlices = False, layoutDenormalised = False}+++-- | Windows of longer columns and denormalised validity.+untidyLayout :: Layout+untidyLayout = Layout {layoutSlices = True, layoutDenormalised = True}+++{- | One batch (0 to 4 of them) of 0 to 8 rows per table; every+batch draws its own data, so dictionary columns carry different+dictionaries in different batches. Layouts are 'untidyLayout'.+-}+genTable :: Gen (Schema, [V.Vector ColumnArray])+genTable = do+  sch <- genSchema+  nBatches <- Gen.int (Range.linear 0 4)+  batches <- replicateM nBatches (genBatchFor sch)+  pure (sch, batches)+++-- | A batch for the schema with 0 to 8 rows ('untidyLayout').+genBatchFor :: Schema -> Gen (V.Vector ColumnArray)+genBatchFor sch = do+  rows <- Gen.frequency [(1, pure 0), (6, Gen.int (Range.linear 1 8))]+  V.mapM (\f -> genColumnForWith untidyLayout f rows) (arrowFields sch)+++{- | A column of exactly @n@ rows for the field, often a window of a+longer column (but never with denormalised validity, so every core+row operation's contract applies).+-}+genColumnFor :: Field -> Int -> Gen ColumnArray+genColumnFor = genColumnForWith untidyLayout {layoutDenormalised = False}+++-- | A column of exactly @n@ rows for the field under a layout policy.+genColumnForWith :: Layout -> Field -> Int -> Gen ColumnArray+genColumnForWith layout f n+  | layoutSlices layout = do+      pad <- Gen.frequency [(3, pure Nothing), (2, Just <$> ((,) <$> Gen.int (Range.linear 0 3) <*> Gen.int (Range.linear 0 3)))]+      case pad of+        Nothing -> buildColumn layout f n >>= denormalise+        Just (pre, post) -> do+          c <- buildColumn layout f (pre + n + post)+          denormalise (sliceColumnArray pre n c)+  | otherwise = buildColumn layout f n >>= denormalise+  where+    denormalise c+      | layoutDenormalised layout && fieldNullable f = Gen.frequency [(5, pure c), (1, pure (allValidValidity c))]+      | otherwise = pure c+++{- | Replace an absent validity of a flat column by an all-valid bitmap+with null count 0 (raw construction: the public API never builds this).+Other columns are returned unchanged.+-}+allValidValidity :: ColumnArray -> ColumnArray+allValidValidity c = case c of+  I.ColPrim t Nothing xs -> withPrim t (I.ColPrim t (allValid (VS.length xs)) xs)+  I.ColBool Nothing bits -> I.ColBool (allValid (bitmapLength bits)) bits+  I.ColUtf8 Nothing o d -> I.ColUtf8 (allValid (VS.length o - 1)) o d+  I.ColBinary Nothing o d -> I.ColBinary (allValid (VS.length o - 1)) o d+  I.ColLargeUtf8 Nothing o d -> I.ColLargeUtf8 (allValid (VS.length o - 1)) o d+  I.ColLargeBinary Nothing o d -> I.ColLargeBinary (allValid (VS.length o - 1)) o d+  I.ColFixedSizeBinary w n Nothing d -> I.ColFixedSizeBinary w n (allValid n) d+  I.ColList Nothing o x -> I.ColList (allValid (VS.length o - 1)) o x+  I.ColStruct n Nothing kids -> I.ColStruct n (allValid n) kids+  _ -> c+  where+    allValid k = Just (I.Validity (bitmapGenerate k (const True)) 0)+++-- | Build a column of exactly @n@ rows (before any windowing by the caller).+buildColumn :: Layout -> Field -> Int -> Gen ColumnArray+buildColumn layout f n = case fieldDictionary f of+  Just de -> do+    emptyDict <- if fieldNullable f then Gen.frequency [(3, pure False), (1, pure True)] else pure False+    k <- if emptyDict then pure 0 else Gen.int (Range.linear 1 5)+    valuesNullable <- if fieldType f == ANull then pure False else Gen.bool+    vals <- sub f {fieldDictionary = Nothing, fieldNullable = valuesNullable} k+    keys <-+      if emptyDict+        then pure (V.replicate n Nothing)+        else V.fromList <$> replicateM n (nullableRow (Gen.int (Range.linear 0 (k - 1))))+    built (mkDictionary (deId de) (keyColumn (deIndexType de) keys) vals)+  Nothing -> case fieldType f of+    ANull -> pure (ColNull n)+    AInt w _ | w `notElem` [8, 16, 32, 64] -> error ("Test.Arrow.Gen: unsupported int width " ++ show w)+    ABinary -> varBytes mkBinary (Just fromByteStrings) fromMaybeByteStrings binaryViaBuilder genBytes+    ALargeBinary -> varBytes mkLargeBinary Nothing fromMaybeLargeByteStrings largeBinaryViaBuilder genBytes+    AUtf8 -> varText mkUtf8 (Just fromTexts) fromMaybeTexts utf8ViaBuilder+    ALargeUtf8 -> varText mkLargeUtf8 Nothing fromMaybeLargeTexts largeUtf8ViaBuilder+    ABinaryView -> fromMaybeBinaryView . V.fromList <$> replicateM n (nullableRow genViewBytes)+    AUtf8View -> fromMaybeUtf8View . V.fromList <$> replicateM n (nullableRow genViewText)+    ABool -> do+      bools <- replicateM n Gen.bool+      mvalid <- genValidityBools+      method <- Gen.int (Range.linear 0 2)+      let rows = rowsOf mvalid bools+      case method of+        0 -> pure (maybe (fromBools (V.fromList bools)) (const (fromMaybeBools (V.fromList rows))) mvalid)+        1 -> pure $ runST $ do+          b <- newBoolBuilder n+          mapM_ (appendBoolMaybe b) rows+          freezeBuilder b+        _ -> built (mkBool (validityOf mvalid) (bitmapFromBools (V.fromList bools)))+    AFixedSizeBinary w -> do+      payloads <- replicateM n (genFixed w)+      mvalid <- genValidityBools+      useMk <- Gen.bool+      if useMk+        then do+          junk <- genBytes+          built (mkFixedSizeBinary w n (validityOf mvalid) (BS.concat payloads <> junk))+        else built (fromMaybeFixedSizeBinary w (V.fromList (rowsOf mvalid payloads)))+    AStruct -> do+      extra <- genExtraRows+      kids <- V.mapM (\c -> (,) (fieldName c) <$> sub c (n + extra)) (fieldChildren f)+      mvalid <- genValidityBools+      built (mkStruct n (validityOf mvalid) kids)+    AList -> do+      (offs, total) <- genOffsets+      c <- sub (onlyChild f) total+      mvalid <- genValidityBools+      built (mkList (validityOf mvalid) (VS.fromList (map fromIntegral offs)) c)+    ALargeList -> do+      (offs, total) <- genOffsets+      c <- sub (onlyChild f) total+      mvalid <- genValidityBools+      built (mkLargeList (validityOf mvalid) (VS.fromList (map fromIntegral offs)) c)+    AFixedSizeList w -> do+      extra <- genExtraRows+      c <- sub (onlyChild f) (n * w + extra)+      mvalid <- genValidityBools+      built (mkFixedSizeList w n (validityOf mvalid) c)+    AListView -> do+      (offs, sizes, c) <- genViews+      mvalid <- genValidityBools+      built (mkListView (validityOf mvalid) (VS.fromList (map fromIntegral offs)) (VS.fromList (map fromIntegral sizes)) c)+    ALargeListView -> do+      (offs, sizes, c) <- genViews+      mvalid <- genValidityBools+      built (mkLargeListView (validityOf mvalid) (VS.fromList (map fromIntegral offs)) (VS.fromList (map fromIntegral sizes)) c)+    AMap sorted -> do+      let entries = onlyChild f+          (kf, vf) = case V.toList (fieldChildren entries) of+            [a, b] -> (a, b)+            _ -> error "Test.Arrow.Gen: map entries must have key and value"+      start <- Gen.frequency [(3, pure 0), (1, Gen.int (Range.linear 1 2))]+      lens <- replicateM n (Gen.int (Range.linear 0 3))+      extra <- genExtraRows+      let offs = scanl (+) start lens+          total = last offs + extra+      keys0 <- sub kf total+      let keys = if sorted then sortKeysWithin offs keys0 else keys0+      vals <- sub vf total+      mvalid <- genValidityBools+      built (mkMap (validityOf mvalid) (VS.fromList (map fromIntegral offs)) keys vals)+    AUnion mode _ -> do+      let kids = fieldChildren f+          k = V.length kids+      tids <- replicateM n (Gen.int (Range.linear 0 (k - 1)))+      let types = VS.fromList (map fromIntegral tids)+      case mode of+        Sparse -> do+          extra <- genExtraRows+          cs <- V.mapM (`sub` (n + extra)) kids+          built (mkSparseUnion types cs)+        Dense -> do+          let counts = map (\c -> length (filter (== c) tids)) [0 .. k - 1]+          extras <- replicateM k (Gen.int (Range.linear 0 1))+          cs <- V.imapM (\i c -> sub c (counts !! i + extras !! i)) kids+          built (mkDenseUnion types (VS.fromList (denseOffsets k tids)) cs)+    ARunEndEncoded -> case V.toList (fieldChildren f) of+      [ref, vf] -> do+        runs <- genRuns n+        let ends = drop 1 (scanl (+) 0 runs)+        vals <- sub vf (length runs)+        let re = case fieldType ref of+              AInt 16 _ -> primColumn PInt16 (VS.fromList (map fromIntegral ends))+              AInt 64 _ -> primColumn PInt64 (VS.fromList (map fromIntegral ends))+              _ -> primColumn PInt32 (VS.fromList (map fromIntegral ends))+        built (mkRunEndEncoded re vals)+      _ -> error "Test.Arrow.Gen: run-end-encoded field needs two children"+    ty -> case primTypeFor ty of+      Just (SomePrimType t) -> primCol t+      Nothing -> error ("Test.Arrow.Gen: no generator for " ++ show ty)+  where+    nullable = fieldNullable f+    sub = genColumnForWith layout+    built = either (\e -> error ("Test.Arrow.Gen: generated column rejected: " ++ e)) pure++    -- Validity as bools (True = valid); nullable columns are sometimes null-free.+    genValidityBools :: Gen (Maybe [Bool])+    genValidityBools+      | not nullable = pure Nothing+      | otherwise =+          Gen.frequency+            [ (1, pure (Just (replicate n True)))+            , (9, Just <$> replicateM n (Gen.frequency [(1, pure False), (2, pure True)]))+            ]+    validityOf = maybe Nothing (validityFromBools . V.fromList)+    rowsOf :: Maybe [Bool] -> [a] -> [Maybe a]+    rowsOf mvalid xs = maybe (map Just xs) (zipWith (\x ok -> if ok then Just x else Nothing) xs) mvalid+    nullableRow :: Gen a -> Gen (Maybe a)+    nullableRow g+      | nullable = Gen.frequency [(1, pure Nothing), (2, Just <$> g)]+      | otherwise = Just <$> g++    -- Fixed width: fromMaybes, a builder, or mkPrim with arbitrary null slots.+    primCol :: PrimType a -> Gen ColumnArray+    primCol t = withPrim t $ do+      xs <- replicateM n (genPrimValue t)+      mvalid <- genValidityBools+      method <- Gen.int (Range.linear 0 2)+      let rows = rowsOf mvalid xs+      case method of+        0 -> pure (maybe (primColumn t (VS.fromList xs)) (const (fromMaybes t (V.fromList rows))) mvalid)+        1 -> pure $ runST $ do+          b <- newPrimBuilder t n+          mapM_ (appendPrimMaybe b) rows+          freezeBuilder b+        _ -> built (mkPrim t (validityOf mvalid) (VS.fromList xs))++    -- Var-length bytes: from*, a builder, or mk* over data with junk+    -- (arbitrary bytes) before and after the referenced range.+    varBytes ::+      (Num o, Storable o) =>+      (Maybe I.Validity -> VS.Vector o -> ByteString -> Either String ColumnArray) ->+      Maybe (V.Vector ByteString -> ColumnArray) ->+      (V.Vector (Maybe ByteString) -> ColumnArray) ->+      ([Maybe ByteString] -> ColumnArray) ->+      Gen ByteString ->+      Gen ColumnArray+    varBytes mk fromAll fromSome viaBuilder genPayload = do+      payloads <- replicateM n genPayload+      mvalid <- genValidityBools+      method <- Gen.int (Range.linear 0 2)+      let rows = rowsOf mvalid payloads+      case method of+        0 -> pure $ case (mvalid, fromAll) of+          (Nothing, Just g) -> g (V.fromList payloads)+          _ -> fromSome (V.fromList rows)+        1 -> pure (viaBuilder rows)+        _ -> do+          prefix <- genJunk+          suffix <- genJunk+          let offs = scanl (+) (BS.length prefix) (map BS.length payloads)+          built (mk (validityOf mvalid) (VS.fromList (map fromIntegral offs)) (BS.concat (prefix : payloads ++ [suffix])))+    varText ::+      (Num o, Storable o) =>+      (Maybe I.Validity -> VS.Vector o -> ByteString -> Either String ColumnArray) ->+      Maybe (V.Vector Text -> ColumnArray) ->+      (V.Vector (Maybe Text) -> ColumnArray) ->+      ([Maybe Text] -> ColumnArray) ->+      Gen ColumnArray+    varText mk fromAll fromSome viaBuilder = do+      texts <- replicateM n genText+      mvalid <- genValidityBools+      method <- Gen.int (Range.linear 0 2)+      let rows = rowsOf mvalid texts+      case method of+        0 -> pure $ case (mvalid, fromAll) of+          (Nothing, Just g) -> g (V.fromList texts)+          _ -> fromSome (V.fromList rows)+        1 -> pure (viaBuilder rows)+        _ -> do+          prefix <- genTextJunk+          suffix <- genTextJunk+          let payloads = map TE.encodeUtf8 texts+              offs = scanl (+) (BS.length prefix) (map BS.length payloads)+          built (mk (validityOf mvalid) (VS.fromList (map fromIntegral offs)) (BS.concat (prefix : payloads ++ [suffix])))+    genJunk = Gen.frequency [(2, pure BS.empty), (1, Gen.bytes (Range.linear 1 5))]+    -- Outside the referenced range of UTF-8 data: valid text, so the+    -- offsets next to it stay on character boundaries.+    genTextJunk = Gen.frequency [(2, pure BS.empty), (1, TE.encodeUtf8 <$> Gen.text (Range.linear 1 3) Gen.unicode)]+    genExtraRows = Gen.frequency [(3, pure 0), (1, Gen.int (Range.linear 1 2))]++    -- Offsets may start past zero (a slice of a larger child), and the+    -- child may extend past the last offset.+    genOffsets = do+      start <- Gen.int (Range.linear 0 2)+      lens <- replicateM n (Gen.int (Range.linear 0 3))+      tailExtra <- Gen.int (Range.linear 0 1)+      let offs = scanl (+) start lens+      pure (offs, last offs + tailExtra)+    genViews = do+      m <- Gen.int (Range.linear 0 6)+      rows <- replicateM n $ do+        o <- Gen.int (Range.linear 0 m)+        s <- Gen.int (Range.linear 0 (m - o))+        pure (o, s)+      c <- sub (onlyChild f) m+      pure (map fst rows, map snd rows, c)+++binaryViaBuilder, largeBinaryViaBuilder :: [Maybe ByteString] -> ColumnArray+binaryViaBuilder rows = runST $ do+  b <- newBinaryBuilder (length rows)+  mapM_ (appendBytesMaybe b) rows+  freezeBuilder b+largeBinaryViaBuilder rows = runST $ do+  b <- newLargeBinaryBuilder (length rows)+  mapM_ (appendBytesMaybe b) rows+  freezeBuilder b+++utf8ViaBuilder, largeUtf8ViaBuilder :: [Maybe Text] -> ColumnArray+utf8ViaBuilder rows = runST $ do+  b <- newUtf8Builder (length rows)+  mapM_ (appendTextMaybe b) rows+  freezeBuilder b+largeUtf8ViaBuilder rows = runST $ do+  b <- newLargeUtf8Builder (length rows)+  mapM_ (appendTextMaybe b) rows+  freezeBuilder b+++-- | Dictionary keys at the index type's wire width (null rows are 'Nothing').+keyColumn :: ArrowType -> V.Vector (Maybe Int) -> ColumnArray+keyColumn idx ks = case idx of+  AInt 8 True -> fromMaybes PInt8 (V.map (fmap fromIntegral) ks)+  AInt 16 True -> fromMaybes PInt16 (V.map (fmap fromIntegral) ks)+  AInt 32 True -> fromMaybes PInt32 (V.map (fmap fromIntegral) ks)+  AInt 64 True -> fromMaybes PInt64 (V.map (fmap fromIntegral) ks)+  AInt 8 False -> fromMaybes PUInt8 (V.map (fmap fromIntegral) ks)+  AInt 16 False -> fromMaybes PUInt16 (V.map (fmap fromIntegral) ks)+  AInt 32 False -> fromMaybes PUInt32 (V.map (fmap fromIntegral) ks)+  AInt 64 False -> fromMaybes PUInt64 (V.map (fmap fromIntegral) ks)+  other -> error ("Test.Arrow.Gen: unsupported dictionary index type " ++ show other)+++-- | Values of every fixed-width element type, edge values included.+genPrimValue :: PrimType a -> Gen a+genPrimValue t = case t of+  PInt8 -> genI8+  PInt16 -> genI16+  PInt32 -> genI32+  PInt64 -> genI64+  PUInt8 -> genW8+  PUInt16 -> genW16+  PUInt32 -> genW32+  PUInt64 -> genW64+  PFloat16 -> Float16 <$> genW16+  PFloat -> genFloat+  PDouble -> genDouble+  PDate32 -> genI32+  PDate64 -> genI64+  PTime32 -> genI32+  PTime64 -> genI64+  PTimestamp -> genI64+  PDuration -> genI64+  PIntervalYearMonth -> genI32+  PIntervalDayTime -> IntervalDayTime <$> genI32 <*> genI32+  PIntervalMonthDayNano -> IntervalMonthDayNano <$> genI32 <*> genI32 <*> genI64+  PDecimal32 _ _ -> genI32+  PDecimal64 _ _ -> genI64+  PDecimal128 _ _ -> Decimal128 <$> genW64 <*> genW64+  PDecimal256 _ _ -> Decimal256 <$> genW64 <*> genW64 <*> genW64 <*> genW64+++onlyChild :: Field -> Field+onlyChild f = case V.toList (fieldChildren f) of+  [c] -> c+  _ -> error ("Test.Arrow.Gen: field " ++ show (fieldName f) ++ " needs exactly one child")+++-- | Per-row offsets of a dense union: the i-th row choosing child c gets that child's next slot.+denseOffsets :: Int -> [Int] -> [Int32]+denseOffsets k = go (replicate k (0 :: Int))+  where+    go _ [] = []+    go counters (t : ts) =+      let o = counters !! t+          counters' = zipWith (\i c -> if i == t then c + 1 else c) [0 ..] counters+      in fromIntegral o : go counters' ts+++-- | Positive run lengths summing to @n@.+genRuns :: Int -> Gen [Int]+genRuns n+  | n <= 0 = pure []+  | otherwise = do+      r <- Gen.int (Range.linear 1 n)+      rest <- genRuns (n - r)+      pure (r : rest)+++{- | Reorder keys so each entry's keys are non-decreasing; rows before+the first offset and after the last one stay where they are.+-}+sortKeysWithin :: [Int] -> ColumnArray -> ColumnArray+sortKeysWithin offs keys =+  case columnValues keys of+    Left e -> error e+    Right vs ->+      let entries = zip offs (drop 1 offs)+          inEntries = concatMap (\(s, e) -> sortBy (\i j -> compareKeyValues (vs V.! i) (vs V.! j)) [s .. e - 1]) entries+          (firstOff, lastOff) = case offs of+            [] -> (0, 0)+            (o : _) -> (o, last offs)+          perm = [0 .. firstOff - 1] ++ inEntries ++ [lastOff .. V.length vs - 1]+      in either error id (takeColumnArray (VS.fromList perm) keys)+++-- ============================================================+-- Leaf value generators+-- ============================================================++genI8 :: Gen Int8+genI8 = Gen.frequency [(1, Gen.element [minBound, maxBound, 0, -1]), (4, Gen.int8 Range.linearBounded)]++genI16 :: Gen Int16+genI16 = Gen.frequency [(1, Gen.element [minBound, maxBound, 0, -1]), (4, Gen.int16 Range.linearBounded)]++genI32 :: Gen Int32+genI32 = Gen.frequency [(1, Gen.element [minBound, maxBound, 0, -1]), (4, Gen.int32 Range.linearBounded)]++genI64 :: Gen Int64+genI64 = Gen.frequency [(1, Gen.element [minBound, maxBound, 0, -1]), (4, Gen.int64 Range.linearBounded)]++genW8 :: Gen Word8+genW8 = Gen.frequency [(1, Gen.element [minBound, maxBound]), (4, Gen.word8 Range.linearBounded)]++genW16 :: Gen Word16+genW16 = Gen.frequency [(1, Gen.element [minBound, maxBound, 0x7e00, 0x7c00, 0xfc00, 0x8000]), (4, Gen.word16 Range.linearBounded)]++genW32 :: Gen Word32+genW32 = Gen.frequency [(1, Gen.element [minBound, maxBound]), (4, Gen.word32 Range.linearBounded)]++genW64 :: Gen Word64+genW64 = Gen.frequency [(1, Gen.element [minBound, maxBound]), (4, Gen.word64 Range.linearBounded)]++genFloat :: Gen Float+genFloat =+  Gen.frequency+    [ (1, Gen.element [0 / 0, 1 / 0, -1 / 0, -0.0, 0, 3.4028235e38, -1.1754944e-38])+    , (4, Gen.float (Range.linearFracFrom 0 (-1e10) 1e10))+    ]++genDouble :: Gen Double+genDouble =+  Gen.frequency+    [ (1, Gen.element [0 / 0, 1 / 0, -1 / 0, -0.0, 0, 1.7976931348623157e308, 5.0e-324])+    , (4, Gen.double (Range.linearFracFrom 0 (-1e100) 1e100))+    ]++genBytes :: Gen ByteString+genBytes = Gen.bytes (Range.linear 0 20)++genFixed :: Int -> Gen ByteString+genFixed w = Gen.bytes (Range.singleton w)++genText :: Gen Text+genText = Gen.frequency [(1, pure ""), (4, Gen.text (Range.linear 0 16) Gen.unicode)]++-- | View payloads straddle the 12-byte inline limit.+genViewBytes :: Gen ByteString+genViewBytes = Gen.choice [Gen.bytes (Range.linear 0 12), Gen.bytes (Range.linear 13 40)]++genViewText :: Gen Text+genViewText =+  Gen.choice+    [ Gen.text (Range.linear 0 3) Gen.unicode+    , Gen.text (Range.linear 0 12) Gen.alphaNum+    , Gen.text (Range.linear 5 30) Gen.unicode+    ]+++-- ============================================================+-- Logical row model+-- ============================================================++{- | The meaning of one row, independent of physical layout. Floats are+kept as bit patterns so NaN rows compare equal to themselves.+-}+data Value+  = VNull+  | VInt Integer+  | VF16 Word16+  | VF32 Word32+  | VF64 Word64+  | VBool Bool+  | VText Text+  | VBytes ByteString+  | VDecimal Integer+  | VPair Int32 Int32+  | VTriple Int32 Int32 Int64+  | VList [Value]+  | VStruct [(Text, Value)]+  | VMap [(Value, Value)]+  | VUnion Int8 Value+  deriving stock (Show, Eq)+++-- | Logical rows of every column of a batch.+batchValues :: V.Vector ColumnArray -> Either String [[Value]]+batchValues = traverse (fmap V.toList . columnValues) . V.toList+++-- | Direct children of a column, run ends and dictionary keys and values included.+childColumns :: ColumnArray -> [ColumnArray]+childColumns = \case+  ColStruct _ _ cs -> map snd (V.toList cs)+  ColList _ _ x -> [x]+  ColLargeList _ _ x -> [x]+  ColFixedSizeList _ _ _ x -> [x]+  ColMap _ _ k v -> [k, v]+  ColDenseUnion _ _ cs -> V.toList cs+  ColSparseUnion _ cs -> V.toList cs+  ColRunEndEncoded _ _ r v -> [r, v]+  ColListView _ _ _ x -> [x]+  ColLargeListView _ _ _ x -> [x]+  ColDictionary _ k v -> [k, v]+  _ -> []+++-- | Whether any column (at any depth) is dictionary-encoded.+hasDictionaries :: ColumnArray -> Bool+hasDictionaries = \case+  ColDictionary {} -> True+  c -> any hasDictionaries (childColumns c)+++{- | Logical rows of a column, read straight from its buffers (not+through the accessors under test); any out-of-range reference or+invalid UTF-8 inside a referenced range is a 'Left'.+-}+columnValues :: ColumnArray -> Either String (V.Vector Value)+columnValues col = case col of+  ColNull n -> Right (V.replicate n VNull)+  ColPrim t v xs -> withPrim t (Right (V.generate (VS.length xs) (\i -> if isValidAt v i then primValue t (xs VS.! i) else VNull)))+  ColBool v bits -> Right (V.generate (bitmapLength bits) (\i -> if isValidAt v i then VBool (bitAt bits i) else VNull))+  ColUtf8 v o d -> varRows utf8 v o d+  ColLargeUtf8 v o d -> varRows utf8 v o d+  ColBinary v o d -> varRows (Right . VBytes) v o d+  ColLargeBinary v o d -> varRows (Right . VBytes) v o d+  ColFixedSizeBinary w n v d ->+    V.generateM n $ \i ->+      if isValidAt v i then VBytes <$> cut "fixed-size binary row" d (i * w) w else Right VNull+  ColUtf8View v views bufs -> viewValues utf8 v views bufs+  ColBinaryView v views bufs -> viewValues (Right . VBytes) v views bufs+  ColStruct n v cs -> do+    kids <- V.mapM (traverse columnValues) cs+    V.generateM n $ \i ->+      if isValidAt v i+        then VStruct . V.toList <$> V.mapM (\(nm, vs) -> (,) nm <$> at "struct child row" vs i) kids+        else Right VNull+  ColList v o c -> listRows v (ints o) c+  ColLargeList v o c -> listRows v (ints o) c+  ColListView v o s c -> viewRows v (ints o) (ints s) c+  ColLargeListView v o s c -> viewRows v (ints o) (ints s) c+  ColFixedSizeList w n v c -> do+    vs <- columnValues c+    V.generateM n $ \i -> if isValidAt v i then VList <$> slice vs (i * w) w else Right VNull+  ColMap v o k x -> do+    ks <- columnValues k+    xs <- columnValues x+    let offs = ints o+    V.generateM (max 0 (VS.length o - 1)) $ \i -> do+      let s = offs V.! i+          e = offs V.! (i + 1)+      if isValidAt v i+        then do+          kk <- slice ks s (e - s)+          vv <- slice xs s (e - s)+          Right (VMap (zip kk vv))+        else Right VNull+  ColDenseUnion ts offs cs -> do+    kids <- V.mapM columnValues cs+    V.generateM (VS.length ts) $ \i -> do+      let t = ts VS.! i+      kid <- at "dense union child" kids (fromIntegral t)+      VUnion t <$> at "dense union offset" kid (fromIntegral (offs VS.! i))+  ColSparseUnion ts cs -> do+    kids <- V.mapM columnValues cs+    V.generateM (VS.length ts) $ \i -> do+      let t = ts VS.! i+      kid <- at "sparse union child" kids (fromIntegral t)+      VUnion t <$> at "sparse union row" kid i+  ColDictionary _ keys vals -> do+    ks <- columnValues keys+    vs <- columnValues vals+    V.mapM+      ( \case+          VNull -> Right VNull+          VInt k -> at "dictionary key" vs (fromIntegral k)+          other -> Left ("Test.Arrow.Gen: dictionary key " ++ show other)+      )+      ks+  ColRunEndEncoded off len re vals -> do+    ends <- columnValues re+    vs <- columnValues vals+    endInts <- traverse (\case VInt e -> Right (fromIntegral e :: Int); other -> Left ("Test.Arrow.Gen: run end " ++ show other)) ends+    V.generateM len $ \i -> case V.findIndex (> off + i) endInts of+      Nothing -> Left "Test.Arrow.Gen: row past the last run end"+      Just r -> at "run value" vs r+  where+    utf8 bs = either (const (Left "Test.Arrow.Gen: invalid UTF-8 in a referenced range")) (Right . VText) (TE.decodeUtf8' bs)+    ints :: (Storable o, Integral o) => VS.Vector o -> V.Vector Int+    ints = V.map fromIntegral . V.convert+    varRows :: (Storable o, Integral o) => (ByteString -> Either String Value) -> Maybe I.Validity -> VS.Vector o -> ByteString -> Either String (V.Vector Value)+    varRows decode v o d =+      let offs = ints o+      in V.generateM (max 0 (V.length offs - 1)) $ \i ->+           if isValidAt v i+             then cut "var-length row" d (offs V.! i) (offs V.! (i + 1) - offs V.! i) >>= decode+             else Right VNull+    viewValues decode v views bufs =+      V.generateM (BS.length views `quot` 16) $ \i ->+        if isValidAt v i+          then do+            let base = 16 * i+                len = le32 views base+            bytes <-+              if len <= 12+                then cut "inline view" views (base + 4) len+                else do+                  buf <- at "view buffer" bufs (le32 views (base + 8))+                  cut "view data" buf (le32 views (base + 12)) len+            decode bytes+          else Right VNull+    listRows v offs c = do+      vs <- columnValues c+      V.generateM (max 0 (V.length offs - 1)) $ \i ->+        if isValidAt v i+          then VList <$> slice vs (offs V.! i) (offs V.! (i + 1) - offs V.! i)+          else Right VNull+    viewRows v offs sizes c = do+      vs <- columnValues c+      V.generateM (V.length offs) $ \i ->+        if isValidAt v i+          then VList <$> slice vs (offs V.! i) (sizes V.! i)+          else Right VNull+    slice :: V.Vector Value -> Int -> Int -> Either String [Value]+    slice vs s l+      | s < 0 || l < 0 || s + l > V.length vs = Left ("Test.Arrow.Gen: slice out of range " ++ show s ++ "+" ++ show l ++ " of " ++ show (V.length vs))+      | otherwise = Right (V.toList (V.slice s l vs))+    cut :: String -> ByteString -> Int -> Int -> Either String ByteString+    cut what bs s l+      | s < 0 || l < 0 || s + l > BS.length bs = Left ("Test.Arrow.Gen: " ++ what ++ " out of range")+      | otherwise = Right (BS.take l (BS.drop s bs))+    at :: String -> V.Vector a -> Int -> Either String a+    at what vs i = maybe (Left ("Test.Arrow.Gen: " ++ what ++ " out of range")) Right (vs V.!? i)+++-- | Little-endian signed 32-bit integer at a byte offset (0 past the end).+le32 :: ByteString -> Int -> Int+le32 bs o+  | o < 0 || o + 4 > BS.length bs = 0+  | otherwise =+      let w = foldr (\k acc -> acc `shiftL` 8 .|. fromIntegral (BS.index bs (o + k))) (0 :: Word32) [0 .. 3]+      in fromIntegral (fromIntegral w :: Int32)+++-- | The logical value of one fixed-width element.+primValue :: PrimType a -> a -> Value+primValue t x = case t of+  PInt8 -> VInt (toInteger x)+  PInt16 -> VInt (toInteger x)+  PInt32 -> VInt (toInteger x)+  PInt64 -> VInt (toInteger x)+  PUInt8 -> VInt (toInteger x)+  PUInt16 -> VInt (toInteger x)+  PUInt32 -> VInt (toInteger x)+  PUInt64 -> VInt (toInteger x)+  PDate32 -> VInt (toInteger x)+  PDate64 -> VInt (toInteger x)+  PTime32 -> VInt (toInteger x)+  PTime64 -> VInt (toInteger x)+  PTimestamp -> VInt (toInteger x)+  PDuration -> VInt (toInteger x)+  PIntervalYearMonth -> VInt (toInteger x)+  PFloat16 -> let Float16 w = x in VF16 w+  PFloat -> VF32 (castFloatToWord32 x)+  PDouble -> VF64 (castDoubleToWord64 x)+  PIntervalDayTime -> let IntervalDayTime d m = x in VPair d m+  PIntervalMonthDayNano -> let IntervalMonthDayNano m d ns = x in VTriple m d ns+  PDecimal32 _ _ -> VDecimal (toInteger x)+  PDecimal64 _ _ -> VDecimal (toInteger x)+  PDecimal128 _ _ -> VDecimal (decimal128ToInteger x)+  PDecimal256 _ _ -> VDecimal (decimal256ToInteger x)+++{- | The map-key order of 'Arrow.Column.validateMapKeysSorted', on+model values: numbers numerically (floats with -0 == 0 and every NaN+equal and largest), text by code point, bytes unsigned-lexicographic,+booleans False first.+-}+compareKeyValues :: Value -> Value -> Ordering+compareKeyValues a b = case (a, b) of+  (VInt x, VInt y) -> compare x y+  (VDecimal x, VDecimal y) -> compare x y+  (VF16 x, VF16 y) -> floating (halfToDouble x) (halfToDouble y)+  (VF32 x, VF32 y) -> floating (realToFrac (castWord32ToFloat x)) (realToFrac (castWord32ToFloat y))+  (VF64 x, VF64 y) -> floating (castWord64ToDouble x) (castWord64ToDouble y)+  (VText x, VText y) -> compare (TE.encodeUtf8 x) (TE.encodeUtf8 y)+  (VBytes x, VBytes y) -> compare x y+  (VBool x, VBool y) -> compare x y+  _ -> EQ+  where+    floating :: Double -> Double -> Ordering+    floating x y = case (isNaN x, isNaN y) of+      (True, True) -> EQ+      (True, False) -> GT+      (False, True) -> LT+      (False, False) -> compare x y+++halfToDouble :: Word16 -> Double+halfToDouble w =+  let sign = if w .&. 0x8000 /= 0 then -1 else 1+      ex = fromIntegral ((w `shiftR` 10) .&. 0x1f) :: Int+      mant = fromIntegral (w .&. 0x3ff) :: Double+  in case ex of+       0 -> sign * mant * 2 ** (-24)+       31 -> if mant == 0 then sign * (1 / 0) else 0 / 0+       _ -> sign * (1 + mant / 1024) * 2 ^^ (ex - 15)
+ test/Test/Arrow/Malformed.hs view
@@ -0,0 +1,1202 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PatternSynonyms #-}++{- | Malformed-input robustness for the Arrow IPC decoders.++The corpus is a set of valid encodings (wireform's stream and file+writers over flat, nullable, nested, dictionary, view, run-end-encoded+and list-view columns, plus the pyarrow goldens). Each property mutates+a corpus entry, or generates raw bytes, and feeds the result to every+public decode entry point. A decoder must return 'Left' or a 'Right'+that can be forced to normal form; it must never throw, hang, or+allocate past 'allocationLimitBytes'.++Decoded columns alias the input and their 'NFData' instance forces+only the outer constructor, so "forced to normal form" means more for+a column-producing decoder: every accessor is run over the column's+rows in bounded windows (children, dictionary values and run ends+included), the column must read back through the row model of+"Test.Arrow.Gen" without an out-of-range reference, and it must equal+itself and its 'copyColumn'. Windows matter because an all-null,+fieldless or run-end-encoded column can name billions of rows for a+few hundred input bytes; walking all of them would measure the test,+not the decoder.+-}+module Test.Arrow.Malformed (tests) where++import Arrow.Column (+  ChildRange (..),+  ColumnArray,+  PrimArray (..),+  PrimType (..),+  anyBytesAt,+  anyTextAt,+  boolAt,+  columnLength,+  copyColumn,+  dictKeyAt,+  expandDictionary,+  fromBools,+  fromByteStrings,+  fromMaybeBinaryView,+  fromMaybeBools,+  fromMaybeFixedSizeBinary,+  fromMaybeLargeTexts,+  fromMaybeTexts,+  fromMaybeUtf8View,+  fromMaybes,+  fromTexts,+  isValidAt,+  listRange,+  mkDenseUnion,+  mkDictionary,+  mkFixedSizeList,+  mkLargeList,+  mkLargeListView,+  mkList,+  mkListView,+  mkMap,+  mkRunEndEncoded,+  mkSparseUnion,+  mkStruct,+  nullCount,+  primAt,+  primColumn,+  sliceColumnArray,+  toBoolVector,+  toBytesVector,+  toTextVector,+  validity,+  validityFromBools,+  withPrim,+  pattern ColNull,+  pattern ColPrim,+ )+import Arrow.Column.Internal qualified as I+import Arrow.File qualified as File+import Arrow.Read.Columns (ReadOptions (..), defaultReadOptions)+import Arrow.FlatBufferIPC (+  ArrowBlock (..),+  DictBatch (..),+  FileIndex (..),+  StreamFrame (..),+  buildDictionaryBatchMessage,+  buildRecordBatchMessage,+  bodyCompressionAvailable,+  buildSchemaMessage,+  compressBufferEither,+  decodeSparseTensorFrame,+  decodeTensorFrame,+  encapsulateMessage,+  readArrowFileFBWithDicts,+  readFileIndex,+  readArrowStreamFBInterleaved,+  writeArrowFileFBWithDicts,+ )+import Arrow.IPC (decodeIPCMessage, validateRecordBatchBuffers)+import Arrow.Stream (+  decodeArrowFile,+  decodeArrowStream,+  WriteOptions (..),+  defaultWriteOptions,+  encodeArrowFile,+  encodeArrowStream,+  fileReaderToList,+  openFileReader,+  openStreamReader,+  streamReaderSchema,+  streamReaderToList,+ )+import Arrow.Types+import Arrow.Vector qualified as AV+import Arrow.Write qualified as Write+import Control.DeepSeq (NFData, force)+import Control.Exception (SomeException, displayException, evaluate, try)+import Data.Bits (complementBit, shiftR)+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.List (foldl')+import Data.Maybe (fromMaybe, mapMaybe)+import Data.Int (Int32, Int64)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Vector qualified as V+import Data.Vector.Generic qualified as VG+import Data.Vector.Storable qualified as VS+import Data.ByteString.Unsafe qualified as BSU+import Foreign.Marshal.Utils (copyBytes)+import Foreign.Ptr (castPtr)+import Data.Word (Word32, Word64, Word8)+import FlatBuffers.Builder qualified as FB+import Hedgehog hiding (Seed)+import Hedgehog.Gen qualified as Gen+import Hedgehog.Range qualified as Range+import Hedgehog.Internal.Property (CoverPercentage)+import System.Mem (disableAllocationLimit, enableAllocationLimit, setAllocationCounter)+import System.Timeout (timeout)+import Test.Arrow.Gen qualified as G+++tests :: IO Bool+tests = do+  goldens <- traverse loadGolden ["pa_int32.arrows", "pa_mixed.arrows", "pa_dict.arrows"]+  let !corpus = builtCorpus ++ goldens+  checkParallel $+    Group+      "Arrow.Malformed"+      [ ("unmutated corpus decodes", prop_corpusDecodes corpus)+      , ("truncation at any point", prop_truncation corpus)+      , ("random bit flips", prop_bitFlips corpus)+      , ("random byte overwrites", prop_overwrites corpus)+      , ("spliced, deleted and duplicated ranges", prop_splices corpus)+      , ("hostile buffer descriptors and lengths", prop_descriptors corpus)+      , ("schema/type confusion over valid bodies", prop_schemaConfusion corpus)+      , ("pure random bytes", prop_randomBytes)+      , ("random bytes behind valid framing", prop_framedRandom)+      , ("aliased FlatBuffers offsets (DAG bomb)", prop_dagBomb)+      , ("schema nested past the depth limit", prop_deepSchema)+      , ("unbacked rows from a tiny header", prop_unbackedRows)+      , ("buffer descriptor validation bounds", prop_bufferValidator)+      , ("compressed buffer length claims", prop_decompressionClaims)+      , ("decode budgets bound rows, values, unbacked rows and decompression", prop_readerBudgets)+      , ("mutated writer output over zero-copy decode never throws from any accessor", prop_zeroCopyAccessorFuzz)+      ]+++-- * Corpus++data Encoding = StreamEnc | FileEnc+  deriving stock (Show, Eq)+++data Seed = Seed+  { seedName :: String+  , seedEncoding :: Encoding+  , seedBytes :: ByteString+  }+++instance Show Seed where+  show s = seedName s ++ " (" ++ show (seedEncoding s) ++ ", " ++ show (BS.length (seedBytes s)) ++ " bytes)"+++loadGolden :: FilePath -> IO Seed+loadGolden name = Seed ("golden " ++ name) StreamEnc <$> BS.readFile ("test/golden/" ++ name)+++plain :: Text -> Bool -> ArrowType -> Field+plain nm nullable ty = Field nm nullable ty V.empty Nothing V.empty+++nested :: Text -> Bool -> ArrowType -> [Field] -> Field+nested nm nullable ty cs = Field nm nullable ty (V.fromList cs) Nothing V.empty+++schemaOf :: [Field] -> Schema+schemaOf fs = Schema (V.fromList fs) Little V.empty V.empty+++-- | Every shape is encoded as a stream and as a file by the high-level+-- writer, uncompressed and with each body-compression codec; the flat+-- shapes also go through 'Arrow.Write'. The shapes match their+-- schemas, so an encoder 'Left' is a bug in the corpus and aborts.+builtCorpus :: [Seed]+builtCorpus = concatMap encodeShape shapes ++ concatMap viaWrite (take 2 shapes)+  where+    encodeShape (nm, sch, batches) = concatMap (encodeWith nm sch batches) codecs+    encodeWith nm sch batches (suffix, codec) =+      let opts = defaultWriteOptions {writeBodyCompression = codec}+      in [ Seed (nm ++ suffix) StreamEnc (built nm (encodeArrowStream opts sch batches))+         , Seed (nm ++ suffix) FileEnc (built nm (encodeArrowFile opts sch batches))+         ]+    codecs = ("", Nothing) : map (\c -> (codecSuffix c, Just c)) availableCodecs+    codecSuffix BodyZstd = " (zstd)"+    codecSuffix LZ4Frame = " (lz4)"+    viaWrite (nm, sch, batches) =+      [ Seed (nm ++ " via Arrow.Write") StreamEnc (built nm (Write.writeArrowStream sch (V.fromList batches)))+      , Seed (nm ++ " via Arrow.Write") FileEnc (built nm (Write.writeArrowFile sch (V.fromList batches)))+      ]+    built nm = either (\e -> error ("Test.Arrow.Malformed: shape " ++ nm ++ " does not encode: " ++ e)) id+++shapes :: [(String, Schema, [V.Vector ColumnArray])]+shapes =+  [ ( "flat"+    , schemaOf+        [ plain "i32" False (AInt 32 True)+        , plain "i64" False (AInt 64 True)+        , plain "u16" False (AInt 16 False)+        , plain "f64" False (AFloatingPoint DoublePrecision)+        , plain "s" False AUtf8+        , plain "b" False ABinary+        , plain "flag" False ABool+        , plain "day" False (ADate DateDay)+        ]+    , [ V.fromList+          [ prim PInt32 [1, -2, 3]+          , prim PInt64 [10, 20, 30]+          , prim PUInt16 [7, 8, 9]+          , prim PDouble [1.5, 2.5, -0.5]+          , fromTexts (V.fromList ["alpha", "", "gamma"])+          , fromByteStrings (V.fromList ["\0\1", "\255", ""])+          , fromBools (V.fromList [True, False, True])+          , prim PDate32 [0, 19000, -1]+          ]+      , V.fromList+          [ prim PInt32 [4]+          , prim PInt64 [40]+          , prim PUInt16 [10]+          , prim PDouble [4.25]+          , fromTexts (V.fromList ["delta"])+          , fromByteStrings (V.fromList ["xyz"])+          , fromBools (V.fromList [False])+          , prim PDate32 [1]+          ]+      ]+    )+  , ( "nullable"+    , schemaOf+        [ plain "i32" True (AInt 32 True)+        , plain "s" True AUtf8+        , plain "flag" True ABool+        , plain "f64" True (AFloatingPoint DoublePrecision)+        , plain "ls" True ALargeUtf8+        , plain "fsb" True (AFixedSizeBinary 3)+        ]+    , [ V.fromList+          [ fromMaybes PInt32 (V.fromList [Just 1, Nothing, Just 3, Nothing])+          , fromMaybeTexts (V.fromList [Just "x", Nothing, Just "zz", Just ""])+          , fromMaybeBools (V.fromList [Nothing, Just True, Just False, Nothing])+          , fromMaybes PDouble (V.fromList [Just 0, Nothing, Nothing, Just 9.5])+          , fromMaybeLargeTexts (V.fromList [Just "large", Nothing, Just "u", Nothing])+          , ok (fromMaybeFixedSizeBinary 3 (V.fromList [Just "abc", Nothing, Just "def", Just "ghi"]))+          ]+      ]+    )+  , ( "struct"+    , schemaOf+        [ nested "s" False AStruct [plain "id" False (AInt 64 True), plain "name" False AUtf8]+        , nested "sm" True AStruct [plain "id" False (AInt 32 True), plain "flag" False ABool]+        ]+    , [ V.fromList+          [ ok (mkStruct 3 Nothing (V.fromList [("id", prim PInt64 [1, 2, 3]), ("name", fromTexts (V.fromList ["a", "b", "c"]))]))+          , ok+              ( mkStruct+                  3+                  (validityFromBools (V.fromList [True, False, True]))+                  (V.fromList [("id", prim PInt32 [1, 2, 3]), ("flag", fromBools (V.fromList [True, False, True]))])+              )+          ]+      ]+    )+  , ( "lists"+    , schemaOf+        [ nested "l" False AList [plain "item" False (AInt 32 True)]+        , nested "lm" True AList [plain "item" False (AInt 32 True)]+        , nested "ll" False ALargeList [plain "item" False AUtf8]+        , nested "fsl" False (AFixedSizeList 2) [plain "item" False (AInt 32 True)]+        , nested "fslm" True (AFixedSizeList 2) [plain "item" False (AInt 32 True)]+        ]+    , [ V.fromList+          [ ok (mkList Nothing (VS.fromList [0, 2, 2, 5]) (prim PInt32 [10, 20, 30, 40, 50]))+          , ok (mkList (validityFromBools (V.fromList [True, False, True])) (VS.fromList [0, 1, 1, 3]) (prim PInt32 [7, 8, 9]))+          , ok (mkLargeList Nothing (VS.fromList [0, 1, 3, 3]) (fromTexts (V.fromList ["p", "q", "r"])))+          , ok (mkFixedSizeList 2 3 Nothing (prim PInt32 [1, 2, 3, 4, 5, 6]))+          , ok (mkFixedSizeList 2 3 (validityFromBools (V.fromList [True, False, True])) (prim PInt32 [1, 2, 3, 4, 5, 6]))+          ]+      ]+    )+  , ( "map"+    , schemaOf+        [ nested "m" False (AMap False) [nested "entries" False AStruct [plain "key" False AUtf8, plain "value" False (AInt 32 True)]]+        ]+    , [V.singleton (ok (mkMap Nothing (VS.fromList [0, 2, 2, 3]) (fromTexts (V.fromList ["a", "b", "c"])) (prim PInt32 [1, 2, 3])))]+    )+  , ( "unions"+    , schemaOf+        [ nested "du" False (AUnion Dense (V.fromList [0, 1])) [plain "v_int" False (AInt 32 True), plain "v_text" False AUtf8]+        , nested "su" False (AUnion Sparse (V.fromList [0, 1])) [plain "flag" False ABool, plain "value" False (AInt 32 True)]+        ]+    , [ V.fromList+          [ ok (mkDenseUnion (VS.fromList [0, 1, 0]) (VS.fromList [0, 0, 1]) (V.fromList [prim PInt32 [100, 200], fromTexts (V.fromList ["hello"])]))+          , ok (mkSparseUnion (VS.fromList [0, 1, 0]) (V.fromList [fromBools (V.fromList [True, False, False]), prim PInt32 [0, 42, 0]]))+          ]+      ]+    )+  , ( "dictionary"+    , schemaOf [Field "d" False AUtf8 V.empty (Just (DictionaryEncoding 0 (AInt 32 True) False)) V.empty]+    , [V.singleton (ok (mkDictionary 0 (prim PInt32 [0, 1, 0, 2, 1]) (fromTexts (V.fromList ["a", "b", "c"]))))]+    )+  , ( "views"+    , schemaOf [plain "v" False AUtf8View, plain "bv" True ABinaryView]+    , [ V.fromList+          [ fromMaybeUtf8View (V.fromList (map Just ["short", "this string is definitely longer than twelve bytes", ""]))+          , fromMaybeBinaryView (V.fromList [Just "tiny", Nothing, Just "another payload well past the inline limit"])+          ]+      ]+    )+  , ( "run-end encoded"+    , schemaOf [nested "ree" True ARunEndEncoded [plain "run_ends" False (AInt 32 True), plain "values" True (AInt 64 True)]]+    , [V.singleton (ok (mkRunEndEncoded (prim PInt32 [3, 5, 8]) (fromMaybes PInt64 (V.fromList [Just 100, Nothing, Just 300]))))]+    )+  , ( "list views"+    , schemaOf+        [ nested "lv" False AListView [plain "item" False (AInt 32 True)]+        , nested "llv" False ALargeListView [plain "item" False (AInt 32 True)]+        ]+    , [ V.fromList+          [ ok (mkListView Nothing (VS.fromList [0, 2, 5]) (VS.fromList [2, 3, 1]) (prim PInt32 [10, 20, 30, 40, 50, 60]))+          , ok (mkLargeListView Nothing (VS.fromList [4, 0, 1]) (VS.fromList [2, 1, 3]) (prim PInt32 [1, 2, 3, 4, 5, 6]))+          ]+      ]+    )+  , ( "null"+    , schemaOf [plain "n" False ANull, plain "i" False (AInt 8 True)]+    , [V.fromList [ColNull 4, prim PInt8 [1, 2, 3, 4]]]+    )+  ]+  where+    prim :: VS.Storable a => PrimType a -> [a] -> ColumnArray+    prim t = primColumn t . VS.fromList+    ok = either (\e -> error ("Test.Arrow.Malformed: bad shape: " ++ e)) id+++-- * Decoder harness++{- | Generous for the corpus (a decode of a few kilobytes allocates a+few megabytes at most), far below what a hostile length claim used to+request.+-}+allocationLimitBytes :: Int64+allocationLimitBytes = 512 * 1024 * 1024+++{- | Evaluate a decoder result to normal form under an allocation limit+and a timeout. 'Nothing' means it behaved; 'Just' describes the crash.+-}+contained :: NFData a => a -> IO (Maybe String)+contained = containedWithin allocationLimitBytes+++containedWithin :: NFData a => Int64 -> a -> IO (Maybe String)+containedWithin limit x = do+  r <- timeout 10000000 . try @SomeException $ do+    setAllocationCounter limit+    enableAllocationLimit+    _ <- evaluate (force x)+    disableAllocationLimit+  disableAllocationLimit+  pure $ case r of+    Nothing -> Just "did not finish within 10s"+    Just (Left e) -> Just ("threw " ++ displayException e)+    Just (Right ()) -> Nothing+++{- | Allocation allowed for exercising the columns of one accepted+decode. A batch may legitimately claim up to a million rows that no+buffer backs (fieldless structs, null children), and the exercise+allocates per row; the decode itself stays under+'allocationLimitBytes'.+-}+exerciseLimitBytes :: Int64+exerciseLimitBytes = 4 * 1024 * 1024 * 1024+++{- | Run a column-producing decoder: force its result under the decoder+limits, then exercise every column it returns ('exerciseColumn'); an+exception there is a crash, and a column the row model rejects is+reported like one.+-}+containedColumns :: NFData a => (a -> [ColumnArray]) -> Either String a -> IO (Maybe String)+containedColumns cols r = do+  crash <- contained r+  case (crash, r) of+    (Just e, _) -> pure (Just e)+    (Nothing, Left _) -> pure Nothing+    (Nothing, Right a) -> do+      let checked = traverse exerciseColumn (cols a)+      exerciseCrash <- containedWithin exerciseLimitBytes checked+      pure $ case (exerciseCrash, checked) of+        (Just e, _) -> Just ("accessor on an accepted column: " ++ e)+        (Nothing, Left e) -> Just ("accepted a column that does not read back: " ++ e)+        _ -> Nothing+++{- | Every accessor over a column and, recursively, over its children:+typed element access, byte/text/list/dictionary/bool access, validity,+the O(n) conversions, dictionary expansion, a slice, 'Show' (its first+64 KiB), reflexive 'Eq' and 'copyColumn'. 'Left' when the column does+not read back through the row model or a copy differs from the+original; the 'Int' only exists to force everything.++Rows are walked in bounded windows ('columnWindows'): a run-end-encoded+array names its rows in run ends and an all-null or fieldless column+has nothing to read, so either can legitimately claim billions of rows+for a few hundred input bytes. The decoder must accept those, and no+per-row pass can walk them, so the head, middle and tail are exercised+instead.+-}+exerciseColumn :: ColumnArray -> Either String Int+exerciseColumn c = do+  kids <- traverse exerciseColumn (G.childColumns c)+  windows <- traverse exerciseWindow (columnWindows c)+  Right $!+    sum windows+      + length (take 65536 (show c))+      + nullCount c+      + sum kids+++{- | Rows one at a time in one column: how many rows 'columnWindows'+walks one by one when a column is longer than that.+-}+exerciseWindowRows :: Int+exerciseWindowRows = 3 * 4096+++{- | The rows 'exerciseColumn' walks: the column itself when it is that+short, otherwise equally sized head, middle and tail slices, so every+accessor still runs against both ends of each buffer and an interior+run.+-}+columnWindows :: ColumnArray -> [ColumnArray]+columnWindows c+  | n <= exerciseWindowRows = [c]+  | otherwise = map (\start -> sliceColumnArray start w c) [0, (n - w) `quot` 2, n - w]+  where+    !n = columnLength c+    !w = exerciseWindowRows `quot` 3+++{- | Upper bound on the values the row model materializes for a column:+its own logical length plus every child's cost. Lists and maps name+child ranges without rebasing the child, so a short parent window can+still reference a huge child; run-end columns and nulls shrink with a+row slice. Cheap, O(columns).+-}+columnCost :: ColumnArray -> Integer+columnCost c = toInteger (columnLength c) + sum (map columnCost (G.childColumns c))+++{- | Values one window may materialize through the row model. A window is+a few thousand rows, so this leaves headroom for wide nested rows while+keeping one accessor call to a few tens of megabytes.+-}+exerciseValueBudget :: Integer+exerciseValueBudget = 262144+++{- | Every accessor over every row of one bounded window. The row model+is only materialized when 'columnCost' says the window is small enough:+a run-end column or a list over one can name billions of values in a few+input bytes, and expanding those would measure the test's memory, not+the decoder's behaviour. The per-row accessors still run over every row+of the window either way.+-}+exerciseWindow :: ColumnArray -> Either String Int+exerciseWindow c =+  if columnCost c > exerciseValueBudget+    then Right $! foldl' (\acc i -> acc + rowScore i) 0 [0 .. n - 1]+    else do+      vs <- G.columnValues c+      expanded <- either (const (Right 0)) (fmap valuesSize . G.columnValues) (expandDictionary c)+      sliced <- valuesSize <$> G.columnValues (sliceColumnArray 1 (n - 2) c)+      if c /= c+        then Left "a decoded column is not equal to itself"+        else+          if copyColumn c /= c+            then Left "copyColumn changed a decoded column"+            else+              Right $!+                foldl' (\acc i -> acc + rowScore i) 0 [0 .. n - 1]+                  + valuesSize vs+                  + conversions+                  + expanded+                  + sliced+  where+    !n = columnLength c+    primScore i = case c of+      ColPrim t v xs -> withPrim t (maybe 0 (const 1) (primAt (PrimArray v xs) i))+      _ -> 0+    rowScore i =+      maybe 0 BS.length (anyBytesAt c i)+        + maybe 0 T.length (anyTextAt c i)+        + maybe 0 (\r -> childStart r + childLength r) (listRange c i)+        + fromMaybe 0 (dictKeyAt c i)+        + maybe 0 fromEnum (boolAt c i)+        + fromEnum (isValidAt (validity c) i)+        + primScore i+    sized :: AV.Element a => (a -> Int) -> Either String (AV.Vector (Maybe a)) -> Int+    sized f = either (const 0) (VG.foldl' (\acc x -> acc + maybe 1 f x) 0)+    conversions =+      sized T.length (toTextVector c)+        + sized BS.length (toBytesVector c)+        + sized fromEnum (toBoolVector c)+++-- | Force every row of the model, deeply.+valuesSize :: V.Vector G.Value -> Int+valuesSize = V.foldl' (\acc v -> acc + valueSize v) 0+  where+    valueSize = \case+      G.VNull -> 1+      G.VInt i -> fromIntegral (i `rem` 7)+      G.VF16 w -> fromIntegral w+      G.VF32 w -> fromIntegral w+      G.VF64 w -> fromIntegral (w `rem` 7)+      G.VBool b -> fromEnum b+      G.VText t -> T.length t+      G.VBytes b -> BS.length b+      G.VDecimal i -> fromIntegral (i `rem` 7)+      G.VPair a b -> fromIntegral a + fromIntegral b+      G.VTriple a b x -> fromIntegral a + fromIntegral b + fromIntegral x+      G.VList xs -> foldl' (\acc x -> acc + valueSize x) 1 xs+      G.VStruct kvs -> foldl' (\acc (k, x) -> acc + T.length k + valueSize x) 1 kvs+      G.VMap kvs -> foldl' (\acc (k, x) -> acc + valueSize k + valueSize x) 1 kvs+      G.VUnion t x -> fromIntegral t + valueSize x+++-- | Every public entry point that accepts untrusted bytes.+decoders :: [(String, ByteString -> IO (Maybe String))]+decoders =+  [ ("Arrow.Stream.decodeArrowStream", containedColumns (concatMap V.toList . snd) . decodeArrowStream defaultReadOptions)+  , ("Arrow.Stream.decodeArrowFile", containedColumns (concatMap V.toList . snd) . decodeArrowFile defaultReadOptions)+  , ("Arrow.Stream.openStreamReader/streamReaderToList", containedColumns (concatMap V.toList . snd) . viaStreamReader)+  , ("Arrow.Stream.openFileReader/fileReaderToList", containedColumns (concatMap V.toList) . viaFileReader)+  , ("Arrow.FlatBufferIPC.readArrowStreamFBInterleaved", contained . fmap (fmap (map frameNF)) . readArrowStreamFBInterleaved)+  , ("Arrow.FlatBufferIPC.readArrowFileFBWithDicts", contained . fmap (\(s, ds, bs) -> (s, map dictNF ds, bs)) . readArrowFileFBWithDicts)+  , ("Arrow.FlatBufferIPC.decodeTensorFrame", contained . fmap (\(t, rest) -> (show t, rest)) . decodeTensorFrame)+  , ("Arrow.FlatBufferIPC.decodeSparseTensorFrame", contained . fmap (\(t, rest) -> (show t, rest)) . decodeSparseTensorFrame)+  , ("Arrow.File.readArrowStream", contained . fmap (\a -> (File.asSchema a, File.asBatches a)) . File.readArrowStream)+  , ("Arrow.File.readArrowFile", contained . fmap (\a -> (File.afSchema a, File.afBatches a)) . File.readArrowFile)+  , ("Arrow.File.readArrowFileColumns", containedColumns (concatMap V.toList . V.toList . snd) . File.readArrowFileColumns defaultReadOptions)+  , ("Arrow.IPC.decodeIPCMessage", contained . decodeIPCMessage)+  , ("Arrow.FlatBufferIPC.readFileIndex", contained . fmap indexNF . readFileIndex)+  , ("Arrow.File.readIPCMessage@0", \bs -> contained (File.readIPCMessage bs 0))+  , ("Arrow.File.readIPCMessage@8", \bs -> contained (File.readIPCMessage bs 8))+  ]+  where+    viaStreamReader bs = do+      rd <- openStreamReader defaultReadOptions bs+      batches <- streamReaderToList rd+      Right (streamReaderSchema rd, batches)+    viaFileReader bs = do+      rd <- openFileReader defaultReadOptions bs+      batches <- fileReaderToList rd+      Right (V.fromList batches)+    indexNF idx = (show (fiSchema idx), map blockNF (V.toList (fiDictBlocks idx)), map blockNF (V.toList (fiRecordBatches idx)))+    blockNF b = (abOffset b, abMetaLen b, abBodyLen b)+    frameNF = \case+      SFDict db -> Left (dictNF db)+      SFBatch rb body -> Right (rb, body)+    dictNF db = (dbId db, dbIsDelta db, dbData db, dbBody db)+++-- | Run every decoder; fail with the decoder name and crash on the first misbehaviour.+noDecoderCrashes :: ByteString -> PropertyT IO ()+noDecoderCrashes bs = mapM_ one decoders+  where+    one (name, run) = do+      r <- evalIO (run bs)+      case r of+        Nothing -> success+        Just err -> do+          annotate (name ++ " " ++ err)+          failure++-- * Byte-level mutations++data Mutation+  = Truncate Int+  | FlipBits [(Int, Int)]+  | Overwrite [(Int, Word8)]+  | Insert Int ByteString+  | Delete Int Int+  | Duplicate Int Int+  deriving stock (Show)+++applyMutation :: Mutation -> ByteString -> ByteString+applyMutation m bs = case m of+  Truncate k -> BS.take k bs+  FlipBits ps -> foldl (\acc (i, b) -> modifyAt i (`complementBit` b) acc) bs ps+  Overwrite ps -> foldl (\acc (i, w) -> modifyAt i (const w) acc) bs ps+  Insert i extra -> BS.take i bs <> extra <> BS.drop i bs+  Delete i n -> BS.take i bs <> BS.drop (i + n) bs+  Duplicate i n -> BS.take (i + n) bs <> BS.take n (BS.drop i bs) <> BS.drop (i + n) bs+  where+    modifyAt i f s+      | i < 0 || i >= BS.length s = s+      | otherwise = BS.take i s <> BS.singleton (f (BS.index s i)) <> BS.drop (i + 1) s+++genPos :: Int -> Gen Int+genPos len = Gen.int (Range.constant 0 (max 0 (len - 1)))+++-- | Byte values that tend to sit on decoder boundaries.+genByte :: Gen Word8+genByte = Gen.frequency [(3, Gen.element [0x00, 0x01, 0x07, 0x08, 0x0C, 0x7F, 0x80, 0xFE, 0xFF]), (2, Gen.word8 Range.constantBounded)]+++genSeed :: [Seed] -> Gen Seed+genSeed = Gen.element+++mutationProperty :: CoverPercentage -> [Seed] -> (Int -> Gen Mutation) -> Property+mutationProperty acceptPct corpus genM = withTests 3000 . property $ do+  seed <- forAll (genSeed corpus)+  m <- forAll (genM (BS.length (seedBytes seed)))+  let mutated = applyMutation m (seedBytes seed)+  noDecoderCrashes mutated+  coverDepth acceptPct mutated+++-- Cutting a frame short is always rejected, so truncations only rarely+-- (at frame boundaries) leave something to accept.+prop_truncation :: [Seed] -> Property+prop_truncation corpus = mutationProperty 0 corpus $ \len -> Truncate <$> genPos len+++prop_bitFlips :: [Seed] -> Property+prop_bitFlips corpus = mutationProperty 5 corpus $ \len ->+  FlipBits <$> Gen.list (Range.linear 1 8) ((,) <$> genPos len <*> Gen.int (Range.constant 0 7))+++prop_overwrites :: [Seed] -> Property+prop_overwrites corpus = mutationProperty 5 corpus $ \len ->+  Overwrite <$> Gen.list (Range.linear 1 8) ((,) <$> genPos len <*> genByte)+++prop_splices :: [Seed] -> Property+-- Length-changing file mutations invalidate footer offsets even when their+-- stream-shaped prefix survived. Do not require malformed files to be accepted:+-- bit flips, overwrites, descriptor/schema mutations and zero-copy mutations+-- independently require materializer coverage. Keep all 3000 structural cases.+prop_splices corpus = mutationProperty 0 corpus $ \len ->+  Gen.choice+    [ Insert <$> genPos len <*> Gen.bytes (Range.linear 1 64)+    , Delete <$> genPos len <*> Gen.int (Range.linear 1 64)+    , Duplicate <$> genPos len <*> Gen.int (Range.linear 1 64)+    ]+++-- * Structured mutations: record batch descriptors++data DescTarget+  = BufOffset+  | BufLength+  | NodeLength+  | NodeNullCount+  | VariadicCount+  | BatchLength+  | ClaimedBodyLength+  deriving stock (Show, Eq, Enum, Bounded)+++data Value+  = Exactly Int64+  | Shift Int64+  | BodyPlus Int64+  deriving stock (Show)+++data DescMutation = DescMutation+  { dmFrame :: Int+  , dmTarget :: DescTarget+  , dmIndex :: Int+  , dmValue :: Value+  }+  deriving stock (Show)+++genValue :: Gen Value+genValue =+  Gen.choice+    [ Exactly <$> Gen.element [0, 1, -1, 7, 8, 13, 255, 2 ^ (31 :: Int) - 1, 2 ^ (31 :: Int), 2 ^ (32 :: Int), 2 ^ (40 :: Int), 2 ^ (62 :: Int), maxBound, minBound]+    , Shift <$> Gen.int64 (Range.linearFrom 0 (-64) 64)+    , BodyPlus <$> Gen.int64 (Range.linearFrom 0 (-16) 16)+    , Exactly <$> Gen.int64 Range.linearBounded+    ]+++resolveValue :: Value -> Int64 -> Int64 -> Int64+resolveValue v original bodyLen = case v of+  Exactly x -> x+  Shift d -> original + d+  BodyPlus d -> bodyLen + d+++-- | A decoded stream whose frames can be edited and re-serialized.+data Parsed = Parsed+  { pSchema :: Schema+  , pFrames :: [StreamFrame]+  }+++parseSeed :: Seed -> Maybe Parsed+parseSeed s = case seedEncoding s of+  StreamEnc -> either (const Nothing) (\(sch, frames) -> Just (Parsed sch frames)) (readArrowStreamFBInterleaved (seedBytes s))+  FileEnc ->+    either+      (const Nothing)+      (\(sch, dicts, batches) -> Just (Parsed sch (map SFDict dicts ++ map (uncurry SFBatch) batches)))+      (readArrowFileFBWithDicts (seedBytes s))+++-- | Re-serialize; @claimed@ overrides record batch body lengths in the metadata.+serialize :: Encoding -> Schema -> [(StreamFrame, Maybe Int64)] -> ByteString+serialize enc sch frames = case enc of+  StreamEnc ->+    BS.concat (encapsulateMessage (buildSchemaMessage sch) BS.empty : map frameBytes frames ++ [eos])+  FileEnc ->+    writeArrowFileFBWithDicts+      sch+      (mapMaybe (\(f, _) -> case f of SFDict db -> Just db; SFBatch {} -> Nothing) frames)+      (mapMaybe (\(f, _) -> case f of SFBatch rb body -> Just (rb, body); SFDict {} -> Nothing) frames)+  where+    eos = BS.pack [0xFF, 0xFF, 0xFF, 0xFF, 0, 0, 0, 0]+    frameBytes = \case+      (SFDict db, _) -> encapsulateMessage (buildDictionaryBatchMessage db) (dbBody db)+      (SFBatch rb body, claimed) ->+        encapsulateMessage (buildRecordBatchMessage rb (maybe (fromIntegral (BS.length body)) id claimed)) body+++applyDesc :: DescMutation -> [(StreamFrame, Maybe Int64)] -> [(StreamFrame, Maybe Int64)]+applyDesc dm frames = zipWith edit [0 ..] frames+  where+    nFrames = length frames+    target = if nFrames == 0 then -1 else dmFrame dm `mod` nFrames+    edit :: Int -> (StreamFrame, Maybe Int64) -> (StreamFrame, Maybe Int64)+    edit i fr@(frame, claimed)+      | i /= target = fr+      | otherwise = case frame of+          SFBatch rb body ->+            let bodyLen = fromIntegral (BS.length body)+            in case dmTarget dm of+                 ClaimedBodyLength -> (frame, Just (resolveValue (dmValue dm) (maybe bodyLen id claimed) bodyLen))+                 _ -> (SFBatch (editRb bodyLen rb) body, claimed)+          SFDict db ->+            let bodyLen = fromIntegral (BS.length (dbBody db))+            in (SFDict db {dbData = editRb bodyLen (dbData db)}, claimed)+    editRb bodyLen rb =+      let v orig = resolveValue (dmValue dm) orig bodyLen+          pick :: Int -> Int+          pick len = dmIndex dm `mod` len+          atIdx :: V.Vector a -> (a -> a) -> V.Vector a+          atIdx xs f+            | V.null xs = xs+            | otherwise = let k = pick (V.length xs) in xs V.// [(k, f (xs V.! k))]+          atIdxS :: VS.Storable a => VS.Vector a -> (a -> a) -> VS.Vector a+          atIdxS xs f+            | VS.null xs = xs+            | otherwise = let k = pick (VS.length xs) in xs VS.// [(k, f (xs VS.! k))]+      in case dmTarget dm of+           BufOffset -> rb {rbBuffers = atIdxS (rbBuffers rb) (\b -> b {bufOffset = v (bufOffset b)})}+           BufLength -> rb {rbBuffers = atIdxS (rbBuffers rb) (\b -> b {bufLength = v (bufLength b)})}+           NodeLength -> rb {rbNodes = atIdxS (rbNodes rb) (\n -> n {fnLength = v (fnLength n)})}+           NodeNullCount -> rb {rbNodes = atIdxS (rbNodes rb) (\n -> n {fnNullCount = v (fnNullCount n)})}+           VariadicCount -> rb {rbVariadicBufferCounts = atIdx (rbVariadicBufferCounts rb) v}+           BatchLength -> rb {rbLength = v (rbLength rb)}+           ClaimedBodyLength -> rb+++-- | Seeds whose frames the structured mutators can edit.+parsedSeeds :: [Seed] -> [(Seed, Parsed)]+parsedSeeds = mapMaybe (\s -> (,) s <$> parseSeed s)+++prop_descriptors :: [Seed] -> Property+prop_descriptors corpus = withTests 3000 . property $ do+  (seed, parsed) <- forAllWith (show . fst) (Gen.element (parsedSeeds corpus))+  muts <-+    forAll . Gen.list (Range.linear 1 3) $+      DescMutation+        <$> Gen.int (Range.constant 0 7)+        <*> Gen.enumBounded+        <*> Gen.int (Range.constant 0 31)+        <*> genValue+  let frames = foldr applyDesc (map (\f -> (f, Nothing)) (pFrames parsed)) muts+  let bytes = serialize (seedEncoding seed) (pSchema parsed) frames+  noDecoderCrashes bytes+  coverDepth 5 bytes+++-- * Structured mutations: schema / type confusion++data SchemaMutation+  = SetType [Int] ArrowType+  | ToggleNullable [Int]+  | DropChildren [Int]+  | AddChild [Int] ArrowType+  | ToggleDictionary [Int] ArrowType+  deriving stock (Show)+++confusingTypes :: [ArrowType]+confusingTypes =+  [ ANull+  , AInt 8 True+  , AInt 16 False+  , AInt 32 True+  , AInt 64 False+  , AInt 7 True+  , ABool+  , AUtf8+  , ABinary+  , ALargeUtf8+  , ALargeBinary+  , AFixedSizeBinary 0+  , AFixedSizeBinary 3+  , AFixedSizeBinary (-1)+  , ADecimal 10 2+  , ADecimal256 40 2+  , ADate DateDay+  , ATime Second 32+  , ATime Nanosecond 64+  , ATimestamp Nanosecond Nothing+  , ADuration Second+  , AInterval YearMonth+  , AInterval DayTime+  , AInterval MonthDayNano+  , AFloatingPoint Half+  , AFloatingPoint DoublePrecision+  , AList+  , ALargeList+  , AStruct+  , AFixedSizeList 0+  , AFixedSizeList 2+  , AFixedSizeList (-3)+  , AMap False+  , AUnion Dense (V.fromList [0, 1])+  , AUnion Sparse V.empty+  , AUnion Dense (V.fromList [5, 127])+  , ARunEndEncoded+  , AUtf8View+  , ABinaryView+  , AListView+  , ALargeListView+  ]+++-- | Preorder paths of every field in a schema.+fieldPaths :: V.Vector Field -> [[Int]]+fieldPaths fs = concat (zipWith (\i f -> [i] : map (i :) (fieldPaths (fieldChildren f))) [0 ..] (V.toList fs))+++editField :: [Int] -> (Field -> Field) -> V.Vector Field -> V.Vector Field+editField [] _ fs = fs+editField (i : rest) f fs = case fs V.!? i of+  Nothing -> fs+  Just fld ->+    let fld' = if null rest then f fld else fld {fieldChildren = editField rest f (fieldChildren fld)}+    in fs V.// [(i, fld')]+++applySchemaMutation :: SchemaMutation -> Schema -> Schema+applySchemaMutation m sch = sch {arrowFields = editField path edit (arrowFields sch)}+  where+    (path, edit) = case m of+      SetType p ty -> (p, \f -> f {fieldType = ty})+      ToggleNullable p -> (p, \f -> f {fieldNullable = not (fieldNullable f)})+      DropChildren p -> (p, \f -> f {fieldChildren = V.empty})+      AddChild p ty -> (p, \f -> f {fieldChildren = V.snoc (fieldChildren f) (plain "extra" True ty)})+      ToggleDictionary p ty -> (p, \f -> f {fieldDictionary = maybe (Just (DictionaryEncoding 0 ty False)) (const Nothing) (fieldDictionary f)})+++genSchemaMutation :: Schema -> Gen SchemaMutation+genSchemaMutation sch = do+  path <- Gen.element (case fieldPaths (arrowFields sch) of [] -> [[0]]; ps -> ps)+  ty <- Gen.element confusingTypes+  Gen.element [SetType path ty, ToggleNullable path, DropChildren path, AddChild path ty, ToggleDictionary path ty]+++prop_schemaConfusion :: [Seed] -> Property+prop_schemaConfusion corpus = withTests 3000 . property $ do+  (seed, parsed) <- forAllWith (show . fst) (Gen.element (parsedSeeds corpus))+  muts <- forAll (Gen.list (Range.linear 1 3) (genSchemaMutation (pSchema parsed)))+  let sch = foldr applySchemaMutation (pSchema parsed) muts+  let bytes = serialize (seedEncoding seed) sch (map (\f -> (f, Nothing)) (pFrames parsed))+  noDecoderCrashes bytes+  coverDepth 5 bytes+++{- | Both outcomes must be common: mutations that every decoder rejects+up front never reach the column materializers, and mutations that+never fail exercise nothing hostile.+-}+coverDepth :: CoverPercentage -> ByteString -> PropertyT IO ()+coverDepth acceptPct bs = do+  let accepted = case (decodeArrowStream defaultReadOptions bs, decodeArrowFile defaultReadOptions bs) of+        (Left _, Left _) -> False+        _ -> True+  cover acceptPct "a decoder accepted the mutated input" accepted+  cover 5 "every decoder rejected the mutated input" (not accepted)+++-- * Random bytes++prop_randomBytes :: Property+prop_randomBytes = withTests 3000 . property $ do+  bs <- forAll (Gen.bytes (Range.linear 0 512))+  noDecoderCrashes bs+++-- | Random payloads behind a valid continuation marker or file magic, so+-- they reach the metadata and footer parsers instead of failing at byte 0.+prop_framedRandom :: Property+prop_framedRandom = withTests 2000 . property $ do+  payload <- forAll (Gen.bytes (Range.linear 0 512))+  metaLen <- forAll (Gen.int32 (Range.linearFrom 0 (-8) 600))+  framing <- forAll (Gen.element [ContinuationFrame, FileMagic])+  noDecoderCrashes $ case framing of+    ContinuationFrame -> BS.pack [0xFF, 0xFF, 0xFF, 0xFF] <> le32 metaLen <> payload+    FileMagic -> "ARROW1\0\0" <> payload <> le32 metaLen <> "ARROW1"+++data Framing = ContinuationFrame | FileMagic+  deriving stock (Show, Eq)+++le32 :: Int32 -> ByteString+le32 n = BS.pack (map (\s -> fromIntegral (fromIntegral n `shiftR` s :: Word32)) [0, 8, 16, 24])+++le64 :: Int64 -> ByteString+le64 n = BS.pack (map (\s -> fromIntegral (fromIntegral n `shiftR` s :: Word64)) [0, 8 .. 56])+++-- * Hand-built hostile metadata++{- | A schema message whose struct fields each list the same child+twice, @depth@ levels deep: a few hundred bytes describing @2^depth@+fields once offsets are followed naively.+-}+aliasedSchemaStream :: Int -> IO ByteString+aliasedSchemaStream depth = do+  b <- FB.newBuilder+  structTy <- FB.writeTable b []+  let fieldTable kids =+        FB.writeTable+          b+          [ Nothing+          , Just (FB.scalar 1 (\bb -> FB.prependU8 bb 1))+          , Just (FB.scalar 1 (\bb -> FB.prependU8 bb 13))+          , Just (FB.voff structTy)+          , Nothing+          , FB.voff <$> kids+          ]+      grow 0 t = pure t+      grow k t = do+        kids <- FB.writeVectorOfOffsets b [t, t]+        t' <- fieldTable (Just kids)+        grow (k - 1 :: Int) t'+  leaf <- fieldTable Nothing+  top <- grow depth leaf+  fields <- FB.writeVectorOfOffsets b [top]+  sch <- FB.writeTable b [Nothing, Just (FB.voff fields)]+  msg <-+    FB.writeTable+      b+      [ Just (FB.scalar 2 (\bb -> FB.prependI16 bb 4))+      , Just (FB.scalar 1 (\bb -> FB.prependU8 bb 1))+      , Just (FB.voff sch)+      , Just (FB.scalar 8 (\bb -> FB.prependI64 bb 0))+      ]+  meta <- FB.finish b msg+  pure (encapsulateMessage meta BS.empty <> BS.pack [0xFF, 0xFF, 0xFF, 0xFF, 0, 0, 0, 0])+++-- | Every decoder survives the aliased schema and rejects it.+rejectsEverywhere :: ByteString -> PropertyT IO ()+rejectsEverywhere bs = do+  noDecoderCrashes bs+  case decodeArrowStream defaultReadOptions bs of+    Left _ -> success+    Right _ -> do+      annotate "decodeArrowStream accepted a hostile schema"+      failure+++prop_dagBomb :: Property+prop_dagBomb = withTests 1 . property $ evalIO (aliasedSchemaStream 40) >>= rejectsEverywhere+++prop_deepSchema :: Property+prop_deepSchema = withTests 1 . property $ do+  let deep = foldr (\_ f -> nested "s" False AStruct [f]) (plain "leaf" False (AInt 32 True)) [1 .. 100 :: Int]+  rejectsEverywhere (encapsulateMessage (buildSchemaMessage (schemaOf [deep])) BS.empty)+++{- | A nullable struct over a null-typed child claiming @2^30@ rows with+no body: every row it would materialize is backed by nothing.+-}+prop_unbackedRows :: Property+prop_unbackedRows = withTests 1 . property $ do+  let sch = schemaOf [nested "s" True AStruct [plain "n" True ANull]]+      rows = 2 ^ (30 :: Int)+      rb = RecordBatchDef rows (VS.fromList [FieldNode rows 0, FieldNode rows rows]) (VS.singleton (Buffer 0 0)) V.empty Nothing+  rejectsEverywhere (serialize StreamEnc sch [(SFBatch rb BS.empty, Nothing)])+++{- | The SIMD buffer validator: an offset and length that each fit but+whose sum overflows 'Int64' must be rejected, and offsets or lengths+between 2 and 4 GiB inside a large enough body must be accepted.+-}+prop_bufferValidator :: Property+prop_bufferValidator = withTests 1 . property $ do+  let batch bufs = RecordBatchDef 0 VS.empty (VS.fromList bufs) V.empty Nothing+      huge = 2 ^ (62 :: Int)+      threeGiB = 3 * 2 ^ (30 :: Int)+  -- One pair takes the scalar tail, two pairs the SIMD path.+  validateRecordBatchBuffers (batch [Buffer huge huge]) 1024 === False+  validateRecordBatchBuffers (batch [Buffer 0 8, Buffer huge huge]) 1024 === False+  validateRecordBatchBuffers (batch [Buffer threeGiB 8]) (threeGiB + 64) === True+  validateRecordBatchBuffers (batch [Buffer 0 8, Buffer threeGiB 8]) (threeGiB + 64) === True+  validateRecordBatchBuffers (batch [Buffer 0 threeGiB, Buffer threeGiB 8]) (threeGiB + 64) === True+++{- | One non-null int64 column of 8 rows whose data buffer is a+compressed envelope claiming @claim@ uncompressed bytes; the payload+really holds 64 bytes of @0x07@.+-}+compressedClaim :: BodyCompressionCodec -> Int64 -> ByteString+compressedClaim codec claim =+  let payload = either (const BS.empty) id (compressBufferEither codec (BS.replicate 64 7))+      env = le64 claim <> payload+      envLen = fromIntegral (BS.length env)+      rb = RecordBatchDef 8 (VS.singleton (FieldNode 8 0)) (VS.fromList [Buffer 0 0, Buffer 0 envLen]) V.empty (Just codec)+  in serialize StreamEnc (schemaOf [plain "x" False (AInt 64 True)]) [(SFBatch rb env, Nothing)]+++-- | The body-compression codecs this build was compiled with (cabal+-- flags @zstd@ / @lz4@). Checks that need a real compressed payload run+-- only over these.+availableCodecs :: [BodyCompressionCodec]+availableCodecs = filter bodyCompressionAvailable [BodyZstd, LZ4Frame]+++{- | A buffer's uncompressed-length header is a claim: one the payload+cannot expand to, or that disagrees with what it does expand to, is+rejected without inflating anything. The honest claim decodes. A codec+that is not compiled in is rejected whatever the claim.+-}+prop_decompressionClaims :: Property+prop_decompressionClaims = withTests 1 . property $ do+  mapM_ (\(codec, claim) -> rejectsEverywhere (compressedClaim codec claim)) $+    concatMap (\codec -> map ((,) codec) [2 ^ (40 :: Int), maxBound, 600, 65, 63, -2]) [BodyZstd, LZ4Frame]+  mapM_+    (\codec -> fmap snd (decodeArrowStream defaultReadOptions (compressedClaim codec 64)) === Right [V.singleton (primColumn PInt64 (VS.replicate 8 0x0707070707070707))])+    availableCodecs+++{- | Two non-null int64 columns of 8 rows whose data buffers are+compressed envelopes honestly claiming @64@ uncompressed bytes each;+the payloads really hold 64 bytes of @0x07@.+-}+compressedTwo :: BodyCompressionCodec -> ByteString+compressedTwo codec =+  let payload = either (const BS.empty) id (compressBufferEither codec (BS.replicate 64 7))+      env = le64 64 <> payload+      envLen = fromIntegral (BS.length env)+      body = env <> env+      rb =+        RecordBatchDef+          8+          (VS.fromList [FieldNode 8 0, FieldNode 8 0])+          (VS.fromList [Buffer 0 0, Buffer 0 envLen, Buffer envLen 0, Buffer envLen envLen])+          V.empty+          (Just codec)+  in serialize StreamEnc (schemaOf [plain "x" False (AInt 64 True), plain "y" False (AInt 64 True)]) [(SFBatch rb body, Nothing)]+++{- | Every 'ReadOptions' budget rejects a batch that claims past it and+accepts the same batch at its exact boundary: the declared row count,+the total declared values across the field nodes, the rows with no+buffer backing (a zero-width fixed-size binary names them all), and+the total uncompressed size a compressed body expands to.+-}+prop_readerBudgets :: Property+prop_readerBudgets = withTests 1 . property $ do+  let sch = schemaOf [plain "z" False (AFixedSizeBinary 0)]+      rows = 600 :: Int64+      rb = RecordBatchDef rows (VS.singleton (FieldNode rows 0)) (VS.fromList [Buffer 0 0, Buffer 0 0]) V.empty Nothing+      bytes = serialize StreamEnc sch [(SFBatch rb BS.empty, Nothing)]+      mustAccept what opts = case decodeArrowStream opts bytes of+        Left e -> annotate (what ++ " at its boundary was rejected: " ++ e) >> failure+        Right _ -> success+      mustReject what opts = case decodeArrowStream opts bytes of+        Left _ -> success+        Right _ -> annotate (what ++ " over budget was accepted") >> failure+      mustAcceptBytes what bs opts = case decodeArrowStream opts bs of+        Left e -> annotate (what ++ " at its boundary was rejected: " ++ e) >> failure+        Right _ -> success+      mustRejectBytes what bs opts = case decodeArrowStream opts bs of+        Left _ -> success+        Right _ -> annotate (what ++ " over budget was accepted") >> failure+  mustAccept "rows" defaultReadOptions+  mustAccept "rows" defaultReadOptions {roMaxRows = rows}+  mustReject "rows" defaultReadOptions {roMaxRows = rows - 1}+  mustAccept "values" defaultReadOptions {roMaxValues = rows}+  mustReject "values" defaultReadOptions {roMaxValues = rows - 1}+  mustAccept "unbacked rows" defaultReadOptions {roMaxUnbackedRows = rows}+  mustReject "unbacked rows" defaultReadOptions {roMaxUnbackedRows = rows - 1}+  mapM_+    ( \codec -> do+        let compressed = compressedTwo codec+        mustAcceptBytes "decompression" compressed defaultReadOptions+        mustAcceptBytes "decompression" compressed defaultReadOptions {roMaxDecompressedBytes = 2 * 64}+        mustRejectBytes "decompression" compressed defaultReadOptions {roMaxDecompressedBytes = 2 * 64 - 1}+    )+    availableCodecs+++-- * Sanity: the corpus itself is valid++prop_corpusDecodes :: [Seed] -> Property+prop_corpusDecodes corpus = withTests 1 . property $ mapM_ checkSeed corpus+  where+    checkSeed seed = do+      let decoded = case seedEncoding seed of+            StreamEnc -> () <$ decodeArrowStream defaultReadOptions (seedBytes seed)+            FileEnc -> () <$ decodeArrowFile defaultReadOptions (seedBytes seed)+      case decoded of+        Right () -> noDecoderCrashes (seedBytes seed)+        Left e -> do+          annotate (show seed ++ " does not decode: " ++ e)+          failure+++-- * Zero-copy decode under mutation++{- | Generated tables (every type, nested, sliced and denormalised+inputs) written by the stream and file writers, placed in an aligned+buffer so the decoders alias it, then mutated in place by bit flips and+byte overwrites. Every decoder must behave, and every accessor over+every row of an accepted column must neither throw nor read out of+range ('exerciseColumn').+-}+prop_zeroCopyAccessorFuzz :: Property+prop_zeroCopyAccessorFuzz = withTests 1000 . property $ do+  (sch, batches) <- forAll G.genTable+  enc <- forAll (Gen.element [StreamEnc, FileEnc])+  bytes <- evalEither $ case enc of+    StreamEnc -> encodeArrowStream defaultWriteOptions sch batches+    FileEnc -> encodeArrowFile defaultWriteOptions sch batches+  let len = BS.length bytes+  m <-+    forAll $+      Gen.choice+        [ FlipBits <$> Gen.list (Range.linear 1 8) ((,) <$> genPos len <*> Gen.int (Range.constant 0 7))+        , Overwrite <$> Gen.list (Range.linear 1 8) ((,) <$> genPos len <*> genByte)+        ]+  let mutated = alignedCopy (applyMutation m bytes)+  noDecoderCrashes mutated+  coverDepth 5 mutated+++-- | The bytes in a fresh 64-byte aligned buffer (so decoders alias them).+alignedCopy :: ByteString -> ByteString+alignedCopy bs =+  I.createAligned (BS.length bs) $ \p ->+    BSU.unsafeUseAsCStringLen bs $ \(src, n) -> copyBytes p (castPtr src) n
+ test/Test/Arrow/Props.hs view
@@ -0,0 +1,1022 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PatternSynonyms #-}++{- | Generated round-trip and algebraic properties for wireform-arrow.++Columns are compared through the logical row model of+"Test.Arrow.Gen" ('columnValues') and through the logical 'Eq' of+'ColumnArray', so dictionary unification, offset rebasing, slicing and+NaN bit patterns do not cause spurious mismatches while any change of+meaning does.++Besides the round trips, these properties pin the representation's+guarantees: decoded columns alias the input buffer (and stay correct+when the input is misaligned), 'copyColumn' detaches from it, the+writer's output is a fixed point of decode then encode, and the lazy+encoders produce the strict encoders' bytes.+-}+module Test.Arrow.Props (tests) where++import Arrow.Column (+  ColumnArray,+  PrimType (..),+  bitmapOffset,+  columnLength,+  columnTag,+  concatColumnArray,+  copyColumn,+  expandDictionary,+  fromBools,+  fromTexts,+  hasValiditySlot,+  mkDictionary,+  mkFixedSizeList,+  mkList,+  mkStruct,+  nullCount,+  primColumn,+  sliceColumnArray,+  takeColumnArray,+  validateMapKeysSorted,+  validity,+  validityBits,+  validityFromBools,+  validityNullCount,+  pattern ColDictionary,+  pattern ColFixedSizeList,+  pattern ColLargeList,+  pattern ColLargeBinary,+  pattern ColLargeUtf8,+  pattern ColBinary,+  pattern ColList,+  pattern ColMap,+  pattern ColNull,+  pattern ColRunEndEncoded,+  pattern ColStruct,+  pattern ColUtf8,+ )+import Arrow.Column.Internal qualified as I+import Arrow.File (readArrowFileColumns)+import Arrow.Read.Columns (defaultReadOptions)+import Arrow.FlatBufferIPC (+  DictBatch (..),+  StreamFrame (..),+  buildDictionaryBatchMessage,+  buildRecordBatchBytes,+  buildRecordBatchMessage,+  buildSchemaMessage,+  encapsulateMessage,+  readArrowStreamFBInterleaved,+  readArrowStreamFBWithDicts,+  writeArrowStreamFBWithDicts,+ )+import Arrow.Stream (+  DictHandling (..),+  WriteOptions (..),+  bodyCompressionAvailable,+  decodeArrowFile,+  decodeArrowStream,+  defaultWriteOptions,+  encodeArrowFile,+  encodeArrowFileLazy,+  encodeArrowStream,+  encodeArrowStreamLazy,+  openStreamReader,+  streamReaderProjected,+  streamReaderToList,+ )+import Arrow.Types+import Arrow.Write (writeArrowFile, writeArrowStream)+import Control.Monad (forM_, replicateM)+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Internal qualified as BSI+import Data.ByteString.Lazy qualified as BL+import Data.ByteString.Unsafe qualified as BSU+import Data.Either (isLeft, isRight)+import Data.Int (Int64)+import Data.List (elemIndex)+import Data.Maybe (isNothing)+import Data.String (fromString)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Vector qualified as V+import Data.Vector.Storable qualified as VS+import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)+import Foreign.Marshal.Utils (copyBytes)+import Foreign.Ptr (castPtr, minusPtr, plusPtr)+import Hedgehog+import Hedgehog.Gen qualified as Gen+import Hedgehog.Range qualified as Range+import Test.Arrow.Gen+++tests :: IO Bool+tests =+  checkParallel $+    Group+      "Arrow.Props"+      [ ("stream round-trip preserves every row", prop_streamRoundTrip)+      , ("big-endian strict/lazy stream/file and dictionary round-trips", prop_bigEndianRoundTrip)+      , ("file round-trip (encodeArrowFile / writeArrowFile / readArrowFileColumns)", prop_fileRoundTrip)+      , ("Arrow.Write.writeArrowStream round-trip", prop_writeStreamRoundTrip)+      , ("replacement-dictionary stream round-trip (eager and iterator)", prop_replaceDictsRoundTrip)+      , ("body compression round-trip", prop_compressionRoundTrip)+      , ("schema fingerprint survives round-trip", prop_fingerprintStable)+      , ("projection equals projecting the decoded table", prop_projection)+      , ("slice length and contents", prop_sliceContents)+      , ("slice of slice composes", prop_sliceOfSlice)+      , ("concat of adjacent slices is the original", prop_concatAdjacentSlices)+      , ("concat appends rows and lengths add", prop_concatAppends)+      , ("slice then round-trip equals round-trip then slice", prop_sliceCommutesWithRoundTrip)+      , ("take gathers rows by index", prop_take)+      , ("dictionary expansion keeps row values", prop_expandDictionary)+      , ("validateMapKeysSorted agrees with the model", prop_mapKeysSorted)+      , ("delta dictionary batches append", prop_deltaDictionary)+      , ("dictionary placeholder has the value type", prop_dictionaryPlaceholder)+      , ("reads a pyarrow-written Arrow file", prop_pyarrowFile)+      , ("fieldless structs and size-0 fixed-size lists keep their row count", prop_zeroWidthRowCount)+      , ("nested dictionaries round-trip, inner dictionaries written first", prop_nestedDictionaries)+      , ("nested dictionary batches resolve across deltas and replacements", prop_nestedDictionaryDeltaReplace)+      , ("emit-once rejects a combined dictionary its index type cannot address", prop_emitOnceIndexOverflow)+      , ("dictionary value columns that cannot be combined are rejected", prop_inconsistentDictionaryValues)+      , ("batches that do not fit the schema are rejected", prop_schemaMismatch)+      , ("decoded columns alias an aligned input and are correct for a misaligned one", prop_zeroCopyDecode)+      , ("copyColumn detaches a decoded column from the input", prop_copyColumnDetaches)+      , ("decode then encode reproduces writer output byte for byte", prop_canonicalRoundTrip)+      , ("lazy and strict encoders produce identical bytes", prop_lazyStrictIdentical)+      ]+++-- ============================================================+-- Helpers+-- ============================================================++-- | Compare two tables by logical rows, logical 'Eq' and top-level tag.+sameTable :: MonadTest m => [V.Vector ColumnArray] -> [V.Vector ColumnArray] -> m ()+sameTable expected actual = do+  length actual === length expected+  forM_ (zip expected actual) $ \(e, a) -> do+    map tag (V.toList a) === map tag (V.toList e)+    evalEither (batchValues a) >>= \av -> evalEither (batchValues e) >>= \ev -> av === ev+    V.toList a === V.toList e+++tag :: ColumnArray -> String+tag = columnTag+++values :: MonadTest m => ColumnArray -> m [Value]+values c = V.toList <$> evalEither (columnValues c)+++-- | Fixture construction that cannot fail unless the test itself is wrong.+fixture :: Either String ColumnArray -> ColumnArray+fixture = either (error . ("Test.Arrow.Props: bad fixture: " ++)) id+++-- | Constructor names of a column and all its children (dictionary keys and values included).+allTags :: ColumnArray -> [String]+allTags c = tag c : concatMap allTags (childColumns c)+++-- | Whether any column in the tree (dictionary values included) satisfies the predicate.+anyColumn :: (ColumnArray -> Bool) -> ColumnArray -> Bool+anyColumn p c = p c || any (anyColumn p) (childColumns c)+++fieldlessStruct, sizeZeroList, nestedDictionary, emptyDictionary, withNulls, offsetBase, bitOffset, denormalisedValidity, windowedRuns :: ColumnArray -> Bool+fieldlessStruct = \case+  ColStruct _ _ cs -> V.null cs+  _ -> False+sizeZeroList = \case+  ColFixedSizeList 0 _ _ _ -> True+  _ -> False+nestedDictionary = \case+  ColDictionary _ _ v -> hasDictionaries v+  _ -> False+emptyDictionary = \case+  ColDictionary _ keys v -> columnLength v == 0 && columnLength keys > 0+  _ -> False+withNulls c = nullCount c > 0+-- Offsets that do not start at zero (a window of a longer column).+offsetBase = \case+  ColUtf8 _ o _ -> firstNonZero o+  ColBinary _ o _ -> firstNonZero o+  ColLargeUtf8 _ o _ -> firstNonZero o+  ColLargeBinary _ o _ -> firstNonZero o+  ColList _ o _ -> firstNonZero o+  ColLargeList _ o _ -> firstNonZero o+  ColMap _ o _ _ -> firstNonZero o+  _ -> False+  where+    firstNonZero :: (VS.Storable o, Eq o, Num o) => VS.Vector o -> Bool+    firstNonZero o = not (VS.null o) && VS.head o /= 0+bitOffset c = maybe False ((/= 0) . bitmapOffset . validityBits) (validity c)+denormalisedValidity c = maybe False ((== 0) . validityNullCount) (validity c)+windowedRuns = \case+  ColRunEndEncoded off _ _ _ -> off /= 0+  _ -> False+++{- | Require the generator to reach the shapes the properties are about.+Nullability is a property of the field (a nullable field's column may+happen to hold no nulls), so the nullable shapes are counted on the+schema of a table with at least one batch; the physical shapes are+counted on the columns.+-}+coverShapes :: MonadTest m => Schema -> [ColumnArray] -> m ()+coverShapes sch cols = do+  let tags = concatMap allTags cols+      has t = t `elem` tags+      some p = any (anyColumn p) cols+      fields = concatMap fieldTree (V.toList (arrowFields sch))+      fieldTree f = f : concatMap fieldTree (V.toList (fieldChildren f))+      nullableField p = not (null cols) && any (\f -> fieldNullable f && p f) fields+      plainType p f = isNothing (fieldDictionary f) && p (fieldType f)+  forM_+    ["ColDenseUnion", "ColSparseUnion", "ColRunEndEncoded", "ColListView", "ColFixedSizeList", "ColDictionary", "ColUtf8View", "ColNull"]+    (\t -> cover 1 (fromString t) (has t))+  cover 1 "nullable large list view" (nullableField (plainType (== ALargeListView)))+  cover 1 "nullable map" (nullableField (plainType (\case AMap _ -> True; _ -> False)))+  cover 1 "nullable dictionary" (nullableField (\f -> fieldDictionary f /= Nothing))+  cover 1 "nullable decimal128" (nullableField (plainType (\case ADecimal {} -> True; _ -> False)))+  cover 1 "nullable decimal256" (nullableField (plainType (\case ADecimal256 {} -> True; _ -> False)))+  cover 1 "nullable binary view" (nullableField (plainType (== ABinaryView)))+  cover 1 "nullable struct" (nullableField (plainType (== AStruct)))+  cover 1 "nullable interval" (nullableField (plainType (\case AInterval _ -> True; _ -> False)))+  cover 1 "nulls in a column" (some withNulls)+  cover 1 "struct with no fields" (some fieldlessStruct)+  cover 1 "fixed-size list of size 0" (some sizeZeroList)+  cover 1 "dictionary inside dictionary values" (some nestedDictionary)+  cover 1 "all-null rows over an empty dictionary" (some emptyDictionary)+  cover 1 "offsets with a non-zero base" (some offsetBase)+  cover 1 "validity at a non-zero bit offset" (some bitOffset)+  cover 1 "validity present with zero nulls" (some denormalisedValidity)+  cover 1 "run-end-encoded window with a logical offset" (some windowedRuns)+++-- | A single generated column (any type, nested to depth 3) with its field.+genColumn :: Gen (Field, ColumnArray)+genColumn = do+  f <- genField 3 "x"+  n <- Gen.int (Range.linear 0 10)+  c <- genColumnFor f n+  pure (f, c)+++-- | Whether the bytes of @inner@ lie inside the memory of @outer@.+regionWithin :: ByteString -> ByteString -> Bool+regionWithin outer inner =+  let (fo, lo) = BSI.toForeignPtr0 outer+      (fi, li) = BSI.toForeignPtr0 inner+      d = unsafeForeignPtrToPtr fi `minusPtr` unsafeForeignPtrToPtr fo+  in d >= 0 && d + li <= lo+++-- | The bytes copied to a fresh 64-byte aligned buffer at byte offset @shift@.+placedAt :: Int -> ByteString -> ByteString+placedAt shift bs =+  BS.drop shift $ I.createAligned (BS.length bs + shift) $ \p ->+    BSU.unsafeUseAsCStringLen bs $ \(src, n) -> copyBytes (p `plusPtr` shift) (castPtr src) n+++-- | Every column of a table, children included.+allColumns :: [V.Vector ColumnArray] -> [ColumnArray]+allColumns = concatMap (concatMap tree . V.toList)+  where+    tree c = c : concatMap tree (childColumns c)+++-- ============================================================+-- Round trips+-- ============================================================++prop_streamRoundTrip :: Property+prop_streamRoundTrip = withTests 1000 . property $ do+  (sch, batches) <- forAll genTable+  coverShapes sch (concatMap V.toList batches)+  bytes <- evalEither (encodeArrowStream defaultWriteOptions sch batches)+  (_, got) <- evalEither (decodeArrowStream defaultReadOptions bytes)+  sameTable batches got+++-- Exercise the complete generated schema/layout vocabulary in the non-native+-- byte order, including sliced offsets, composite intervals, decimals, views,+-- nested dictionary values and replacement dictionaries.+prop_bigEndianRoundTrip :: Property+prop_bigEndianRoundTrip = withTests 1000 . property $ do+  (base, batches) <- forAll genTable+  codec <- forAll (Gen.element (Nothing : map Just (filter bodyCompressionAvailable [BodyZstd, LZ4Frame])))+  let sch = base {arrowEndianness = Big}+      opts = defaultWriteOptions {writeBodyCompression = codec}+      outputs =+        [ (False, encodeArrowStream opts sch batches)+        , (False, BL.toStrict <$> encodeArrowStreamLazy opts sch batches)+        , (True, encodeArrowFile opts sch batches)+        , (True, BL.toStrict <$> encodeArrowFileLazy opts sch batches)+        , (False, encodeArrowStream opts {writeDictHandling = DictReplaceOnChange} sch batches)+        ]+  forM_ outputs $ \(file, encoded) -> do+    bytes <- evalEither encoded+    (actualSchema, got) <- evalEither ((if file then decodeArrowFile else decodeArrowStream) defaultReadOptions bytes)+    actualSchema === sch+    sameTable batches got+++prop_fileRoundTrip :: Property+prop_fileRoundTrip = withTests 200 . property $ do+  (sch, batches) <- forAll genTable+  file <- evalEither (encodeArrowFile defaultWriteOptions sch batches)+  (_, got1) <- evalEither (decodeArrowFile defaultReadOptions file)+  sameTable batches got1+  written <- evalEither (writeArrowFile sch (V.fromList batches))+  (_, got2) <- evalEither (readArrowFileColumns defaultReadOptions written)+  sameTable batches (V.toList got2)+  (_, got3) <- evalEither (readArrowFileColumns defaultReadOptions file)+  sameTable batches (V.toList got3)+++prop_writeStreamRoundTrip :: Property+prop_writeStreamRoundTrip = withTests 100 . property $ do+  (sch, batches) <- forAll genTable+  bytes <- evalEither (writeArrowStream sch (V.fromList batches))+  (_, got) <- evalEither (decodeArrowStream defaultReadOptions bytes)+  sameTable batches got+++prop_replaceDictsRoundTrip :: Property+prop_replaceDictsRoundTrip = withTests 200 . property $ do+  (sch, batches) <- forAll genTable+  bytes <- evalEither (encodeArrowStream defaultWriteOptions {writeDictHandling = DictReplaceOnChange} sch batches)+  (_, got) <- evalEither (decodeArrowStream defaultReadOptions bytes)+  sameTable batches got+  rd <- evalEither (openStreamReader defaultReadOptions bytes)+  got' <- evalEither (streamReaderToList rd)+  sameTable batches got'++{- | Compression round-trips. A codec that is not compiled in (cabal+flags @zstd@ / @lz4@) is documented to fall back to uncompressed+batches; when the codec is available every batch must say so.+-}+prop_compressionRoundTrip :: Property+prop_compressionRoundTrip = withTests 100 . property $ do+  (sch, batches) <- forAll genTable+  codec <- forAll (Gen.element [LZ4Frame, BodyZstd])+  bytes <- evalEither (encodeArrowStream defaultWriteOptions {writeBodyCompression = Just codec} sch batches)+  let expected = if bodyCompressionAvailable codec then Just codec else Nothing+  (_, got) <- evalEither (decodeArrowStream defaultReadOptions bytes)+  sameTable batches got+  (_, dicts, frames) <- evalEither (readArrowStreamFBWithDicts bytes)+  map (rbBodyCompression . fst) frames === map (const expected) frames+  map (rbBodyCompression . dbData) dicts === map (const expected) dicts+++prop_fingerprintStable :: Property+prop_fingerprintStable = withTests 200 . property $ do+  (sch, batches) <- forAll genTable+  (sch1, _) <- evalEither (encodeArrowStream defaultWriteOptions sch batches >>= decodeArrowStream defaultReadOptions)+  (sch2, _) <- evalEither (encodeArrowFile defaultWriteOptions sch batches >>= decodeArrowFile defaultReadOptions)+  schemaFingerprint sch1 === schemaFingerprint sch+  schemaFingerprint sch2 === schemaFingerprint sch+  assert (schemaEquivalent sch1 sch)+++prop_projection :: Property+prop_projection = withTests 200 . property $ do+  (sch, batches) <- forAll genTable+  let names = V.toList (V.map fieldName (arrowFields sch))+  picked <- forAll (Gen.subsequence names >>= Gen.shuffle)+  bytes <- evalEither (encodeArrowStream defaultWriteOptions sch batches)+  rd <- evalEither (openStreamReader defaultReadOptions bytes)+  (psch, got) <- evalEither (streamReaderProjected picked rd)+  (_, full) <- evalEither (decodeArrowStream defaultReadOptions bytes)+  let idx nm = V.findIndex ((== nm) . fieldName) (arrowFields sch)+      project cols = V.fromList (map (\nm -> maybe (error "missing") (cols V.!) (idx nm)) picked)+  map fieldName (V.toList (arrowFields psch)) === picked+  sameTable (map project full) got+++-- ============================================================+-- Slicing, concatenation, gathering+-- ============================================================++prop_sliceContents :: Property+prop_sliceContents = withTests 500 . property $ do+  (_, c) <- forAll genColumn+  let n = columnLength c+  s <- forAll (Gen.int (Range.linear (-2) (n + 2)))+  l <- forAll (Gen.int (Range.linear (-2) (n + 2)))+  let s' = min n (max 0 s)+      l' = max 0 (min l (n - s'))+      sl = sliceColumnArray s l c+  columnLength sl === l'+  vs <- values c+  svs <- values sl+  svs === take l' (drop s' vs)+  tag sl === tag c+++prop_sliceOfSlice :: Property+prop_sliceOfSlice = withTests 300 . property $ do+  (_, c) <- forAll genColumn+  let n = columnLength c+  s1 <- forAll (Gen.int (Range.linear 0 n))+  l1 <- forAll (Gen.int (Range.linear 0 (n - s1)))+  s2 <- forAll (Gen.int (Range.linear 0 l1))+  l2 <- forAll (Gen.int (Range.linear 0 (l1 - s2)))+  let twice = sliceColumnArray s2 l2 (sliceColumnArray s1 l1 c)+      once = sliceColumnArray (s1 + s2) l2 c+  a <- values twice+  b <- values once+  a === b+  twice === once+++prop_concatAdjacentSlices :: Property+prop_concatAdjacentSlices = withTests 500 . property $ do+  (_, c) <- forAll genColumn+  let n = columnLength c+  k <- forAll (Gen.int (Range.linear 0 n))+  joined <- evalEither (concatColumnArray (sliceColumnArray 0 k c) (sliceColumnArray k (n - k) c))+  columnLength joined === n+  a <- values joined+  b <- values c+  a === b+  joined === c+++prop_concatAppends :: Property+prop_concatAppends = withTests 300 . property $ do+  f <- forAll (genField 3 "x")+  n1 <- forAll (Gen.int (Range.linear 0 8))+  n2 <- forAll (Gen.int (Range.linear 0 8))+  c1 <- forAll (genColumnFor f n1)+  c2 <- forAll (genColumnFor f n2)+  joined <- evalEither (concatColumnArray c1 c2)+  columnLength joined === n1 + n2+  a <- values joined+  v1 <- values c1+  v2 <- values c2+  a === v1 ++ v2+++prop_sliceCommutesWithRoundTrip :: Property+prop_sliceCommutesWithRoundTrip = withTests 200 . property $ do+  sch <- forAll genSchema+  batch <- forAll (genBatchFor sch)+  let n = if V.null batch then 0 else columnLength (V.head batch)+  s <- forAll (Gen.int (Range.linear 0 n))+  l <- forAll (Gen.int (Range.linear 0 (n - s)))+  (_, sliced) <- evalEither (encodeArrowStream defaultWriteOptions sch [V.map (sliceColumnArray s l) batch] >>= decodeArrowStream defaultReadOptions)+  (_, whole) <- evalEither (encodeArrowStream defaultWriteOptions sch [batch] >>= decodeArrowStream defaultReadOptions)+  sameTable (map (V.map (sliceColumnArray s l)) whole) sliced+++prop_take :: Property+prop_take = withTests 300 . property $ do+  (_, c) <- forAll genColumn+  let n = columnLength c+  ix <- forAll (if n == 0 then pure [] else Gen.list (Range.linear 0 12) (Gen.int (Range.linear 0 (n - 1))))+  taken <- evalEither (takeColumnArray (VS.fromList ix) c)+  vs <- values c+  tvs <- values taken+  tvs === map (vs !!) ix+  if n > 0 then pure () else assert (not (isRight (takeColumnArray (VS.singleton 0) c)))+++{- | Expansion keeps every row's value; a nullable dictionary column+whose rows are all null over an empty dictionary expands to an+all-null column of the value type. A dictionary's null rows stay null+in the expansion (the expanded column can hold them).+-}+prop_expandDictionary :: Property+prop_expandDictionary = withTests 300 . property $ do+  f <- forAll (Gen.choice [genField 0 "x", genDictionaryField 0 "x"])+  n <- forAll (Gen.int (Range.linear 0 10))+  c <- forAll (genColumnFor f n)+  cover 1 "all-null rows over an empty dictionary" (emptyDictionary c)+  e <- evalEither (expandDictionary c)+  columnLength e === n+  a <- values e+  b <- values c+  a === b+  assert (not (hasDictionaries e))+  case c of+    ColDictionary _ keys _ -> do+      assert (hasValiditySlot e || nullCount keys == 0)+      assert (nullCount e >= nullCount keys)+    _ -> success+++-- ============================================================+-- Map key order+-- ============================================================++{- | 'validateMapKeysSorted' accepts a map exactly when every non-null+entry's keys are non-null and non-decreasing under the documented+order, and rejects key types without an order.+-}+prop_mapKeysSorted :: Property+prop_mapKeysSorted = withTests 1000 . property $ do+  kind <- forAll (Gen.element ["orderable", "dictionary", "nullable", "unorderable" :: Text])+  lens <- forAll (Gen.list (Range.linear 0 5) (Gen.int (Range.linear 0 4)))+  let offs = scanl (+) 0 lens+      total = last offs+  keyField <- forAll $ case kind of+    "unorderable" -> do+      ty <- Gen.element [AInterval YearMonth, AInterval DayTime, AInterval MonthDayNano, ANull]+      pure (Field "key" False ty V.empty Nothing V.empty)+    "dictionary" -> do+      f <- genOrderableKeyField "key"+      pure f {fieldDictionary = Just (DictionaryEncoding 0 (AInt 16 True) False)}+    "nullable" -> (\f -> f {fieldNullable = True}) <$> genOrderableKeyField "key"+    _ ->+      Gen.choice+        [ genOrderableKeyField "key"+        , pure (Field "key" False (AFloatingPoint Half) V.empty Nothing V.empty)+        , pure (Field "key" False (AFloatingPoint Single) V.empty Nothing V.empty)+        , pure (Field "key" False (ADecimal256 40 0) V.empty Nothing V.empty)+        , pure (Field "key" False ABinaryView V.empty Nothing V.empty)+        , pure (Field "key" False (ADuration Second) V.empty Nothing V.empty)+        ]+  keys <- forAll (genColumnFor keyField total)+  mapValid <- forAll (Gen.list (Range.singleton (length lens)) (Gen.frequency [(1, pure False), (4, pure True)]))+  nullableMap <- forAll Gen.bool+  let o = VS.fromList (map fromIntegral offs)+      vals = ColNull total+      -- Raw construction: null keys are not something mkMap has to accept.+      col = I.ColMap (if nullableMap then validityFromBools (V.fromList mapValid) else Nothing) o keys vals+      entryOn i = not nullableMap || (mapValid !! i)+  kvs <- values keys+  let entries = zip [0 ..] (zip offs (drop 1 offs))+      entryKeys (s, e) = take (e - s) (drop s kvs)+      ok = all (\(i, se) -> not (entryOn i) || entrySorted (entryKeys se)) entries+      entrySorted ks = notElem VNull ks && and (zipWith (\x y -> compareKeyValues x y /= GT) ks (drop 1 ks))+      expected = kind /= "unorderable" && ok+  annotateShow (validateMapKeysSorted col)+  isRight (validateMapKeysSorted col) === expected+++-- ============================================================+-- Dictionaries at the wire level+-- ============================================================++{- | A dictionary followed by a delta dictionary batch: the record+batch's indices address the concatenation of both value columns.+-}+prop_deltaDictionary :: Property+prop_deltaDictionary = withTests 300 . property $ do+  ty <- forAll (Gen.filter (/= ANull) genFieldType)+  valuesNullable <- forAll Gen.bool+  idxTy <- forAll (AInt <$> Gen.element [8, 16, 32, 64] <*> Gen.bool)+  let field = Field "d" False ty V.empty (Just (DictionaryEncoding 7 idxTy False)) V.empty+      valuesField = field {fieldDictionary = Nothing, fieldNullable = valuesNullable}+      valuesSchema = defaultSchema (V.singleton valuesField)+      sch = defaultSchema (V.singleton field)+  k1 <- forAll (Gen.int (Range.linear 1 4))+  k2 <- forAll (Gen.int (Range.linear 1 4))+  part1 <- forAll (genColumnFor valuesField k1)+  part2 <- forAll (genColumnFor valuesField k2)+  n <- forAll (Gen.int (Range.linear 0 10))+  ix <- forAll (replicateM n (Gen.int (Range.linear 0 (k1 + k2 - 1))))+  let dict isDelta vals =+        let (drb, dbody) = layoutBatch valuesSchema (V.singleton vals)+            db = DictBatch {dbId = 7, dbIsDelta = isDelta, dbData = drb, dbBody = dbody}+        in encapsulateMessage (buildDictionaryBatchMessage db) dbody+      -- Raw: the keys address values this batch does not carry yet.+      keys = keyColumn idxTy (V.fromList (map Just ix))+      (rb, body) = layoutBatch sch (V.singleton (I.ColDictionary 7 keys part1))+      bytes =+        BS.concat+          [ encapsulateMessage (buildSchemaMessage sch) BS.empty+          , dict False part1+          , dict True part2+          , encapsulateMessage (buildRecordBatchMessage rb (fromIntegral (BS.length body))) body+          , BS.pack [0xff, 0xff, 0xff, 0xff, 0, 0, 0, 0]+          ]+  (_, got) <- evalEither (decodeArrowStream defaultReadOptions bytes)+  v1 <- values part1+  v2 <- values part2+  let allVals = v1 ++ v2+  case got of+    [cols] | [c] <- V.toList cols -> do+      gv <- values c+      gv === map (allVals !!) ix+    _ -> failure+++{- | A batch with no dictionary rows and no dictionary batch keeps a+typed empty placeholder whose constructor matches the value type, and+empty keys at the index type's width.+-}+prop_dictionaryPlaceholder :: Property+prop_dictionaryPlaceholder = withTests 200 . property $ do+  ty <- forAll (Gen.filter (/= ANull) genFieldType)+  nullable <- forAll Gen.bool+  let field = Field "d" nullable ty V.empty (Just (DictionaryEncoding 3 (AInt 32 True) False)) V.empty+      sch = defaultSchema (V.singleton field)+  empty <- forAll (genColumnFor field 0)+  sample <- forAll (genColumnFor field {fieldDictionary = Nothing, fieldNullable = False} 0)+  let bytes = writeArrowStreamFBWithDicts sch [] [layoutBatch sch (V.singleton empty)]+  (_, got) <- evalEither (decodeArrowStream defaultReadOptions bytes)+  case got of+    [cols] | [c] <- V.toList cols -> case c of+      ColDictionary _ keys vals -> do+        columnLength keys === 0+        tag keys === "ColInt32"+        tag vals === tag sample+      other -> annotateShow other >> failure+    _ -> failure+++{- | @test/golden/pa_file.arrow@ was written by pyarrow 25.0.1 with+@pa.ipc.new_file@: two batches over+@i: int32, s: string, d: dictionary<string, int32>, n: decimal128(10, 2), l: list<int64>@+(all fields nullable, as pyarrow defaults), batch 1 =+@i=[1,2,3], s=["a",null,"h\\u00e9llo"], d=["x","y",null], n=[1.50,null,-2.25], l=[[1,2],null,[]]@,+batch 2 = @i=[-7], s=[""], d=["y"], n=[null], l=[[9]]@, dictionary @["x","y"]@.+-}+prop_pyarrowFile :: Property+prop_pyarrowFile = withTests 1 . property $ do+  bytes <- evalIO (BS.readFile "test/golden/pa_file.arrow")+  (sch, batches) <- evalEither (readArrowFileColumns defaultReadOptions bytes)+  map fieldName (V.toList (arrowFields sch)) === ["i", "s", "d", "n", "l"]+  map fieldNullable (V.toList (arrowFields sch)) === replicate 5 True+  (_, viaStream) <- evalEither (decodeArrowFile defaultReadOptions bytes)+  map (map tag . V.toList) viaStream+    === replicate 2 ["ColInt32", "ColUtf8", "ColDictionary", "ColDecimal128", "ColList"]+  map (map nullCount . V.toList) viaStream === [[0, 1, 1, 1, 1], [0, 0, 0, 1, 0]]+  got <- evalEither (traverse batchValues (V.toList batches))+  got+    === [ [ [VInt 1, VInt 2, VInt 3]+          , [VText "a", VNull, VText "h\233llo"]+          , [VText "x", VText "y", VNull]+          , [VDecimal 150, VNull, VDecimal (-225)]+          , [VList [VInt 1, VInt 2], VNull, VList []]+          ]+        , [ [VInt (-7)]+          , [VText ""]+          , [VText "y"]+          , [VNull]+          , [VList [VInt 9]]+          ]+        ]+++-- ============================================================+-- Fieldless structs, size-0 fixed-size lists, nested dictionaries,+-- writer errors+-- ============================================================++{- | A struct with no fields and a fixed-size list of size 0 carry their+row count explicitly: it survives stream and file round trips at the+top level and nested in a list, nullable or not, and slicing and+concatenation keep it.+-}+prop_zeroWidthRowCount :: Property+prop_zeroWidthRowCount = withTests 200 . property $ do+  n <- forAll (Gen.int (Range.linear 0 20))+  structNullable <- forAll Gen.bool+  listNullable <- forAll Gen.bool+  valid <- forAll (V.fromList <$> Gen.list (Range.singleton n) Gen.bool)+  let item = Field "item" False (AInt 32 True) V.empty Nothing V.empty+      fieldless nm nullable = Field nm nullable AStruct V.empty Nothing V.empty+      sch =+        defaultSchema+          ( V.fromList+              [ fieldless "s" structNullable+              , Field "f" listNullable (AFixedSizeList 0) (V.singleton item) Nothing V.empty+              , Field "l" False AList (V.singleton (fieldless "item" False)) Nothing V.empty+              ]+          )+      mask nullable = if nullable then validityFromBools valid else Nothing+      structCol = fixture (mkStruct n (mask structNullable) V.empty)+      listCol = fixture (mkFixedSizeList 0 n (mask listNullable) (primColumn PInt32 VS.empty))+      -- Every list row holds two fieldless structs.+      nestedCol = fixture (mkList Nothing (VS.generate (n + 1) (fromIntegral . (* 2))) (fixture (mkStruct (2 * n) Nothing V.empty)))+      batch = V.fromList [structCol, listCol, nestedCol]+  stream <- evalEither (encodeArrowStream defaultWriteOptions sch [batch])+  (_, got) <- evalEither (decodeArrowStream defaultReadOptions stream)+  sameTable [batch] got+  map (map columnLength . V.toList) got === [[n, n, n]]+  file <- evalEither (encodeArrowFile defaultWriteOptions sch [batch])+  (_, gotFile) <- evalEither (decodeArrowFile defaultReadOptions file)+  sameTable [batch] gotFile+  k <- forAll (Gen.int (Range.linear 0 n))+  forM_ [structCol, listCol] $ \c -> do+    joined <- evalEither (concatColumnArray (sliceColumnArray 0 k c) (sliceColumnArray k (n - k) c))+    columnLength joined === n+    columnLength (sliceColumnArray k n c) === n - k+++{- | Dictionaries nested in dictionary values,+@o: dictionary\<struct\<e: dictionary\<utf8\>\>\>@ and+@p: dictionary\<list\<dictionary\<utf8\>\>\>@, round-trip through both+dictionary modes (eager and iterator reads) and the file writer, and+every inner dictionary batch is written before the outer dictionary+batch that needs it (in replacement mode, an outer dictionary is+re-sent with every inner one).+-}+prop_nestedDictionaries :: Property+prop_nestedDictionaries = withTests 200 . property $ do+  innerIdx <- forAll (AInt <$> Gen.element [8, 16, 32, 64] <*> Gen.bool)+  outerIdx <- forAll (AInt <$> Gen.element [8, 16, 32, 64] <*> Gen.bool)+  oNull <- forAll Gen.bool+  eNull <- forAll Gen.bool+  pNull <- forAll Gen.bool+  qNull <- forAll Gen.bool+  let dict did idx = Just (DictionaryEncoding did idx False)+      e = Field "e" eNull AUtf8 V.empty (dict 2 innerIdx) V.empty+      o = Field "o" oNull AStruct (V.singleton e) (dict 1 outerIdx) V.empty+      q = Field "item" qNull AUtf8 V.empty (dict 4 innerIdx) V.empty+      p = Field "p" pNull AList (V.singleton q) (dict 3 outerIdx) V.empty+      sch = defaultSchema (V.fromList [o, p])+  nBatches <- forAll (Gen.int (Range.linear 1 4))+  batches <- forAll (replicateM nBatches (genBatchFor sch))+  mode <- forAll (Gen.element [DictEmitOnce, DictReplaceOnChange])+  let opts = defaultWriteOptions {writeDictHandling = mode}+  bytes <- evalEither (encodeArrowStream opts sch batches)+  (_, got) <- evalEither (decodeArrowStream defaultReadOptions bytes)+  sameTable batches got+  rd <- evalEither (openStreamReader defaultReadOptions bytes)+  got' <- evalEither (streamReaderToList rd)+  sameTable batches got'+  file <- evalEither (encodeArrowFile opts sch batches)+  (_, gotFile) <- evalEither (decodeArrowFile defaultReadOptions file)+  sameTable batches gotFile+  (_, frames) <- evalEither (readArrowStreamFBInterleaved bytes)+  let groups = dictGroups frames+  annotateShow groups+  forM_ groups $ \ids ->+    forM_ [(2, 1), (4, 3)] $ \(inner, outer) -> case elemIndex inner ids of+      Nothing -> success+      Just i -> assert (maybe False (> i) (elemIndex outer ids))+  where+    -- Dictionary ids sent before each record batch, in stream order.+    dictGroups :: [StreamFrame] -> [[Int64]]+    dictGroups = go []+      where+        go acc [] = if null acc then [] else [reverse acc]+        go acc (SFDict db : rest) = go (dbId db : acc) rest+        go acc (SFBatch {} : rest) = reverse acc : go [] rest+++{- | Wire-level nested dictionaries: inner dictionary 2 (utf8) inside the+struct values of outer dictionary 1. The stream sends inner, outer, a+batch; an inner delta, an outer delta, a batch; an inner replacement,+an outer replacement, a batch; and a second inner replacement alone,+then a batch. Each outer dictionary batch resolves against the inner+dictionary in force when it arrives, so the last batch still reads the+values of the third.+-}+prop_nestedDictionaryDeltaReplace :: Property+prop_nestedDictionaryDeltaReplace = withTests 200 . property $ do+  let word = Gen.text (Range.linear 0 4) Gen.alphaNum+      indexInto k = Gen.int (Range.linear 0 (k - 1))+      rowsOver k = Gen.list (Range.linear 0 6) (indexInto k)+  a <- forAll (Gen.list (Range.linear 1 4) word)+  aDelta <- forAll (Gen.list (Range.linear 1 3) word)+  b <- forAll (Gen.list (Range.linear 1 4) word)+  c <- forAll (Gen.list (Range.linear 1 4) word)+  o1 <- forAll (Gen.list (Range.linear 1 4) (indexInto (length a)))+  o1Delta <- forAll (Gen.list (Range.linear 1 3) (indexInto (length a + length aDelta)))+  o2 <- forAll (Gen.list (Range.linear 1 4) (indexInto (length b)))+  r1 <- forAll (rowsOver (length o1))+  r2 <- forAll (rowsOver (length o1 + length o1Delta))+  r3 <- forAll (rowsOver (length o2))+  r4 <- forAll (rowsOver (length o2))+  let e = Field "e" False AUtf8 V.empty (Just (DictionaryEncoding 2 (AInt 8 True) False)) V.empty+      o = Field "o" False AStruct (V.singleton e) (Just (DictionaryEncoding 1 (AInt 16 True) False)) V.empty+      sch = defaultSchema (V.singleton o)+      innerKeys = primColumn PInt8 . VS.fromList . map fromIntegral+      outerKeys = primColumn PInt16 . VS.fromList . map fromIntegral+      -- Raw: keys address dictionaries sent in separate batches.+      innerPlaceholder = I.ColDictionary 2 (innerKeys []) (fromTexts V.empty)+      dictMsg schema did isDelta col =+        let (drb, dbody) = layoutBatch schema (V.singleton col)+            db = DictBatch {dbId = did, dbIsDelta = isDelta, dbData = drb, dbBody = dbody}+        in encapsulateMessage (buildDictionaryBatchMessage db) dbody+      inner isDelta ws = dictMsg (defaultSchema (V.singleton e {fieldDictionary = Nothing})) 2 isDelta (fromTexts (V.fromList ws))+      outer isDelta ix =+        dictMsg+          (defaultSchema (V.singleton o {fieldDictionary = Nothing}))+          1+          isDelta+          (I.ColStruct (length ix) Nothing (V.singleton ("e", I.ColDictionary 2 (innerKeys ix) (fromTexts V.empty))))+      batch ix =+        let (rb, body) = layoutBatch sch (V.singleton (I.ColDictionary 1 (outerKeys ix) (I.ColStruct 0 Nothing (V.singleton ("e", innerPlaceholder)))))+        in encapsulateMessage (buildRecordBatchMessage rb (fromIntegral (BS.length body))) body+      bytes =+        BS.concat+          [ encapsulateMessage (buildSchemaMessage sch) BS.empty+          , inner False a+          , outer False o1+          , batch r1+          , inner True aDelta+          , outer True o1Delta+          , batch r2+          , inner False b+          , outer False o2+          , batch r3+          , inner False c+          , batch r4+          , BS.pack [0xff, 0xff, 0xff, 0xff, 0, 0, 0, 0]+          ]+      structRows rows outerVals = map (\i -> VStruct [("e", VText (outerVals !! i))]) rows+      afterDelta = map ((a ++ aDelta) !!) (o1 ++ o1Delta)+      replaced = map (b !!) o2+      expected =+        [ [structRows r1 (map (a !!) o1)]+        , [structRows r2 afterDelta]+        , [structRows r3 replaced]+        , [structRows r4 replaced]+        ]+  (_, got) <- evalEither (decodeArrowStream defaultReadOptions bytes)+  evalEither (traverse batchValues got) >>= (=== expected)+  rd <- evalEither (openStreamReader defaultReadOptions bytes)+  got' <- evalEither (streamReaderToList rd)+  evalEither (traverse batchValues got') >>= (=== expected)+++{- | With emit-once dictionaries an int8-indexed field whose batches+combine to more than 128 dictionary values cannot be written (the+shifted indices would wrap); replacement dictionaries write each+batch's own dictionary and round-trip.+-}+prop_emitOnceIndexOverflow :: Property+prop_emitOnceIndexOverflow = withTests 100 . property $ do+  k1 <- forAll (Gen.int (Range.constant 1 120))+  k2 <- forAll (Gen.int (Range.constant 1 120))+  let field = Field "d" False AUtf8 V.empty (Just (DictionaryEncoding 0 (AInt 8 True) False)) V.empty+      sch = defaultSchema (V.singleton field)+      dictCol prefix k = fixture (mkDictionary 0 (primColumn PInt8 (VS.generate k fromIntegral)) (fromTexts (V.generate k (\i -> prefix <> T.pack (show i)))))+      batches = [V.singleton (dictCol "a" k1), V.singleton (dictCol "b" k2)]+      fits = k1 + k2 <= 128+  cover 20 "combined dictionary overflows int8" (not fits)+  cover 20 "combined dictionary fits int8" fits+  isRight (encodeArrowStream defaultWriteOptions sch batches) === fits+  isRight (encodeArrowFile defaultWriteOptions sch batches) === fits+  isRight (writeArrowStream sch (V.fromList batches)) === fits+  replaced <- evalEither (encodeArrowStream defaultWriteOptions {writeDictHandling = DictReplaceOnChange} sch batches)+  (_, got) <- evalEither (decodeArrowStream defaultReadOptions replaced)+  sameTable batches got+++{- | Value columns of one dictionary id that cannot be concatenated+(struct values whose child names differ between batches) are rejected+by the emit-once and file writers, and a values column of the wrong+type is rejected by every writer. Replacement dictionaries never+concatenate, so the first pair still writes in that mode.+-}+prop_inconsistentDictionaryValues :: Property+prop_inconsistentDictionaryValues = withTests 100 . property $ do+  n1 <- forAll (Gen.int (Range.linear 1 5))+  n2 <- forAll (Gen.int (Range.linear 1 5))+  let child = Field "x" False (AInt 32 True) V.empty Nothing V.empty+      field = Field "d" False AStruct (V.singleton child) (Just (DictionaryEncoding 0 (AInt 32 True) False)) V.empty+      sch = defaultSchema (V.singleton field)+      ix k = primColumn PInt32 (VS.generate k fromIntegral)+      structBatch nm k = V.singleton (fixture (mkDictionary 0 (ix k) (fixture (mkStruct k Nothing (V.singleton (nm, ix k))))))+      renamed = [structBatch "x" n1, structBatch "y" n2]+      wrongType = [structBatch "x" n1, V.singleton (fixture (mkDictionary 0 (ix n2) (fromTexts (V.replicate n2 "v"))))]+      replace = defaultWriteOptions {writeDictHandling = DictReplaceOnChange}+  assert (isLeft (encodeArrowStream defaultWriteOptions sch renamed))+  assert (isLeft (encodeArrowFile defaultWriteOptions sch renamed))+  assert (isRight (encodeArrowStream replace sch renamed))+  assert (isLeft (encodeArrowStream defaultWriteOptions sch wrongType))+  assert (isLeft (encodeArrowStream replace sch wrongType))+  assert (isLeft (encodeArrowFile defaultWriteOptions sch wrongType))+++{- | A batch that does not fit its schema is rejected by every writer:+a missing or extra column, a column whose type its field cannot hold,+columns of unequal length, and a struct or fixed-size list whose row+count disagrees with its child (built raw: the public constructors+refuse them). The unbroken batch writes.+-}+prop_schemaMismatch :: Property+prop_schemaMismatch = withTests 300 . property $ do+  sch <- forAll genSchema+  batch <- forAll (Gen.filter (\b -> columnLength (V.head b) > 0) (genBatchFor sch))+  let n = columnLength (V.head batch)+      f0 = V.head (arrowFields sch)+      c0 = V.head batch+      withFields fs = sch {arrowFields = V.fromList fs}+  breakage <- forAll (Gen.element ["missing", "extra", "type", "length", "struct", "fixed-size list" :: Text])+  doubled <- evalEither (concatColumnArray c0 c0)+  let (sch', bad) = case breakage of+        "missing" -> (sch, V.tail batch)+        "extra" -> (sch, V.snoc batch (ColNull n))+        "type" ->+          let wrong = if fieldType f0 == ABool && isNothing (fieldDictionary f0) then ColNull n else fromBools (V.replicate n False)+          in (sch, V.cons wrong (V.tail batch))+        "length" -> (withFields (f0 : f0 {fieldName = "dup"} : V.toList (V.tail (arrowFields sch))), V.cons c0 (V.cons doubled (V.tail batch)))+        "struct" ->+          ( withFields [Field "s" False AStruct (V.singleton f0) Nothing V.empty]+          , V.singleton (I.ColStruct (n + 1) Nothing (V.singleton (fieldName f0, c0)))+          )+        _ ->+          ( withFields [Field "l" False (AFixedSizeList 1) (V.singleton f0) Nothing V.empty]+          , V.singleton (I.ColFixedSizeList 1 (n + 1) Nothing c0)+          )+      replace = defaultWriteOptions {writeDictHandling = DictReplaceOnChange}+  assert (isRight (encodeArrowStream defaultWriteOptions sch [batch]))+  annotateShow (encodeArrowStream defaultWriteOptions sch' [bad])+  assert (isLeft (encodeArrowStream defaultWriteOptions sch' [bad]))+  assert (isLeft (encodeArrowStream replace sch' [bad]))+  assert (isLeft (encodeArrowFile defaultWriteOptions sch' [bad]))+  assert (isLeft (writeArrowStream sch' (V.singleton bad)))+  assert (isLeft (writeArrowFile sch' (V.singleton bad)))+++-- ============================================================+-- Representation guarantees+-- ============================================================++{- | Flat fixed-width, boolean, var-length and view columns decode+without copying: placed in an aligned buffer, every buffer of every+decoded column lies inside the input. Placed one byte off alignment+(so 2, 4, 8 and 16-byte elements are misaligned) the decode still+yields the same rows.+-}+prop_zeroCopyDecode :: Property+prop_zeroCopyDecode = withTests 300 . property $ do+  k <- forAll (Gen.int (Range.linear 1 4))+  fields <- forAll $ forM' [0 .. k - 1] $ \i -> do+    ty <- Gen.filter (/= ANull) genFieldType+    nullable <- Gen.bool+    pure (Field (T.pack ("c" ++ show i)) nullable ty V.empty Nothing V.empty)+  let sch = defaultSchema (V.fromList fields)+  nBatches <- forAll (Gen.int (Range.linear 1 3))+  batches <- forAll (replicateM nBatches (genBatchFor sch))+  bytes <- evalEither (encodeArrowStream defaultWriteOptions sch batches)+  let aligned = placedAt 0 bytes+      misaligned = placedAt 1 bytes+  (_, got) <- evalEither (decodeArrowStream defaultReadOptions aligned)+  sameTable batches got+  let outside = filter (\b -> not (BS.null b) && not (regionWithin aligned b)) (concatMap I.columnBuffers (allColumns got))+  annotateShow (length outside)+  assert (null outside)+  cover 10 "has a validity bitmap" (any (any (\c -> nullCount c > 0) . V.toList) got)+  (_, gotMis) <- evalEither (decodeArrowStream defaultReadOptions misaligned)+  sameTable batches gotMis+++{- | 'copyColumn' of any decoded column is equal to it and shares no+memory with the input bytes.+-}+prop_copyColumnDetaches :: Property+prop_copyColumnDetaches = withTests 200 . property $ do+  (sch, batches) <- forAll genTable+  bytes <- evalEither (encodeArrowStream defaultWriteOptions sch batches)+  (_, got) <- evalEither (decodeArrowStream defaultReadOptions bytes)+  let copies = map (V.map copyColumn) got+  sameTable got copies+  let attached = filter (\b -> not (BS.null b) && regionWithin bytes b) (concatMap I.columnBuffers (allColumns copies))+  annotateShow (length attached)+  assert (null attached)+++{- | Writer output is canonical: decoding it and encoding the decoded+schema and columns again with the same options gives the same bytes,+for streams and files, every dictionary mode and body compression.+-}+prop_canonicalRoundTrip :: Property+prop_canonicalRoundTrip = withTests 300 . property $ do+  (sch, batches) <- forAll genTable+  mode <- forAll (Gen.element [DictEmitOnce, DictReplaceOnChange])+  codec <- forAll (Gen.element [Nothing, Just LZ4Frame, Just BodyZstd])+  let opts = defaultWriteOptions {writeDictHandling = mode, writeBodyCompression = codec}+  stream <- evalEither (encodeArrowStream opts sch batches)+  (sch1, got1) <- evalEither (decodeArrowStream defaultReadOptions stream)+  again <- evalEither (encodeArrowStream opts sch1 got1)+  again === stream+  file <- evalEither (encodeArrowFile opts sch batches)+  (sch2, got2) <- evalEither (decodeArrowFile defaultReadOptions file)+  againFile <- evalEither (encodeArrowFile opts sch2 got2)+  againFile === file+++-- | The lazy encoders' output, made strict, is the strict encoders' output.+prop_lazyStrictIdentical :: Property+prop_lazyStrictIdentical = withTests 200 . property $ do+  (sch, batches) <- forAll genTable+  mode <- forAll (Gen.element [DictEmitOnce, DictReplaceOnChange])+  codec <- forAll (Gen.element [Nothing, Just LZ4Frame, Just BodyZstd])+  let opts = defaultWriteOptions {writeDictHandling = mode, writeBodyCompression = codec}+  strict <- evalEither (encodeArrowStream opts sch batches)+  lazy <- evalEither (encodeArrowStreamLazy opts sch batches)+  BL.toStrict lazy === strict+  strictFile <- evalEither (encodeArrowFile opts sch batches)+  lazyFile <- evalEither (encodeArrowFileLazy opts sch batches)+  BL.toStrict lazyFile === strictFile+++forM' :: Monad m => [a] -> (a -> m b) -> m [b]+forM' xs f = mapM f xs+++-- | 'buildRecordBatchBytes' for hand-built fixtures that must lay out.+layoutBatch :: Schema -> V.Vector ColumnArray -> (RecordBatchDef, ByteString)+layoutBatch sch cols = either (error . ("Test.Arrow.Props: fixture does not lay out: " ++)) id (buildRecordBatchBytes sch cols)
+ test/Test/Arrow/Sections.hs view
@@ -0,0 +1,279 @@+{-# LANGUAGE OverloadedStrings #-}++{- | Section-at-a-time reading of Arrow IPC containers: the footer+'FileIndex', the pure 'Arrow.Stream.FileReader', and the+semi-streaming 'Handle' iterators that decode one section live at a+time.++The parity tests compare every sectioning path against the eager+decoders over the same bytes; the corruption test proves the handle+iterator really works section by section (earlier batches arrive+before a broken later section is reported).+-}+module Test.Arrow.Sections (tests) where++import Arrow.Column (ColumnArray, PrimType (..), primColumn)+import Arrow.Column qualified as C+import Arrow.FlatBufferIPC (ArrowBlock (..), FileIndex (..), readFileIndex)+import Arrow.Stream (+  ReadOptions,+  decodeArrowFile,+  decodeArrowStream,+  defaultReadOptions,+  defaultWriteOptions,+  encodeArrowFile,+  encodeArrowStream,+  fileBatchesHandle,+  fileReaderSchema,+  fileReaderToList,+  openFileReader,+  streamBatchesHandle,+ )+import Arrow.Types (ArrowType (..), Schema, defaultLeafField, defaultSchema)+import Arrow.Types qualified as A+import Columnar.Stream (IterIO, IterIOStep (..), iterIOStep, iterIOToList)+import Control.Exception (bracket)+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Builder qualified as BB+import Data.ByteString.Lazy qualified as BL+import Data.Int (Int64)+import Data.List (isPrefixOf)+import Data.Vector qualified as V+import Data.Vector.Storable qualified as VS+import FlatBuffers.Reader qualified as FB+import Hedgehog+import System.Directory (getTemporaryDirectory, removeFile)+import System.IO (Handle, SeekMode (..), hClose, hSeek, openBinaryTempFile)+++tests :: IO Bool+tests =+  checkParallel $+    Group+      "Arrow.Sections"+      [ ("file sections decode to the eager batches", prop_fileSectionsMatchEager)+      , ("pyarrow file sections decode to the eager batches", prop_pyarrowFileSections)+      , ("handle iteration decodes one file section at a time", prop_fileBatchesHandle)+      , ("handle iteration decodes one stream frame at a time (dictionaries included)", prop_streamBatchesHandle)+      , ("a corrupt section fails only when its turn comes", prop_corruptSectionFailsLate)+      , ("Rust 64-byte aligned files use their footer index", prop_rustAlignedFile)+      , ("Rust empty 64-byte aligned files keep their schema", prop_rustEmptyFile)+      , ("footer blocks cannot address footer bytes", prop_blockCannotPointIntoFooter)+      , ("negative message body lengths fail every file reader", prop_negativeMessageBody)+      ]+++prop_rustAlignedFile :: Property+prop_rustAlignedFile = withTests 1 $ property $ do+  bytes <- evalIO (BS.readFile "test/golden/rs_file_alignment64.arrow")+  -- The old eager reader mistook this alignment padding for stream EOS.+  BS.take 4 (BS.drop 8 bytes) === BS.replicate 4 0+  let expected = [V.singleton (C.fromMaybes PInt32 (V.fromList [Just 1, Nothing, Just (-7), Just maxBound]))]+  (schema, batches) <- evalEither (decodeArrowFile defaultReadOptions bytes)+  V.map A.fieldName (A.arrowFields schema) === V.singleton "value"+  V.map A.fieldType (A.arrowFields schema) === V.singleton (AInt 32 True)+  batches === expected+  reader <- evalEither (openFileReader defaultReadOptions bytes)+  fileReaderSchema reader === schema+  streamed <- evalEither (fileReaderToList reader)+  streamed === expected+++prop_rustEmptyFile :: Property+prop_rustEmptyFile = withTests 1 $ property $ do+  bytes <- evalIO (BS.readFile "test/golden/rs_file_empty_alignment64.arrow")+  (schema, batches) <- evalEither (decodeArrowFile defaultReadOptions bytes)+  V.map A.fieldName (A.arrowFields schema) === V.singleton "value"+  V.map A.fieldType (A.arrowFields schema) === V.singleton (AInt 32 True)+  batches === []+  reader <- evalEither (openFileReader defaultReadOptions bytes)+  fileReaderSchema reader === schema+  streamed <- evalEither (fileReaderToList reader)+  streamed === []+++prop_blockCannotPointIntoFooter :: Property+prop_blockCannotPointIntoFooter = withTests 1 $ property $ do+  bytes <- evalIO (BS.readFile "test/golden/rs_file_alignment64.arrow")+  footerLength <- evalEither (FB.peekI32 bytes (BS.length bytes - 10))+  let footerStart = BS.length bytes - 10 - fromIntegral footerLength+      footer = BS.take (fromIntegral footerLength) (BS.drop footerStart bytes)+  root <- evalEither (FB.peekU32 footer 0)+  slots <- evalEither (FB.resolveTable footer (fromIntegral root))+  slot <- evalMaybe (slots 3)+  vector <- evalEither (FB.followUOffset footer slot)+  (_, positions) <- evalEither (FB.readVectorOfStructs footer vector 24)+  position <- evalMaybe (positions V.!? 0)+  let patchPosition = footerStart + position+      -- A tiny block fits within the whole file but lies in its footer.+      encoded = BL.toStrict (BB.toLazyByteString (BB.int64LE (fromIntegral footerStart :: Int64) <> BB.int32LE 8 <> BB.int32LE 0 <> BB.int64LE 0))+      corrupted = BS.take patchPosition bytes <> encoded <> BS.drop (patchPosition + 24) bytes+  case readFileIndex corrupted of+    Left _ -> success+    Right _ -> failure+++prop_negativeMessageBody :: Property+prop_negativeMessageBody = withTests 1 $ property $ do+  bytes <- evalIO (BS.readFile "test/golden/rs_file_alignment64.arrow")+  index <- evalEither (readFileIndex bytes)+  descriptor <- evalMaybe (fiRecordBatches index V.!? 0)+  let metadataStart = fromIntegral (abOffset descriptor) + 8+      metadata = BS.take (fromIntegral (abMetaLen descriptor) - 8) (BS.drop metadataStart bytes)+  root <- evalEither (FB.peekU32 metadata 0)+  slots <- evalEither (FB.resolveTable metadata (fromIntegral root))+  bodyLengthSlot <- evalMaybe (slots 3)+  let position = metadataStart + bodyLengthSlot+      encoded = BL.toStrict (BB.toLazyByteString (BB.int64LE (-1)))+      corrupted = BS.take position bytes <> encoded <> BS.drop (position + 8) bytes+      rejected result = case result of Left _ -> True; Right _ -> False+  assert (rejected (decodeArrowFile defaultReadOptions corrupted))+  reader <- evalEither (openFileReader defaultReadOptions corrupted)+  assert (rejected (fileReaderToList reader))+  handled <- evalIO (withTempFile corrupted (drainFile defaultReadOptions))+  assert (rejected handled)+++-- * Fixtures+++-- | A single-column schema and three batches of different lengths.+smallBatches :: (Schema, [V.Vector ColumnArray])+smallBatches =+  ( defaultSchema (V.fromList [defaultLeafField "a" False (AInt 32 True)])+  ,+    [ V.singleton (primColumn PInt32 (VS.fromList [1, 2, 3]))+    , V.singleton (primColumn PInt32 (VS.fromList [4]))+    , V.singleton (primColumn PInt32 (VS.fromList [5, 6, 7, 8]))+    ]+  )+++fileBytesOf :: (Schema, [V.Vector ColumnArray]) -> Either String ByteString+fileBytesOf (sch, batches) = encodeArrowFile defaultWriteOptions sch batches+++-- * Properties+++{- | The pure footer-indexed reader and the eager decoder agree on+every batch, and the footer's block table names one block per record+batch and no dictionary blocks.+-}+prop_fileSectionsMatchEager :: Property+prop_fileSectionsMatchEager = withTests 1 . property $ do+  let (sch, batches) = smallBatches+  bytes <- evalEither (fileBytesOf (sch, batches))+  (_, eager) <- evalEither (decodeArrowFile defaultReadOptions bytes)+  length eager === length batches+  idx <- evalEither (readFileIndex bytes)+  V.length (fiRecordBatches idx) === length batches+  V.null (fiDictBlocks idx) === True+  rd <- evalEither (openFileReader defaultReadOptions bytes)+  fileReaderSchema rd === sch+  got <- evalEither (fileReaderToList rd)+  got === eager+++-- | The same parity against @test/golden/pa_file.arrow@, written by pyarrow.+prop_pyarrowFileSections :: Property+prop_pyarrowFileSections = withTests 1 . property $ do+  bytes <- evalIO (BS.readFile "test/golden/pa_file.arrow")+  (_, eager) <- evalEither (decodeArrowFile defaultReadOptions bytes)+  idx <- evalEither (readFileIndex bytes)+  V.length (fiRecordBatches idx) === length eager+  rd <- evalEither (openFileReader defaultReadOptions bytes)+  got <- evalEither (fileReaderToList rd)+  got === eager+++-- | 'fileBatchesHandle' over a temp file yields exactly the eager batches.+prop_fileBatchesHandle :: Property+prop_fileBatchesHandle = withTests 1 . property $ do+  let (sch, batches) = smallBatches+  bytes <- evalEither (fileBytesOf (sch, batches))+  (_, eager) <- evalEither (decodeArrowFile defaultReadOptions bytes)+  got <- evalIO (withTempFile bytes (drainFile defaultReadOptions))+  got === Right eager+++{- | 'streamBatchesHandle' over a temp stream yields exactly the eager+batches, including for a stream with dictionary batches (which the+reader must apply as it walks).+-}+prop_streamBatchesHandle :: Property+prop_streamBatchesHandle = withTests 1 . property $ do+  let (sch, batches) = smallBatches+  streamBytes <- evalEither (encodeArrowStream defaultWriteOptions sch batches)+  (_, eager) <- evalEither (decodeArrowStream defaultReadOptions streamBytes)+  got <- evalIO (withTempFile streamBytes (drainStream defaultReadOptions))+  got === Right eager+  dictBytes <- evalIO (BS.readFile "test/golden/pa_dict.arrows")+  (_, eagerDicts) <- evalEither (decodeArrowStream defaultReadOptions dictBytes)+  gotDicts <- evalIO (withTempFile dictBytes (drainStream defaultReadOptions))+  gotDicts === Right eagerDicts+++{- | Zeroing the last section's message header costs nothing before+its turn: the handle iterator hands over every earlier batch and only+then reports the broken section.+-}+prop_corruptSectionFailsLate :: Property+prop_corruptSectionFailsLate = withTests 1 . property $ do+  let (sch, batches) = smallBatches+  bytes <- evalEither (fileBytesOf (sch, batches))+  (_, eager) <- evalEither (decodeArrowFile defaultReadOptions bytes)+  idx <- evalEither (readFileIndex bytes)+  let lastBlock = V.unsafeIndex (fiRecordBatches idx) (V.length (fiRecordBatches idx) - 1)+      off = fromIntegral (abOffset lastBlock)+      metaLen = fromIntegral (abMetaLen lastBlock)+      broken = BS.take off bytes <> BS.replicate metaLen 0 <> BS.drop (off + metaLen) bytes+  (tag, got) <- evalIO $ withTempFile broken $ \h -> do+    eiter <- fileBatchesHandle defaultReadOptions h+    case eiter of+      Left e -> pure ("open failed: " ++ e, [])+      Right it -> drainSteps it []+  take (length eager - 1) got === take (length eager - 1) eager+  length got === length eager - 1+  True === ("error" `isPrefixOf` tag)+++-- * Helpers+++-- | Step an iterator to its end, collecting the batches it yields.+drainSteps :: IterIO (V.Vector ColumnArray) -> [V.Vector ColumnArray] -> IO (String, [V.Vector ColumnArray])+drainSteps it acc = do+  estep <- iterIOStep it+  case estep of+    Left e -> pure ("error: " ++ e, reverse acc)+    Right IterIODone -> pure ("done", reverse acc)+    Right (IterIOYield batch next) -> drainSteps next (batch : acc)+++-- | Write bytes to a temp file, position the handle at 0, run the action.+withTempFile :: ByteString -> (Handle -> IO a) -> IO a+withTempFile bs act =+  bracket+    (getTemporaryDirectory >>= \dir -> openBinaryTempFile dir "wf-arrow-sections.tmp")+    (\(path, h) -> hClose h >> removeFile path)+    (\(_path, h) -> BS.hPut h bs >> hSeek h AbsoluteSeek 0 >> act h)+++-- | Drain a file iterator into the batches it yields.+drainFile :: ReadOptions -> Handle -> IO (Either String [V.Vector ColumnArray])+drainFile opts h = do+  eiter <- fileBatchesHandle opts h+  case eiter of+    Left e -> pure (Left e)+    Right it -> iterIOToList it+++-- | Drain a stream iterator into the batches it yields.+drainStream :: ReadOptions -> Handle -> IO (Either String [V.Vector ColumnArray])+drainStream opts h = do+  eiter <- streamBatchesHandle opts h+  case eiter of+    Left e -> pure (Left e)+    Right it -> iterIOToList it
+ test/Test/Arrow/Vector.hs view
@@ -0,0 +1,391 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PatternSynonyms #-}++{- | Properties of "Arrow.Vector" against the boxed "Data.Vector" model.++For every element type the generic API ('VG.fromList', 'VG.generate',+slices at arbitrary offsets, 'VG.map', 'VG.filter', 'VG.concat',+'VG.thaw' / 'VG.freeze' / 'VG.modify', growth, 'VG.force', 'Eq', 'Ord')+must agree with the same operation on a list. Mutable operations run on+a slice of a larger vector so bitmaps sit at bit offsets that are not a+multiple of 8 and var-length rows share their store with neighbours.++The column conversions are checked against the logical row model of+"Test.Arrow.Gen" on generated columns, on the same columns after an IPC+round trip (decoded columns alias the input; one input lives on the+GHC heap and one in malloc memory, which takes the text copy path), and+on slices of both.+-}+module Test.Arrow.Vector (tests) where++import Arrow.Column+import Arrow.Stream (decodeArrowStream, defaultReadOptions, defaultWriteOptions, encodeArrowStream)+import Arrow.Types (ArrowType (..), DictionaryEncoding (..), Field (..), IntervalUnit (..), Precision (..), defaultField, defaultLeafField, defaultSchema)+import Arrow.Vector (Element, Vector, maybeValues)+import Control.Monad (forM_)+import Control.Monad.ST (ST, runST)+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Internal qualified as BSI+import Data.Foldable (foldl')+import Data.Either (isLeft)+import Data.Int (Int32, Int64)+import Data.List (uncons)+import Data.Maybe (isNothing)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Text.Encoding qualified as TE+import Data.Vector qualified as V+import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as GM+import Foreign.ForeignPtr (newForeignPtr)+import Foreign.Marshal.Alloc (finalizerFree, mallocBytes)+import Foreign.Marshal.Utils (copyBytes)+import Foreign.Ptr (castPtr)+import Hedgehog+import Hedgehog.Gen qualified as Gen+import Hedgehog.Range qualified as Range+import Test.Arrow.Gen (Value (..), columnValues, genColumnFor)+++tests :: IO Bool+tests =+  checkParallel $+    Group+      "Arrow.Vector"+      [ ("Int8", laws (Gen.int8 Range.linearBounded))+      , ("Int16", laws (Gen.int16 Range.linearBounded))+      , ("Int32", laws i32)+      , ("Int64", laws i64)+      , ("Int", laws (Gen.int Range.linearBounded))+      , ("Word8", laws (Gen.word8 Range.linearBounded))+      , ("Word16", laws (Gen.word16 Range.linearBounded))+      , ("Word32", laws (Gen.word32 Range.linearBounded))+      , ("Word64", laws (Gen.word64 Range.linearBounded))+      , ("Word", laws (Gen.word Range.linearBounded))+      , ("Float", laws (Gen.float (Range.linearFrac (-1e6) 1e6)))+      , ("Double", laws (Gen.double (Range.linearFrac (-1e12) 1e12)))+      , ("Float16", laws (Float16 <$> Gen.word16 Range.linearBounded))+      , ("IntervalDayTime", laws (IntervalDayTime <$> i32 <*> i32))+      , ("IntervalMonthDayNano", laws (IntervalMonthDayNano <$> i32 <*> i32 <*> i64))+      , ("Decimal128", laws genDecimal128)+      , ("Decimal256", laws (decimal256FromInteger <$> Gen.integral (Range.linearFrom 0 (negate (2 ^ (255 :: Int))) (2 ^ (255 :: Int) - 1))))+      , ("Bool", laws Gen.bool)+      , ("Text", laws genText)+      , ("ByteString", laws genBytes)+      , ("Maybe Int64", laws (maybeOf i64))+      , ("Maybe Bool", laws (maybeOf Gen.bool))+      , ("Maybe Decimal128", laws (maybeOf genDecimal128))+      , ("Maybe Text", laws (maybeOf genText))+      , ("Maybe ByteString", laws (maybeOf genBytes))+      , ("Maybe (Maybe Int32)", laws (maybeOf (maybeOf i32)))+      , ("Vector (Maybe Int32)", laws (genVec (maybeOf i32)))+      , ("Maybe (Vector (Maybe Int32))", laws (maybeOf (genVec (maybeOf i32))))+      , ("Vector Text", laws (genVec genText))+      , ("Maybe (Vector (Maybe Text))", laws (maybeOf (genVec (maybeOf genText))))+      , ("Vector (Vector Int64)", laws (genVec (genVec i64)))+      , ("new initialises: Nothing, False, empty rows", prop_newInitialises)+      , ("unsafeThaw of a converted text vector keeps rows apart", prop_unsafeThawConverted)+      , ("conversions agree with the row model on generated, decoded and sliced columns", prop_conversions)+      ]+++-- ============================================================+-- Generators+-- ============================================================++i32 :: Gen Int32+i32 = Gen.int32 Range.linearBounded+++i64 :: Gen Int64+i64 = Gen.int64 Range.linearBounded+++genDecimal128 :: Gen Decimal128+genDecimal128 = decimal128FromInteger <$> Gen.integral (Range.linearFrom 0 (negate (2 ^ (127 :: Int))) (2 ^ (127 :: Int) - 1))+++maybeOf :: Gen a -> Gen (Maybe a)+maybeOf g = Gen.frequency [(1, pure Nothing), (3, Just <$> g)]+++-- | Often a suffix of a longer text, so rows start inside their array.+genText :: Gen Text+genText = do+  t <- Gen.text (Range.linear 0 10) Gen.unicode+  k <- Gen.int (Range.linear 0 (T.length t))+  pure (T.drop k t)+++genBytes :: Gen ByteString+genBytes = do+  b <- Gen.bytes (Range.linear 0 10)+  k <- Gen.int (Range.linear 0 (BS.length b))+  pure (BS.drop k b)+++-- | Often a slice of a longer vector.+genVec :: Element a => Gen a -> Gen (Vector a)+genVec g = do+  xs <- Gen.list (Range.linear 0 6) g+  k <- Gen.int (Range.linear 0 (length xs))+  pure (VG.drop k (VG.fromList xs))+++-- ============================================================+-- Generic API against the list model+-- ============================================================++data Op a+  = Write Int a+  | Set Int Int a+  | -- | destination, source, length (the ranges may overlap)+    Move Int Int Int+  | -- | copy from a separate mutable vector+    CopyIn Int [a]+  | -- | copy from an immutable vector+    CopyFrom Int [a]+  deriving stock (Show)+++genOps :: Int -> Gen a -> Gen [Op a]+genOps n g+  | n <= 0 = pure []+  | otherwise = Gen.list (Range.linear 0 8) (Gen.choice [write, set, move, copy CopyIn, copy CopyFrom])+  where+    write = Write <$> Gen.int (Range.constant 0 (n - 1)) <*> g+    set = do+      s <- Gen.int (Range.constant 0 n)+      l <- Gen.int (Range.constant 0 (n - s))+      Set s l <$> g+    move = do+      l <- Gen.int (Range.constant 0 n)+      Move <$> Gen.int (Range.constant 0 (n - l)) <*> Gen.int (Range.constant 0 (n - l)) <*> pure l+    copy k = do+      l <- Gen.int (Range.constant 0 n)+      k <$> Gen.int (Range.constant 0 (n - l)) <*> Gen.list (Range.singleton l) g+++applyOp :: Element a => VG.Mutable Vector s a -> Op a -> ST s ()+applyOp mv = \case+  Write i x -> GM.write mv i x+  Set s l x -> GM.set (GM.slice s l mv) x+  Move d s l -> GM.move (GM.slice d l mv) (GM.slice s l mv)+  CopyIn d xs -> do+    src <- VG.thaw (VG.fromList xs)+    GM.copy (GM.slice d (length xs) mv) src+  CopyFrom d xs -> VG.copy (GM.slice d (length xs) mv) (VG.fromList xs)+++modelOp :: [a] -> Op a -> [a]+modelOp xs = \case+  Write i x -> splice i 1 [x]+  Set s l x -> splice s l (replicate l x)+  Move d s l -> splice d l (take l (drop s xs))+  CopyIn d ys -> splice d (length ys) ys+  CopyFrom d ys -> splice d (length ys) ys+  where+    splice at l ys = take at xs ++ ys ++ drop (at + l) xs+++laws :: (Element a, Ord a, Show a) => Gen a -> Property+laws gen = withTests 150 . property $ do+  xs <- forAll (Gen.list (Range.linear 0 70) gen)+  ys <- forAll (Gen.list (Range.linear 0 20) gen)+  y <- forAll gen+  let n = length xs+      v = VG.fromList xs+      w = VG.fromList ys+      bv = V.fromList xs+  -- construction, indexing, conversion to and from boxed vectors+  VG.toList v === xs+  VG.length v === n+  map (v VG.!) [0 .. n - 1] === xs+  VG.convert v === bv+  VG.convert bv === v+  VG.generate n (bv V.!) === v+  VG.toList (VG.unfoldr uncons xs `asTypeOf` v) === xs+  VG.toList (VG.replicate 11 y `asTypeOf` v) === replicate 11 y+  VG.toList (VG.snoc (VG.cons y v) y) === [y] ++ xs ++ [y]+  -- slices at any offset, and slices of slices+  i <- forAll (Gen.int (Range.linear 0 n))+  k <- forAll (Gen.int (Range.linear 0 (n - i)))+  let s = VG.slice i k v+      sm = take k (drop i xs)+  VG.toList s === sm+  j <- forAll (Gen.int (Range.linear 0 k))+  l <- forAll (Gen.int (Range.linear 0 (k - j)))+  VG.toList (VG.slice j l s) === take l (drop j sm)+  -- map, filter and concat over a slice+  keep <- V.fromList <$> forAll (Gen.list (Range.singleton k) Gen.bool)+  VG.toList (VG.imap (\ix x -> if keep V.! ix then x else y) s) === zipWith (\b x -> if b then x else y) (V.toList keep) sm+  VG.toList (VG.ifilter (\ix _ -> keep V.! ix) s) === map snd (filter fst (zip (V.toList keep) sm))+  VG.toList (VG.concat [s, w, v, s]) === sm ++ ys ++ xs ++ sm+  VG.toList (v <> w) === xs ++ ys+  -- force copies only the referenced rows and keeps them+  VG.force s === s+  VG.toList (VG.force s) === sm+  -- mutable operations on a window of a larger vector+  ops <- forAll (genOps k gen)+  let modified = VG.modify (\mv -> forM_ ops (applyOp (GM.slice i k mv))) v+  VG.toList modified === take i xs ++ foldl' modelOp sm ops ++ drop (i + k) xs+  let thawed = runST $ do+        mv <- VG.thaw s+        forM_ ops (applyOp mv)+        VG.freeze mv+  VG.toList thawed === foldl' modelOp sm ops+  VG.toList s === sm+  -- growing keeps the rows and appends writable slots+  extra <- forAll (Gen.int (Range.linear 0 20))+  let grown = runST $ do+        mv <- VG.thaw s+        mv' <- GM.grow mv extra+        forM_ [k .. k + extra - 1] $ \ix -> GM.write mv' ix y+        VG.freeze mv'+  VG.toList grown === sm ++ replicate extra y+  -- Eq and Ord are the list's+  (v == w) === (xs == ys)+  compare v w === compare xs ys+  compare s w === compare sm ys+  (VG.fromList sm == s) === True+++prop_newInitialises :: Property+prop_newInitialises = withTests 50 . property $ do+  k <- forAll (Gen.int (Range.linear 0 70))+  let fresh :: Element a => Vector a+      fresh = runST (GM.new k >>= VG.freeze)+  assert (VG.all isNothing (fresh :: Vector (Maybe Int64)))+  assert (VG.all isNothing (fresh :: Vector (Maybe Text)))+  assert (VG.all isNothing (fresh :: Vector (Maybe (Vector Int32))))+  assert (VG.all not (fresh :: Vector Bool))+  assert (VG.all T.null (fresh :: Vector Text))+  assert (VG.all BS.null (fresh :: Vector ByteString))+  assert (VG.all VG.null (fresh :: Vector (Vector Int64)))+++{- | A utf8 column converts with starts and ends sharing one offsets+buffer; mutating the unsafely thawed vector must still treat rows as+independent.+-}+prop_unsafeThawConverted :: Property+prop_unsafeThawConverted = withTests 100 . property $ do+  rows <- forAll (Gen.list (Range.linear 1 20) genText)+  x <- forAll genText+  i <- forAll (Gen.int (Range.linear 0 (length rows - 1)))+  converted <- evalEither (toTextVector (fromTexts (V.fromList rows)))+  let values = maybeValues converted+      written = runST $ do+        mv <- VG.unsafeThaw values+        GM.write mv i x+        VG.unsafeFreeze mv+  VG.toList written === take i rows ++ [x] ++ drop (i + 1) rows+++-- ============================================================+-- Conversions against the row model+-- ============================================================++-- | Fields whose columns some conversion accepts; lists hold int32.+genConvField :: Gen Field+genConvField = do+  nullable <- Gen.bool+  let leaf = defaultLeafField "c" nullable+      list ty = defaultField "c" nullable ty (V.singleton (defaultLeafField "item" True (AInt 32 True)))+  Gen.choice+    [ leaf <$> Gen.element [AInt 8 True, AInt 16 False, AInt 32 True, AInt 64 True, AInt 64 False, AFloatingPoint Half, AFloatingPoint Single, AFloatingPoint DoublePrecision, ADecimal 38 4, ADecimal256 76 10, AInterval DayTime, AInterval MonthDayNano]+    , leaf <$> Gen.element [ABool, AUtf8, ALargeUtf8, ABinary, ALargeBinary, AUtf8View, ABinaryView]+    , leaf . AFixedSizeBinary <$> Gen.int (Range.linear 0 4)+    , list <$> Gen.element [AList, ALargeList, AListView, ALargeListView]+    , list . AFixedSizeList <$> Gen.int (Range.linear 0 3)+    , do+        ty <- Gen.element [AUtf8, ABinary, ALargeUtf8]+        idx <- AInt <$> Gen.element [8, 16, 32, 64] <*> Gen.bool+        pure (leaf ty) {fieldDictionary = Just (DictionaryEncoding 0 idx False)}+    ]+++prop_conversions :: Property+prop_conversions = withTests 400 . property $ do+  f <- forAll genConvField+  n <- forAll (Gen.int (Range.linear 0 40))+  c <- forAll (genColumnFor f n)+  bytes <- evalEither (encodeArrowStream defaultWriteOptions (defaultSchema (V.singleton f)) [V.singleton c])+  heap <- decodedFrom bytes+  malloced <- evalIO (mallocCopy bytes) >>= decodedFrom+  forM_ [c, heap, malloced] $ \col -> do+    checkConversions f col+    i <- forAll (Gen.int (Range.linear 0 n))+    k <- forAll (Gen.int (Range.linear 0 (n - i)))+    checkConversions f (sliceColumnArray i k col)+  where+    decodedFrom bs = do+      (_, batches) <- evalEither (decodeArrowStream defaultReadOptions bs)+      case batches of+        [batch] | V.length batch == 1 -> pure (V.head batch)+        _ -> annotate "expected one batch of one column" >> failure+++-- | A copy of the bytes in malloc memory (not the GHC heap).+mallocCopy :: ByteString -> IO ByteString+mallocCopy bs = do+  let len = BS.length bs+  p <- mallocBytes (max 1 len)+  BS.useAsCStringLen bs $ \(src, _) -> copyBytes p (castPtr src) len+  fp <- newForeignPtr finalizerFree p+  pure (BSI.BS fp len)+++checkConversions :: Field -> ColumnArray -> PropertyT IO ()+checkConversions f c = do+  model <- evalEither (columnValues c)+  let rows = V.toList model+      ty = fieldType f+      textKind = ty `elem` [AUtf8, ALargeUtf8, AUtf8View]+      bytesKind = textKind || ty `elem` [ABinary, ALargeBinary, ABinaryView] || isFixedBinary ty+      listKind = ty `elem` [AList, ALargeList, AListView, ALargeListView] || isFixedList ty+  case c of+    ColPrim t _ _ -> withPrim t $ do+      arr <- maybe (annotate "asPrim" >> failure) pure (asPrim t c)+      let v = toMaybeVector arr+      VG.length v === length rows+      -- every row through the generic index path, rebuilt as a column+      evalEither (columnValues (fromMaybes t (VG.convert v))) >>= (=== model)+      evalEither (columnValues (fromMaybeVector t v)) >>= (=== model)+      evalEither (columnValues (fromMaybeVector t (VG.force v))) >>= (=== model)+    ColBool {} -> do+      v <- evalEither (toBoolVector c)+      map (maybe VNull VBool) (VG.toList v) === rows+      evalEither (columnValues (fromBoolVector (VG.force v))) >>= (=== model)+    _ -> success+  if textKind+    then do+      v <- evalEither (toTextVector c)+      map (maybe VNull VText) (VG.toList v) === rows+      VG.force v === v+    else assert (isLeft (toTextVector c))+  if bytesKind+    then do+      v <- evalEither (toBytesVector c)+      map (maybe VNull VBytes) (VG.toList v) === map asBytes rows+      VG.force v === v+    else success+  if listKind+    then do+      v <- evalEither (toListVector int32s c)+      map (maybe VNull (VList . map (maybe VNull (VInt . toInteger)) . VG.toList)) (VG.toList v) === rows+      VG.force v === v+    else success+  where+    int32s ch = maybe (Left ("not an int32 child: " ++ columnTag ch)) (Right . toMaybeVector) (asPrim PInt32 ch)+    asBytes = \case+      VText t -> VBytes (TE.encodeUtf8 t)+      other -> other+    isFixedBinary = \case+      AFixedSizeBinary _ -> True+      _ -> False+    isFixedList = \case+      AFixedSizeList _ -> True+      _ -> False+
+ test/golden/pa_decimal32.arrow view

binary file changed (absent → 506 bytes)

+ test/golden/pa_decimal32.arrows view

binary file changed (absent → 320 bytes)

+ test/golden/pa_decimal64.arrow view

binary file changed (absent → 522 bytes)

+ test/golden/pa_decimal64.arrows view

binary file changed (absent → 336 bytes)

+ test/golden/pa_dict.arrows view

binary file changed (absent → 528 bytes)

+ test/golden/pa_file.arrow view

binary file changed (absent → 2146 bytes)

+ test/golden/pa_int32.arrows view

binary file changed (absent → 304 bytes)

+ test/golden/pa_mixed.arrows view

binary file changed (absent → 576 bytes)

+ test/golden/rs_file_alignment64.arrow view

binary file changed (absent → 762 bytes)

+ test/golden/rs_file_empty_alignment64.arrow view

binary file changed (absent → 418 bytes)

+ wireform-arrow.cabal view
@@ -0,0 +1,190 @@+cabal-version: 3.0+name: wireform-arrow+version: 0.1.0.0+synopsis: Apache Arrow IPC schema framing and record batch materialization+description:+  Arrow column arrays and IPC file/stream codecs, with schema and record+  derivation and optional LZ4 and Zstandard buffer compression.+category: Data, Codec+license: BSD-3-Clause+license-file: LICENSE+author: Ian Duncan+maintainer: ian@iankduncan.com+copyright: 2026 Ian Duncan+build-type: Simple+extra-source-files:+  test-interop/pyarrow_interop.py+  test/golden/*.arrow+  test/golden/*.arrows++flag zstd+  description: Enable Zstandard body-buffer compression for Arrow IPC.+  default: False+  manual: True++flag lz4+  description: Enable LZ4 frame-format body-buffer compression for Arrow IPC.+  default: False+  manual: True++common defaults+  ghc-options:+    -Wall+    -Wcompat+    -Widentities+    -Wincomplete-record-updates+    -Wincomplete-uni-patterns+    -Wpartial-fields+    -Wredundant-constraints+    -O2+  default-language: GHC2021+  default-extensions:+    StrictData+    DerivingStrategies+    DeriveGeneric+    DeriveAnyClass+    OverloadedStrings+    LambdaCase++library+  import: defaults+  hs-source-dirs: src+  exposed-modules:+    Arrow.Column+    Arrow.Column.Buffer+    Arrow.Column.Builder+    Arrow.Column.Internal+    Arrow.Derive+    Arrow.File+    Arrow.FlatBufferIPC+    Arrow.FlatBufferIPC.Common+    Arrow.FlatBufferIPC.Read+    Arrow.FlatBufferIPC.Write+    Arrow.IPC+    Arrow.Read.Columns+    Arrow.Record+    Arrow.Record.Generic+    Arrow.Record.TH+    Arrow.Stream+    Arrow.Types+    Arrow.Vector+    Arrow.Vector.Internal+    Arrow.Write+    Arrow.Write.Columns+  build-depends:+    base >= 4.20 && < 5,+    wireform-core >= 0.3 && < 0.4,+    wireform-columnar-core == 0.1.*,+    wireform-derive == 0.2.*,+    wireform-flatbuffers == 0.1.*,+    bytestring >= 0.11 && < 0.13,+    text >= 2.0 && < 2.2,+    vector >= 0.13 && < 0.14,+    primitive >= 0.7 && < 0.10,+    containers >= 0.6 && < 0.9,+    crypton >= 2.2.6 && < 2.3,+    deepseq >= 1.4 && < 1.6,+    ram >= 0.20.1 && < 0.23,+    template-haskell >= 2.18 && < 2.25+  if flag(zstd)+    cpp-options: -DHAVE_ZSTD+    build-depends: zstd >= 0.1 && < 0.2+  if flag(lz4)+    cpp-options: -DHAVE_LZ4+    -- lz4-hs bundles LZ4 1.9.3 and exposes Codec.Lz4's frame-format+    -- compress / decompress functions. Older Hackage "lz4" only+    -- implements the block format (incompatible with arrow-cpp).+    build-depends: lz4-hs >= 0.1 && < 0.2++test-suite wireform-arrow-test+  import: defaults+  type: exitcode-stdio-1.0+  main-is: Main.hs+  hs-source-dirs: test+  other-modules:+    Test.Arrow.Core+    Test.Arrow.Decimal+    Test.Arrow.Features+    Test.Arrow.Gen+    Test.Arrow.Malformed+    Test.Arrow.Props+    Test.Arrow.Sections+    Test.Arrow.Vector+  build-depends:+    base >= 4.20 && < 5,+    wireform-arrow,+    wireform-columnar-core,+    wireform-flatbuffers,+    bytestring >= 0.11 && < 0.13,+    containers >= 0.6 && < 0.9,+    deepseq >= 1.4 && < 1.6,+    directory,+    hedgehog >= 1.2 && < 1.8,+    text >= 2.0 && < 2.2,+    vector >= 0.13 && < 0.14,+    primitive >= 0.7 && < 0.10++test-suite wireform-arrow-derive-test+  import: defaults+  type: exitcode-stdio-1.0+  main-is: Main.hs+  hs-source-dirs: test-derive+  other-modules:+    Test.Arrow.Derive+    Test.Arrow.Derive.Types+    Test.Arrow.Derive.Instances+  default-extensions:+    ScopedTypeVariables+    TemplateHaskell+  build-depends:+    base,+    wireform-arrow,+    wireform-derive,+    template-haskell,+    text,+    vector,+    sydtest++benchmark wireform-arrow-bench+  import: defaults+  type: exitcode-stdio-1.0+  hs-source-dirs: bench+  main-is: Bench.hs+  -- A 64 MB nursery: the materializing and typed rows are otherwise+  -- dominated by minor GCs (see README, Performance).+  ghc-options: -O2 -rtsopts "-with-rtsopts=-A64m"+  build-depends:+    base,+    wireform-arrow,+    bytestring,+    deepseq,+    text,+    vector,+    criterion >= 1.5 && < 1.7++-- Bidirectional pyarrow interop (wireform-arrow writes / pyarrow reads,+-- pyarrow writes / wireform-arrow reads). Needs a Python with pyarrow:+-- WIREFORM_ARROW_PYTHON (default python3). Prints SKIP and passes when+-- pyarrow is missing unless WIREFORM_ARROW_REQUIRE_PYARROW=1.+test-suite wireform-arrow-pyarrow-interop+  import: defaults+  type: exitcode-stdio-1.0+  main-is: Main.hs+  hs-source-dirs: test-interop+  build-depends:+    base >= 4.20 && < 5,+    wireform-arrow,+    bytestring >= 0.11 && < 0.13,+    containers >= 0.6 && < 0.9,+    deepseq >= 1.4 && < 1.6,+    directory >= 1.3 && < 1.4,+    filepath >= 1.4 && < 1.6,+    primitive >= 0.7 && < 0.10,+    process >= 1.6 && < 1.8,+    temporary >= 1.3 && < 1.4,+    text >= 2.0 && < 2.2,+    vector >= 0.13 && < 0.14+  if flag(zstd)+    cpp-options: -DHAVE_ZSTD+  if flag(lz4)+    cpp-options: -DHAVE_LZ4