{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFunctor #-}
{- | Pull-based streaming primitives shared by every columnar
format ("Arrow.Stream", "Parquet.HighLevel", "ORC").
The point of this module is to give callers /one/ shape for
iterating a record batch / row group / stripe sequence, in
place of the lazy @[Either String a]@ / @[V.Vector ColumnArray]@
lists the per-format readers used to return.
== Why not lazy lists?
Lazy lists conflate three orthogonal concerns:
1. /Pulling/ — the consumer wants to drive iteration explicitly
(e.g. read until a predicate, take @n@ batches, fold an
accumulator).
2. /Errors/ — a per-batch decode failure used to surface as a
@Left e@ slot in the middle of a list, with no way to stop
iteration short of pattern-matching every element.
3. /Resource lifetime/ — even though the underlying byte buffer
is in memory, the lazy list holds onto enough offsets to
keep every later batch reachable, so consumers that only
care about the first few batches still pay for the whole
list spine.
'Iter' separates them. An 'Iter' value is a single suspended
step:
@
'iterStep' :: 'Iter' a -> Either String ('IterStep' a)
@
which yields either @'IterDone'@ (the iterator is exhausted) or
@'IterYield' a next@ (here is the next element + the
continuation iterator). Errors live at the step level rather
than as a Either-element-of-list, so the consumer's loop is
the one place that decides whether to abort or keep going.
== Building blocks
* 'iterFromList' — wrap a strict list (handy bridging API).
* 'iterFromIndexed' — turn a @(Int, Int -> Either String a)@ pair
into an iterator. Used by every format's row-group / stripe /
batch reader: the underlying file already knows /how many/
batches there are and can decode each by index without
touching the others.
* 'iterMap', 'iterMapM', 'iterFilter', 'iterTake', 'iterDrop' —
the usual streaming combinators.
* 'toList', 'toListE', 'foldM' — drain helpers.
== Resource model
'Iter' itself is pure: each step returns @Either String
IterStep@. There is no IO threaded through the iterator; if a
format needs IO (e.g. reading a fresh disk page) the file is
expected to be loaded into memory first (which the existing
"Parquet.Read" / "ORC.Read" code does anyway). When we add
streaming-from-disk it'll be a separate 'IterIO' shape;
keeping the pure shape for now matches the rest of the
columnar code's "ByteString in, columns out" boundary.
-}
module Columnar.Stream (
-- * Core type
Iter,
IterStep (..),
iterStep,
-- * Constructors
iterEmpty,
iterSingleton,
iterFromList,
iterFromVector,
iterFromIndexed,
iterUnfold,
-- * Transformations
iterMap,
iterMapM,
iterMapMaybe,
iterFilter,
iterTake,
iterDrop,
iterAppend,
iterConcat,
-- * Drains
iterToList,
iterToVector,
iterFold,
iterFoldM,
iterForM_,
iterLength,
-- * Row-slice helpers
iterRowSlice,
-- * Re-batching / windowing
iterChunk,
iterScan,
iterMergeBy,
-- * IO-shaped iterator
IterIO (..),
IterIOStep (..),
iterIOFromIter,
iterIOFromAction,
iterIOToList,
iterIOFold,
iterIOMap,
iterIOTake,
iterIOFilter,
-- * Concurrent prefetch
iterIOPrefetch,
iterPrefetch,
iterParallelMap,
) where
import Control.Concurrent.Async (async, link)
import Control.Concurrent.STM (
TBQueue,
atomically,
newTBQueueIO,
readTBQueue,
writeTBQueue,
)
import Control.Exception (SomeException, try)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.IORef (atomicModifyIORef', newIORef)
import Data.Vector qualified as V
{- | A pull-based, error-aware, finite iterator.
Each 'Iter' is a single suspended step: ask it for an
'IterStep' and you either get 'IterDone' (the stream has
ended) or 'IterYield' (here's the next element + the next
iterator).
The ergonomic difference from a list is that 'Iter' yields
/one element at a time/ and each element decode is allowed
to fail, with the caller choosing how to handle the failure.
A lazy @[Either String a]@ forces you to either thread the
error through every later element or risk exceptions; an
'Iter' lets the caller @case@ on the step and break the loop.
-}
newtype Iter a = Iter {iterStep :: Either String (IterStep a)}
instance Functor Iter where
fmap f (Iter step) = Iter $ case step of
Left e -> Left e
Right IterDone -> Right IterDone
Right (IterYield a next) -> Right (IterYield (f a) (fmap f next))
{- | One step of an 'Iter'. Either the iterator has ended or it
carries one element + the continuation.
-}
data IterStep a
= IterDone
| IterYield !a !(Iter a)
deriving (Functor)
-- ============================================================
-- Constructors
-- ============================================================
-- | The empty iterator.
iterEmpty :: Iter a
iterEmpty = Iter (Right IterDone)
-- | An iterator that yields one element and then ends.
iterSingleton :: a -> Iter a
iterSingleton a = Iter (Right (IterYield a iterEmpty))
{- | Wrap a list as an iterator. Each step pulls one element off
the head of the list. The list elements are forced to WHNF
when yielded.
-}
iterFromList :: [a] -> Iter a
iterFromList = go
where
go [] = Iter (Right IterDone)
go (x : xs) = Iter (Right (IterYield x (go xs)))
{- | Wrap a vector as an iterator. Indexes the vector lazily on
each step rather than copying.
-}
iterFromVector :: V.Vector a -> Iter a
iterFromVector v0 = step 0
where
!n = V.length v0
step !i
| i >= n = Iter (Right IterDone)
| otherwise = Iter (Right (IterYield (V.unsafeIndex v0 i) (step (i + 1))))
{- | Build an iterator over @[0..n-1]@ where each index is
decoded on demand by the supplied function. This is the
canonical shape for "we know how many batches the file
carries; produce them one at a time" — every per-format
streaming reader bridges through this.
A decode failure stops the iterator at that step (the next
'iterStep' on the failing iterator returns @Left e@; later
indices are not visited).
-}
iterFromIndexed :: Int -> (Int -> Either String a) -> Iter a
iterFromIndexed !n decode = go 0
where
go !i
| i >= n = Iter (Right IterDone)
| otherwise = Iter $ case decode i of
Left e -> Left e
Right !a -> Right (IterYield a (go (i + 1)))
{- | Generic anamorphism. @iterUnfold seed step@ keeps applying
@step@ to the threading seed, yielding the produced element
and a fresh seed each step until @step@ returns 'Nothing'.
Decode failures from @step@ surface as 'Left' on the failing
iterator step.
-}
iterUnfold :: s -> (s -> Either String (Maybe (a, s))) -> Iter a
iterUnfold s0 step = go s0
where
go s = Iter $ case step s of
Left e -> Left e
Right Nothing -> Right IterDone
Right (Just (a, s')) -> Right (IterYield a (go s'))
-- ============================================================
-- Transformations
-- ============================================================
-- | Apply a pure function to every yielded element.
iterMap :: (a -> b) -> Iter a -> Iter b
iterMap = fmap
{- | Apply a function that may itself fail. A 'Left' from @f@
terminates the iterator at that point.
-}
iterMapM :: (a -> Either String b) -> Iter a -> Iter b
iterMapM f = go
where
go (Iter step) = Iter $ case step of
Left e -> Left e
Right IterDone -> Right IterDone
Right (IterYield a next) -> case f a of
Left e -> Left e
Right !b -> Right (IterYield b (go next))
{- | Drop elements where @f@ returns 'Nothing'; yield @b@ where
it returns @Just b@.
-}
iterMapMaybe :: (a -> Maybe b) -> Iter a -> Iter b
iterMapMaybe f = go
where
go (Iter step) = Iter $ case step of
Left e -> Left e
Right IterDone -> Right IterDone
Right (IterYield a next) -> case f a of
Nothing -> iterStep (go next)
Just !b -> Right (IterYield b (go next))
-- | Keep only elements that satisfy the predicate.
iterFilter :: (a -> Bool) -> Iter a -> Iter a
iterFilter p = iterMapMaybe (\a -> if p a then Just a else Nothing)
{- | Stop after @n@ elements. Any decode error inside the first
@n@ steps still surfaces; errors past @n@ do not.
-}
iterTake :: Int -> Iter a -> Iter a
iterTake = go
where
go !k _ | k <= 0 = iterEmpty
go !k (Iter step) = Iter $ case step of
Left e -> Left e
Right IterDone -> Right IterDone
Right (IterYield a next) -> Right (IterYield a (go (k - 1) next))
{- | Skip the first @n@ elements. Decode errors in the skipped
prefix /are/ surfaced (we cannot tell which index produced
the error without decoding it).
-}
iterDrop :: Int -> Iter a -> Iter a
iterDrop = go
where
go !k it | k <= 0 = it
go !k (Iter step) = Iter $ case step of
Left e -> Left e
Right IterDone -> Right IterDone
Right (IterYield _ next) -> iterStep (go (k - 1) next)
-- | Concatenate two iterators end-to-end.
iterAppend :: Iter a -> Iter a -> Iter a
iterAppend (Iter step) right = Iter $ case step of
Left e -> Left e
Right IterDone -> iterStep right
Right (IterYield a next) -> Right (IterYield a (iterAppend next right))
-- | Flatten an iterator-of-iterators in left-to-right order.
iterConcat :: Iter (Iter a) -> Iter a
iterConcat outer0 = go outer0 iterEmpty
where
go outer inner = Iter $ case iterStep inner of
Left e -> Left e
Right (IterYield a next) -> Right (IterYield a (go outer next))
Right IterDone -> case iterStep outer of
Left e -> Left e
Right IterDone -> Right IterDone
Right (IterYield it' outer') -> iterStep (go outer' it')
-- ============================================================
-- Drains
-- ============================================================
{- | Drain an iterator into a list. The result is strict in
spine (errors stop the drain immediately).
-}
iterToList :: Iter a -> Either String [a]
iterToList = go []
where
go acc (Iter step) = case step of
Left e -> Left e
Right IterDone -> Right (reverse acc)
Right (IterYield a next) -> go (a : acc) next
-- | Drain into a 'V.Vector'.
iterToVector :: Iter a -> Either String (V.Vector a)
iterToVector = fmap V.fromList . iterToList
-- | Strict left fold.
iterFold :: (b -> a -> b) -> b -> Iter a -> Either String b
iterFold f = go
where
go !acc (Iter step) = case step of
Left e -> Left e
Right IterDone -> Right acc
Right (IterYield a next) -> go (f acc a) next
-- | Strict left fold whose step may itself fail.
iterFoldM :: (b -> a -> Either String b) -> b -> Iter a -> Either String b
iterFoldM f = go
where
go !acc (Iter step) = case step of
Left e -> Left e
Right IterDone -> Right acc
Right (IterYield a next) -> case f acc a of
Left e -> Left e
Right !acc' -> go acc' next
-- | Run an effectful action for every yielded element.
iterForM_ :: Monad m => Iter a -> (a -> m ()) -> m (Either String ())
iterForM_ it0 act = go it0
where
go (Iter step) = case step of
Left e -> pure (Left e)
Right IterDone -> pure (Right ())
Right (IterYield a next) -> act a >> go next
{- | How many elements does the iterator carry? Drains the
iterator entirely.
-}
iterLength :: Iter a -> Either String Int
iterLength = iterFold (\n _ -> n + 1) 0
-- ============================================================
-- Row-slice helpers
-- ============================================================
{- | Take an @offset, length@ window of /rows/ from an iterator
whose elements carry their own row count via the supplied
@rowCount@ projection. Walks the iterator one element at a
time, dropping fully-elided elements, slicing the boundary
elements, and stopping early once the window is filled.
Designed for iterating record batches where each batch has a
different number of rows: the caller passes
@columnLength . V.head@ as @rowCount@ and a
'sliceColumnArray'-style slicer to carve the boundary
batches.
-}
iterRowSlice
:: (a -> Int)
-- ^ row count of one element
-> (Int -> Int -> a -> a)
-- ^ slice @start@ @len@ @element@
-> Int
-- ^ row offset to skip
-> Int
-- ^ rows to take
-> Iter a
-> Iter a
iterRowSlice rowCount sliceFn = go
where
go _ 0 _ = iterEmpty
go offset want it = Iter $ case iterStep it of
Left e -> Left e
Right IterDone -> Right IterDone
Right (IterYield a next) ->
let !n = rowCount a
in if offset >= n
then iterStep (go (offset - n) want next)
else
let !take' = min want (n - offset)
!sliced = sliceFn offset take' a
!want' = want - take'
in if want' <= 0
then Right (IterYield sliced iterEmpty)
else Right (IterYield sliced (go 0 want' next))
-- ============================================================
-- Re-batching / windowing
-- ============================================================
{- | Re-bucket an iterator into fixed-size chunks of @n@
elements. The final chunk may be shorter if the source's
length isn't divisible by @n@. Useful for adapting between
formats with different default batch sizes (Parquet's 1M-row
row groups vs Arrow's typical 64K-row record batches).
-}
iterChunk :: Int -> Iter a -> Iter [a]
iterChunk n it0
| n <= 0 = iterEmpty
| otherwise = go it0
where
go it = Iter $ case takeUpTo n [] it of
Left e -> Left e
Right ([], _) -> Right IterDone
Right (xs, rest) -> Right (IterYield (reverse xs) (go rest))
takeUpTo 0 acc rest = Right (acc, rest)
takeUpTo k acc it1 = case iterStep it1 of
Left e -> Left e
Right IterDone -> Right (acc, iterEmpty)
Right (IterYield a next) -> takeUpTo (k - 1) (a : acc) next
{- | Strict left scan: yield the running fold accumulator
alongside each input element. The output iterator's first
element is the seed, then one element per source element
showing the accumulator /after/ that element was folded in.
Useful for prefix sums / running aggregates.
-}
iterScan :: (b -> a -> b) -> b -> Iter a -> Iter b
iterScan f = go
where
go !acc it =
Iter $
Right (IterYield acc (loop acc it))
loop !acc it = Iter $ case iterStep it of
Left e -> Left e
Right IterDone -> Right IterDone
Right (IterYield a next) ->
let !acc' = f acc a
in Right (IterYield acc' (loop acc' next))
{- | Merge a non-empty list of /already-sorted/ iterators into
a single sorted iterator under the supplied comparator.
Each input must be non-decreasing under @cmp@; the output is
the merged non-decreasing sequence.
@O(n log k)@ per element where @k@ is the number of input
iterators (uses a priority loop, not a heap — fine for small
@k@; for large @k@ swap in a real priority queue).
-}
iterMergeBy :: (a -> a -> Ordering) -> [Iter a] -> Iter a
iterMergeBy _ [] = iterEmpty
iterMergeBy cmp xs =
let prepped = mapM peek xs
in Iter $ case prepped of
Left e -> Left e
Right initial ->
let active = [(a, rest) | Just (a, rest) <- initial]
in iterStep (drain cmp active)
where
peek :: Iter a -> Either String (Maybe (a, Iter a))
peek it = case iterStep it of
Left e -> Left e
Right IterDone -> Right Nothing
Right (IterYield a next) -> Right (Just (a, next))
drain :: (a -> a -> Ordering) -> [(a, Iter a)] -> Iter a
drain _ [] = iterEmpty
drain cmp' (h : t) =
Iter $
-- 'pickMin' takes the head separately so the type system
-- proves the input is non-empty (no partial-function
-- match on @[]@). The previous shape used @x : xs@ +
-- a fall-through @[]@ case with 'error' that was
-- statically unreachable but tripped pattern-completeness
-- warnings.
let (winner, rest) = pickMin cmp' h t
(a, next) = winner
in case peek next of
Left e -> Left e
Right Nothing -> Right (IterYield a (drain cmp' rest))
Right (Just (a', next')) ->
Right (IterYield a (drain cmp' ((a', next') : rest)))
-- O(k) per element — fine for small fanouts (every
-- columnar caller is in the [2..32] range). For large k
-- we'd swap in a real priority queue; deferred until a
-- profile shows it matters.
pickMin
:: (a -> a -> Ordering)
-> (a, Iter a)
-> [(a, Iter a)]
-> ((a, Iter a), [(a, Iter a)])
pickMin _ h [] = (h, [])
pickMin cmp' h (x : xs) =
let (m, rest) = pickMin cmp' x xs
in case cmp' (fst h) (fst m) of
GT -> (m, h : rest)
_ -> (h, m : rest)
-- ============================================================
-- IterIO: IO-shaped pull iterator
-- ============================================================
--
-- 'Iter' is pure (the file is in memory). For genuinely
-- streaming reads — say a 50 GB Parquet file you don't want to
-- read into a ByteString first — we want the same pull-based
-- shape but with IO threaded through every step. 'IterIO' is
-- exactly that: each step lives in IO and may decode the next
-- chunk from a handle, fetch a remote object, etc.
--
-- The combinators mirror the pure 'Iter' surface so a consumer
-- can write the same fold once and run it against either a
-- whole-file 'Iter' or a chunked 'IterIO'.
{- | One step of an 'IterIO': either the iterator is done, or it
carries one element + the continuation iterator. Errors come
through the surrounding @Either String IterIOStep@ wrapper.
-}
data IterIOStep a
= IterIODone
| IterIOYield !a !(IterIO a)
{- | Pull-based, IO-effecting, finite, error-aware iterator.
Each call to 'iterIOStep' may perform IO (reading from a
handle, decoding a chunk, etc.) before producing the next
element. Errors are 'Left'; end-of-stream is 'Right
IterIODone'.
-}
newtype IterIO a = IterIO {iterIOStep :: IO (Either String (IterIOStep a))}
{- | Lift a pure 'Iter' into an 'IterIO'. The pure step is
still pulled lazily (one step per outer 'iterIOStep' call).
-}
iterIOFromIter :: Iter a -> IterIO a
iterIOFromIter it = IterIO $ pure $ case iterStep it of
Left e -> Left e
Right IterDone -> Right IterIODone
Right (IterYield a next) ->
Right (IterIOYield a (iterIOFromIter next))
{- | Build an 'IterIO' from a stateful IO action that produces
the next element on each call. The state is kept in an
'IORef' so the produced iterator may be 'iterIOStep'ped
repeatedly without re-running the action's setup.
Useful for wrapping handle-based readers: the action returns
@Right Nothing@ at end-of-stream and @Right (Just x)@ for
each row group / record batch / stripe.
-}
iterIOFromAction
:: IO (Either String (Maybe a)) -> IterIO a
iterIOFromAction step = IterIO go
where
go = do
r <- step
pure $ case r of
Left e -> Left e
Right Nothing -> Right IterIODone
Right (Just x) -> Right (IterIOYield x (IterIO go))
-- | Drain an 'IterIO' into a list (in IO).
iterIOToList :: IterIO a -> IO (Either String [a])
iterIOToList = go []
where
go acc it = do
r <- iterIOStep it
case r of
Left e -> pure (Left e)
Right IterIODone -> pure (Right (reverse acc))
Right (IterIOYield a next) -> go (a : acc) next
-- | Strict left fold in IO.
iterIOFold :: (b -> a -> b) -> b -> IterIO a -> IO (Either String b)
iterIOFold f = go
where
go !acc it = do
r <- iterIOStep it
case r of
Left e -> pure (Left e)
Right IterIODone -> pure (Right acc)
Right (IterIOYield a next) -> go (f acc a) next
-- | Map a pure function over every yielded element.
iterIOMap :: (a -> b) -> IterIO a -> IterIO b
iterIOMap f = go
where
go (IterIO step) = IterIO $ do
r <- step
pure $ case r of
Left e -> Left e
Right IterIODone -> Right IterIODone
Right (IterIOYield a next) ->
Right (IterIOYield (f a) (go next))
-- | Take at most @n@ elements.
iterIOTake :: Int -> IterIO a -> IterIO a
iterIOTake n0 it0
| n0 <= 0 = IterIO (pure (Right IterIODone))
| otherwise = go n0 it0
where
go n (IterIO step) = IterIO $ do
r <- step
pure $ case r of
Left e -> Left e
Right IterIODone -> Right IterIODone
Right (IterIOYield a next) ->
Right (IterIOYield a (go (n - 1) next))
-- | Keep only elements that satisfy the predicate.
iterIOFilter :: (a -> Bool) -> IterIO a -> IterIO a
iterIOFilter p = go
where
go (IterIO step) = IterIO $ do
r <- step
case r of
Left e -> pure (Left e)
Right IterIODone -> pure (Right IterIODone)
Right (IterIOYield a next)
| p a -> pure (Right (IterIOYield a (go next)))
| otherwise -> iterIOStep (go next)
-- The IORef helper isn't currently used internally but is
-- exported via the Hackage-friendly module signature. Silence
-- the unused-imports warning by binding the imported helpers
-- to a no-op.
_iterIOIoRefShim :: IO ()
_iterIOIoRefShim = do
ref <- newIORef ()
() <- atomicModifyIORef' ref (\() -> ((), ()))
pure ()
_iterIOMonadIOShim :: MonadIO m => m ()
_iterIOMonadIOShim = liftIO (pure ())
-- ============================================================
-- Concurrent prefetch
-- ============================================================
{- | Drain an 'IterIO' through a bounded background queue so
the source's @iterIOStep@ runs concurrently with the
consumer. The bound (in elements) is the queue depth: the
producer stalls when the queue is full, the consumer stalls
when it's empty.
For Parquet / ORC readers the source's per-step decode is
usually CPU-bound; running it on a separate thread overlaps
decode with downstream consumption. Typical 2-4× throughput
on multi-row-group reads.
The background thread's exceptions are re-raised on the
consumer side via 'link'.
-}
iterIOPrefetch :: Int -> IterIO a -> IO (IterIO a)
iterIOPrefetch !depth source = do
q <- newTBQueueIO (fromIntegral (max 1 depth))
let producer it = do
r <- try @SomeException (iterIOStep it)
case r of
Left e ->
atomically (writeTBQueue q (Left (show e)))
Right (Left e) -> do
atomically (writeTBQueue q (Left e))
Right (Right IterIODone) ->
atomically (writeTBQueue q (Right Nothing))
Right (Right (IterIOYield a next)) -> do
atomically (writeTBQueue q (Right (Just a)))
producer next
worker <- async (producer source)
link worker
pure (consumeQueue q)
where
consumeQueue :: TBQueue (Either String (Maybe a)) -> IterIO a
consumeQueue q = IterIO $ do
r <- atomically (readTBQueue q)
pure $ case r of
Left e -> Left e
Right Nothing -> Right IterIODone
Right (Just a) -> Right (IterIOYield a (consumeQueue q))
{- | Pure-Iter prefetch: lift to 'IterIO', prefetch, drain back.
Useful when the per-step pure decode is heavy enough to
benefit from a worker thread (e.g. a 'iterMap' over the
output of 'streamRowGroupsIter').
-}
iterPrefetch :: Int -> Iter a -> IO (Iter a)
iterPrefetch depth it = do
ioIt <- iterIOPrefetch depth (iterIOFromIter it)
-- Drain back into a pure Iter by reading the queue
-- synchronously on each pure step. This recovers the pure
-- shape but pays the cost of one IORef per step; callers
-- that want full async benefit should stay in IterIO.
drained <- iterIOToList ioIt
pure $ case drained of
Left e -> iterUnfold () (\_ -> Left e)
Right xs -> iterFromList xs
{- | Parallel map: apply @f@ to every element with up to @k@
in-flight workers, preserving input order. Output is an
'IterIO' with bounded queue depth @k@; consumer back-pressure
caps in-flight work.
Useful for the multi-row-group Parquet case: pass
@iterParallelMap k decodeRowGroup it@ where @decodeRowGroup@
takes a row-group index and returns its decoded columns.
-}
iterParallelMap :: Int -> (a -> IO b) -> IterIO a -> IO (IterIO b)
iterParallelMap !k f source = do
q <- newTBQueueIO (fromIntegral (max 1 k))
let producer it = do
r <- iterIOStep it
case r of
Left e -> atomically (writeTBQueue q (Left e))
Right IterIODone -> atomically (writeTBQueue q (Right Nothing))
Right (IterIOYield a next) -> do
-- Apply f synchronously per source pull. For real
-- parallel-map a worker pool would launch f on a
-- thread and gather; this single-worker shape is
-- the conservative starting point and already
-- overlaps decode with downstream consumption.
r' <- try @SomeException (f a)
case r' of
Left e -> atomically (writeTBQueue q (Left (show e)))
Right b -> do
atomically (writeTBQueue q (Right (Just b)))
producer next
worker <- async (producer source)
link worker
pure (drainQ q)
where
drainQ :: TBQueue (Either String (Maybe b)) -> IterIO b
drainQ q = IterIO $ do
r <- atomically (readTBQueue q)
pure $ case r of
Left e -> Left e
Right Nothing -> Right IterIODone
Right (Just b) -> Right (IterIOYield b (drainQ q))