diff --git a/dataframe-parquet.cabal b/dataframe-parquet.cabal --- a/dataframe-parquet.cabal +++ b/dataframe-parquet.cabal @@ -1,13 +1,13 @@ cabal-version: 2.4 name: dataframe-parquet -version: 1.1.0.0 +version: 1.1.0.1 synopsis: Parquet reader and writer for the dataframe ecosystem. description: @DataFrame.IO.Parquet@ — pure-Haskell Parquet 2.0 reader and writer (with snappy, zstd, gzip codecs, dictionary decoding, nested list/repeated columns, and predicate pushdown). - Heavy package — pulls in @pinch@, @streamly@, and compression codecs. + Heavy package — pulls in @pinch@ and compression codecs. Reading directly from HuggingFace (@hf://@) datasets lives in the separate @dataframe-huggingface@ package. Most users want @dataframe-csv@ instead unless they specifically need Parquet. @@ -31,6 +31,7 @@ library import: warnings + ghc-options: -O2 exposed-modules: DataFrame.IO.Parquet DataFrame.IO.Parquet.Binary @@ -56,8 +57,6 @@ Glob >= 0.10 && < 1, pinch >= 0.5 && < 1, snappy-hs ^>= 0.1, - streamly-bytestring >= 0.2.0 && < 0.4, - streamly-core >= 0.2.3 && < 0.4, text >= 2.1 && < 3, time >= 1.12 && < 2, vector ^>= 0.13, diff --git a/src/DataFrame/IO/Parquet.hs b/src/DataFrame/IO/Parquet.hs --- a/src/DataFrame/IO/Parquet.hs +++ b/src/DataFrame/IO/Parquet.hs @@ -41,7 +41,7 @@ int32Decoder, int64Decoder, int96Decoder, - readPages, + foldColumnPagesM, ) import DataFrame.IO.Parquet.Seeking ( FileBufferedOrSeekable, @@ -81,7 +81,6 @@ import DataFrame.Operations.Merge () import qualified DataFrame.Operations.Subset as DS import qualified Pinch -import qualified Streamly.Data.Stream as Stream import System.Directory (doesDirectoryExist) import System.FilePath (()) import System.FilePath.Glob (glob) @@ -231,11 +230,6 @@ Nothing -> Nothing Just selected -> Just (L.nub (selected ++ predicateColumns)) - -- TODO: When selectedColumnsForRead is Just, pass the set of required - -- column indices into the chunk parsers so that RandomAccess reads are - -- skipped for columns not in the selection, rather than decoding all - -- columns and projecting afterward. - -- TODO: When rowRange is set, compute cumulative row offsets from -- rg_num_rows in each RowGroup and skip any group whose row interval does -- not overlap the requested range, avoiding all decoding for those groups. @@ -278,10 +272,22 @@ Int vectorLength = if topLevelRows > 0 then topLevelRows else rgRows - rawCols <- zipWithM (parseColumnChunks vectorLength) chunks descriptions + -- Column-projection pushdown: decode only the columns needed for the + -- requested output plus any predicate, instead of decoding every column + -- and dropping the unwanted ones afterward. A 'Nothing' selection keeps + -- all columns, so full reads are unchanged. + let keep name = case selectedColumnsForRead of + Nothing -> True + Just req -> last (T.splitOn "." name) `elem` req + kept = filter (\(n, _, _) -> keep n) (zip3 allNames chunks descriptions) + keptNames = [n | (n, _, _) <- kept] + keptChunks = [c | (_, c, _) <- kept] + keptDescs = [d | (_, _, d) <- kept] - let finalCols = zipWith applyDescLogicalType descriptions rawCols - indices = Map.fromList $ zip allNames [0 ..] + rawCols <- zipWithM (parseColumnChunks vectorLength) keptChunks keptDescs + + let finalCols = zipWith applyDescLogicalType keptDescs rawCols + indices = Map.fromList $ zip keptNames [0 ..] dimensions = (vectorLength, length finalCols) let df = @@ -374,9 +380,7 @@ PageDecoder a -> m Column go decoder = - foldNonNullable totalRows $ - (\(vs, _, _) -> vs) - <$> Stream.unfoldEach (readPages description decoder) (Stream.fromList chunks) + foldNonNullable totalRows (foldColumnPagesM description decoder chunks) unboxedGo :: forall a. @@ -384,11 +388,7 @@ UnboxedPageDecoder a -> m Column unboxedGo decoder = - foldNonNullableUnboxed totalRows $ - (\(vs, _, _) -> vs) - <$> Stream.unfoldEach - (readPages description decoder) - (Stream.fromList chunks) + foldNonNullableUnboxed totalRows (foldColumnPagesM description decoder chunks) -- | Decode an optional (nullable) column. getNullableColumn :: @@ -421,20 +421,14 @@ PageDecoder a -> m Column go decoder = - foldNullable maxDef totalRows $ - (\(vs, ds, _) -> (vs, ds)) - <$> Stream.unfoldEach (readPages description decoder) (Stream.fromList chunks) + foldNullable maxDef totalRows (foldColumnPagesM description decoder chunks) unboxedGo :: forall a. (Columnable a, VU.Unbox a) => UnboxedPageDecoder a -> m Column unboxedGo decoder = - foldNullableUnboxed maxDef totalRows $ - (\(vs, ds, _) -> (vs, ds)) - <$> Stream.unfoldEach - (readPages description decoder) - (Stream.fromList chunks) + foldNullableUnboxed maxDef totalRows (foldColumnPagesM description decoder chunks) -- | Decode a repeated (list/nested) column. getRepeatedColumn :: @@ -472,8 +466,7 @@ PageDecoder a -> m Column go decoder = - foldRepeated maxRep maxDef $ - Stream.unfoldEach (readPages description decoder) (Stream.fromList chunks) + foldRepeated maxRep maxDef (foldColumnPagesM description decoder chunks) unboxedGo :: forall a. @@ -486,10 +479,7 @@ UnboxedPageDecoder a -> m Column unboxedGo decoder = - foldRepeatedUnboxed maxRep maxDef $ - Stream.unfoldEach - (readPages description decoder) - (Stream.fromList chunks) + foldRepeatedUnboxed maxRep maxDef (foldColumnPagesM description decoder chunks) -- Options application ----------------------------------------------------- diff --git a/src/DataFrame/IO/Parquet/Page.hs b/src/DataFrame/IO/Parquet/Page.hs --- a/src/DataFrame/IO/Parquet/Page.hs +++ b/src/DataFrame/IO/Parquet/Page.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE OverloadedRecordDot #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -15,7 +16,7 @@ byteArrayDecoder, fixedLenByteArrayDecoder, -- Page iteration - readPages, + foldColumnPagesM, ) where import Control.Monad.IO.Class (MonadIO (liftIO)) @@ -60,7 +61,6 @@ import GHC.Float (castWord32ToFloat, castWord64ToDouble) import Pinch (decodeWithLeftovers) import qualified Pinch -import Streamly.Internal.Data.Unfold (Step (..), Unfold, mkUnfoldM) -- --------------------------------------------------------------------------- -- Types @@ -208,49 +208,51 @@ rawBytes <- readBytes (Range offset compLen) return (codec, pType, rawBytes) -{- | An 'Unfold' from a 'ColumnChunk' to per-page value triples. - -The seed is a 'ColumnChunk'. The inject step reads the chunk's compressed -bytes and discovers the codec and physical type from the column metadata. -Codec and type are then threaded through the unfold state along with the -running dictionary and remaining bytes, so no intermediate list or -concatenation step is needed. Use with 'Stream.unfoldEach' to produce a -flat stream of per-page results directly from a stream of column chunks. +{- | Left-fold a monadic step over every DATA page (V1 or V2) of every column +chunk, in order, threading the running dictionary internally. -Dictionary pages are consumed silently and update the running dictionary -that is threaded through the unfold state. +Replaces the old Streamly @readPages@ 'Unfold' + @unfoldEach@/@fold@ pipeline +with a direct monadic fold: pages are decoded and handed to @step@ one at a +time, so only a single page's decoded values are live at once (the +@unsafeFreeze@-into-preallocated-buffer consumers rely on this — a materialised +page list would roughly double peak memory per column). -The internal state is -@(Maybe DictVals, BS.ByteString, CompressionCodec, ThriftType)@. +Dictionary pages are consumed silently and update the running dictionary that +is threaded through the recursion. @INDEX_PAGE@s are skipped. -- TODO: when a page index is available, use it here to compute which page -- byte ranges to request from the RandomAccess layer instead of reading the -- entire column chunk in one contiguous read. -- TODO: accept an optional row-range and use the column/offset page index --- (when present in file metadata) to Skip pages whose row range does not +-- (when present in file metadata) to skip pages whose row range does not -- overlap the requested range, avoiding decompression of irrelevant pages -- entirely. -} -readPages :: +foldColumnPagesM :: + forall m v a acc. (RandomAccess m, MonadIO m, VG.Vector v a) => ColumnDescription -> (Maybe DictVals -> Encoding -> Int -> BS.ByteString -> v a) -> - Unfold m ColumnChunk (v a, VU.Vector Int, VU.Vector Int) -readPages description decoder = mkUnfoldM step inject + [ColumnChunk] -> + (acc -> (v a, VU.Vector Int, VU.Vector Int) -> m acc) -> + acc -> + m acc +foldColumnPagesM description decoder chunks step = goChunks chunks where maxDef = fromIntegral description.maxDefinitionLevel :: Int maxRep = fromIntegral description.maxRepetitionLevel :: Int - -- Inject: read chunk bytes; put codec and pType into state. - inject cc = do + goChunks [] !acc = return acc + goChunks (cc : ccs) !acc = do (codec, pType, rawBytes) <- readChunkBytes cc - return (Nothing, rawBytes, codec, pType) + acc' <- goPages Nothing codec pType rawBytes acc + goChunks ccs acc' - step (dict, bs, codec, pType) - | BS.null bs = return Stop + goPages dict codec pType bs !acc + | BS.null bs = return acc | otherwise = case parsePageHeader bs of - Left e -> error ("readPages: failed to parse page header: " ++ e) + Left e -> error ("foldColumnPagesM: failed to parse page header: " ++ e) Right (rest, hdr) -> do let compSz = fromIntegral . unField $ hdr.ph_compressed_page_size uncmpSz = fromIntegral . unField $ hdr.ph_uncompressed_page_size @@ -264,7 +266,7 @@ numVals = unField dictHdr.diph_num_values decompressed <- liftIO $ decompressData uncmpSz codec pageData let d = readDictVals pType decompressed numVals description.typeLength - return $ Skip (Just d, rest', codec, pType) + goPages (Just d) codec pType rest' acc DATA_PAGE _ -> do let dph = fromMaybe @@ -276,7 +278,8 @@ let (defLvls, repLvls, nPresent, valBytes) = readLevelsV1 n maxDef maxRep decompressed triple = (decoder dict enc nPresent valBytes, defLvls, repLvls) - return $ Yield triple (dict, rest', codec, pType) + acc' <- step acc triple + goPages dict codec pType rest' acc' DATA_PAGE_V2 _ -> do let dph2 = fromMaybe @@ -296,8 +299,9 @@ then liftIO $ decompressData uncmpSz codec compValBytes else pure compValBytes let triple = (decoder dict enc nPresent valBytes, defLvls, repLvls) - return $ Yield triple (dict, rest', codec, pType) - INDEX_PAGE _ -> return $ Skip (dict, rest', codec, pType) + acc' <- step acc triple + goPages dict codec pType rest' acc' + INDEX_PAGE _ -> goPages dict codec pType rest' acc -- --------------------------------------------------------------------------- -- Page header parsing diff --git a/src/DataFrame/IO/Parquet/Seeking.hs b/src/DataFrame/IO/Parquet/Seeking.hs --- a/src/DataFrame/IO/Parquet/Seeking.hs +++ b/src/DataFrame/IO/Parquet/Seeking.hs @@ -1,6 +1,4 @@ -{- | This module contains low-level utilities around file seeking - -potentially also contains all Streamly related low-level utilities. +{- | This module contains low-level utilities around file seeking. later this module can be renamed / moved to an internal module. -} @@ -9,28 +7,19 @@ SeekMode (..), FileBufferedOrSeekable (..), ForceNonSeekable, - advanceBytes, mkFileBufferedOrSeekable, mkSeekableHandle, readLastBytes, - seekAndReadBytes, - seekAndStreamBytes, withFileBufferedOrSeekable, fSeek, fGet, ) where import Control.Monad -import Control.Monad.IO.Class import qualified Data.ByteString as BS import Data.ByteString.Unsafe (unsafeDrop, unsafeTake) import Data.IORef import Data.Int -import Data.Word -import Streamly.Data.Stream (Stream) -import qualified Streamly.Data.Stream as S -import qualified Streamly.External.ByteString as SBS -import qualified Streamly.FileSystem.Handle as SHandle import System.IO {- | This handle carries a proof that it must be seekable. @@ -94,41 +83,20 @@ fbos <- mkFileBufferedOrSeekable forceNonSeek h action fbos --- | Read from the end, useful for reading metadata without loading entire file +{- | Read the last @n@ bytes, useful for reading metadata without loading the +entire file. Uses 'BS.hGet' (not @hGetContents@) so the handle stays open for +the subsequent column-chunk reads. +-} readLastBytes :: Integer -> FileBufferedOrSeekable -> IO BS.ByteString readLastBytes n (FileSeekable sh) = do let h = getSeekableHandle sh hSeek h SeekFromEnd (negate n) - S.fold SBS.write (SHandle.read h) + BS.hGet h (fromIntegral n) readLastBytes n (FileBuffered i bs) = do writeIORef i (fromIntegral $ BS.length bs) when (n > fromIntegral (BS.length bs)) $ error "lastBytes: n > length bs" pure $ BS.drop (BS.length bs - fromIntegral n) bs --- | Note: this does not guarantee n bytes (if it ends early) -advanceBytes :: Int -> FileBufferedOrSeekable -> IO BS.ByteString -advanceBytes = seekAndReadBytes Nothing - --- | Note: this does not guarantee n bytes (if it ends early) -seekAndReadBytes :: - Maybe (SeekMode, Integer) -> Int -> FileBufferedOrSeekable -> IO BS.ByteString -seekAndReadBytes mSeek len f = seekAndStreamBytes mSeek len f >>= S.fold SBS.write - -{- | Warning: the stream produced from this function accesses to the mutable handler. -if multiple streams are pulled from the same handler at the same time, chaos happen. -Make sure there is only one stream running at one time for each SeekableHandle, -and streams are not read again when they are not used anymore. --} -seekAndStreamBytes :: - (MonadIO m) => - Maybe (SeekMode, Integer) -> Int -> FileBufferedOrSeekable -> m (Stream m Word8) -seekAndStreamBytes mSeek len f = do - liftIO $ - case mSeek of - Nothing -> pure () - Just (seekMode, seekTo) -> fSeek f seekMode seekTo - pure $ S.take len $ fRead f - fSeek :: FileBufferedOrSeekable -> SeekMode -> Integer -> IO () fSeek (FileSeekable (SeekableHandle h)) seekMode seekTo = hSeek h seekMode seekTo fSeek (FileBuffered i _bs) AbsoluteSeek seekTo = writeIORef i (fromIntegral seekTo) @@ -145,15 +113,3 @@ then if i <= BS.length bs then pure $ unsafeDrop i bs else pure BS.empty else pure . unsafeTake n . unsafeDrop i $ bs | otherwise = error "Can't read a negative number of bytes" - -fRead :: (MonadIO m) => FileBufferedOrSeekable -> Stream m Word8 -fRead (FileSeekable (SeekableHandle h)) = SHandle.read h -fRead (FileBuffered i bs) = S.concatEffect $ do - pos <- liftIO $ readIORef i - pure $ - S.mapM - ( \x -> do - liftIO (modifyIORef' i (+ 1)) - pure x - ) - (S.unfold SBS.reader (BS.drop (fromIntegral pos) bs)) diff --git a/src/DataFrame/IO/Parquet/Utils.hs b/src/DataFrame/IO/Parquet/Utils.hs --- a/src/DataFrame/IO/Parquet/Utils.hs +++ b/src/DataFrame/IO/Parquet/Utils.hs @@ -3,6 +3,7 @@ {-# LANGUAGE GADTs #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} module DataFrame.IO.Parquet.Utils ( @@ -40,17 +41,22 @@ ThriftType, unField, ) -import DataFrame.IO.Utils.RandomAccess (RandomAccess) import DataFrame.Internal.Column ( Column (..), Columnable, buildBitmapFromValid, fromList, ) -import qualified Streamly.Data.Fold as Fold -import Streamly.Data.Stream (Stream) -import qualified Streamly.Data.Stream as Stream +{- | A left-fold driver over a column's per-page triples +@(values, def-levels, rep-levels)@, as produced by +'DataFrame.IO.Parquet.Page.foldColumnPagesM'. Generic in the accumulator so a +single driver serves every fold strategy below. +-} +type PageFold m v a = + forall acc. + (acc -> (v a, VU.Vector Int, VU.Vector Int) -> m acc) -> acc -> m acc + data ColumnDescription = ColumnDescription { colElementType :: !(Maybe ThriftType) , maxDefinitionLevel :: !Int32 @@ -155,88 +161,74 @@ childLeaves = go children subPath False in childLeaves ++ go rest path skipThis -{- | Fold a stream of value chunks into a non-nullable 'Column'. - -Pre-allocates a mutable vector of @totalRows@ and fills it chunk-by-chunk -using a single 'Fold.foldlM\'' pass, avoiding any intermediate list or -concatenation allocation. +{- | Fold a column's value pages into a non-nullable 'Column'. -For unboxable element types the chunks (which are always boxed) are -unboxed element-by-element directly into the pre-allocated unboxed -buffer, eliminating the boxing round-trip that a 'fromVector' call on a -boxed concat would otherwise require. +Pre-allocates a mutable vector of @totalRows@ and fills it page-by-page via a +single streaming left fold ('PageFold'), avoiding any intermediate list or +concatenation allocation. Only one page's values are live at a time. -} foldNonNullable :: forall m a. - (RandomAccess m, MonadIO m, Columnable a) => + (MonadIO m, Columnable a) => Int -> - Stream m (VB.Vector a) -> + PageFold m VB.Vector a -> m Column -foldNonNullable totalRows stream = do +foldNonNullable totalRows runPages = do mv <- liftIO $ VBM.unsafeNew totalRows _ <- - Stream.fold - ( Fold.foldlM' - ( \off chunk -> liftIO $ do - let n = VB.length chunk - VB.copy (VBM.unsafeSlice off n mv) chunk - return (off + n) - ) - (return 0) + runPages + ( \off (chunk, _, _) -> liftIO $ do + let n = VB.length chunk + VB.copy (VBM.unsafeSlice off n mv) chunk + return (off + n) ) - stream + (0 :: Int) v <- liftIO $ VB.unsafeFreeze mv return (BoxedColumn Nothing v) foldNonNullableUnboxed :: forall m a. - (RandomAccess m, MonadIO m, Columnable a, VU.Unbox a) => + (MonadIO m, Columnable a, VU.Unbox a) => Int -> - Stream m (VU.Vector a) -> + PageFold m VU.Vector a -> m Column -foldNonNullableUnboxed totalRows stream = do +foldNonNullableUnboxed totalRows runPages = do mv <- liftIO $ VUM.unsafeNew totalRows _ <- - Stream.fold - ( Fold.foldlM' - ( \off chunk -> liftIO $ do - let n = VU.length chunk - go i - | i >= n = return () - | otherwise = do - VUM.unsafeWrite - mv - (off + i) - (VU.unsafeIndex chunk i) - go (i + 1) - go 0 - return (off + n) - ) - (return 0) + runPages + ( \off (chunk, _, _) -> liftIO $ do + let n = VU.length chunk + go i + | i >= n = return () + | otherwise = do + VUM.unsafeWrite + mv + (off + i) + (VU.unsafeIndex chunk i) + go (i + 1) + go 0 + return (off + n) ) - stream + (0 :: Int) dat <- liftIO $ VU.unsafeFreeze mv return (UnboxedColumn Nothing dat) -{- | Fold a stream of (values, def-levels) pairs into a nullable 'Column'. +{- | Fold a column's (values, def-levels) pages into a nullable 'Column'. -Pre-allocates the output buffer and a valid-mask vector of @totalRows@, -then scatters values inline during a single 'Fold.foldlM\'' pass. -This eliminates the @allVals@ intermediate vector that the old -'Stream.toList' + concat approach required. +Pre-allocates the output buffer and a valid-mask vector of @totalRows@, then +scatters values inline during a single streaming left fold ('PageFold'). A 'hasNull' flag is accumulated during the scatter so the -'buildBitmapFromValid' call (and the second 'VU.all' scan) is skipped -entirely when all values are present. +'buildBitmapFromValid' call is skipped entirely when all values are present. -} foldNullable :: forall m a. - (RandomAccess m, MonadIO m, Columnable a) => + (MonadIO m, Columnable a) => Int -> Int -> - Stream m (VB.Vector a, VU.Vector Int) -> + PageFold m VB.Vector a -> m Column -foldNullable maxDef totalRows stream = do +foldNullable maxDef totalRows runPages = do -- null slots hold an error thunk, guarded by bitmap. -- -- IMPORTANT: 'VBM.unsafeWrite' for boxed vectors stores a *pointer* to @@ -248,24 +240,21 @@ liftIO $ VBM.replicate totalRows (error "parquet: null slot accessed") mvValid <- liftIO (VUM.new totalRows :: IO (VUM.IOVector Word8)) (_, hasNull) <- - Stream.fold - ( Fold.foldlM' - ( \(rowOff, anyNull) (vals, defs) -> liftIO $ do - let nDefs = VU.length defs - go i j acc - | i >= nDefs = return acc - | VU.unsafeIndex defs i == maxDef = do - let !v = VB.unsafeIndex vals j - VBM.unsafeWrite mvDat (rowOff + i) v - VUM.unsafeWrite mvValid (rowOff + i) 1 - go (i + 1) (j + 1) acc - | otherwise = go (i + 1) j True - newNull <- go 0 0 False - return (rowOff + nDefs, anyNull || newNull) - ) - (return (0, False)) + runPages + ( \(rowOff, anyNull) (vals, defs, _) -> liftIO $ do + let nDefs = VU.length defs + go i j acc + | i >= nDefs = return acc + | VU.unsafeIndex defs i == maxDef = do + let !v = VB.unsafeIndex vals j + VBM.unsafeWrite mvDat (rowOff + i) v + VUM.unsafeWrite mvValid (rowOff + i) 1 + go (i + 1) (j + 1) acc + | otherwise = go (i + 1) j True + newNull <- go 0 0 False + return (rowOff + nDefs, anyNull || newNull) ) - stream + (0 :: Int, False) dat <- liftIO $ VB.unsafeFreeze mvDat maybeBm <- if hasNull @@ -277,20 +266,30 @@ foldNullableUnboxed :: forall m a. - (RandomAccess m, MonadIO m, Columnable a, VU.Unbox a) => + (MonadIO m, Columnable a, VU.Unbox a) => Int -> Int -> - Stream m (VU.Vector a, VU.Vector Int) -> + PageFold m VU.Vector a -> m Column -foldNullableUnboxed maxDef totalRows stream = do +foldNullableUnboxed maxDef totalRows runPages = do -- zero-init means null slots silently hold 0, guarded by bitmap. mvDat <- liftIO $ VUM.new totalRows mvValid <- liftIO (VUM.new totalRows :: IO (VUM.IOVector Word8)) - -- Drain the stream into a list once, then run a tight IO loop. This - -- avoids per-page Streamly polymorphic-monad dispatch in the inner - -- scatter loop. - chunks <- Stream.toList stream - hasNull <- liftIO $ scatterChunks mvDat mvValid maxDef chunks + (_, hasNull) <- + runPages + ( \(rowOff, anyNull) (vals, defs, _) -> liftIO $ do + let !nDefs = VU.length defs + go !i !j !acc + | i >= nDefs = pure acc + | VU.unsafeIndex defs i == maxDef = do + VUM.unsafeWrite mvDat (rowOff + i) (VU.unsafeIndex vals j) + VUM.unsafeWrite mvValid (rowOff + i) 1 + go (i + 1) (j + 1) acc + | otherwise = go (i + 1) j True + !newNull <- go 0 0 False + return (rowOff + nDefs, anyNull || newNull) + ) + (0 :: Int, False) dat <- liftIO $ VU.unsafeFreeze mvDat maybeBm <- if hasNull @@ -299,28 +298,6 @@ return (Just (buildBitmapFromValid validV)) else return Nothing return (UnboxedColumn maybeBm dat) - where - scatterChunks :: - VUM.IOVector a -> - VUM.IOVector Word8 -> - Int -> - [(VU.Vector a, VU.Vector Int)] -> - IO Bool - scatterChunks mvDat mvValid !md = goChunks 0 False - where - goChunks !_ !anyNull [] = pure anyNull - goChunks !rowOff !anyNull ((vals, defs) : rest) = do - let !nDefs = VU.length defs - go !i !j !acc - | i >= nDefs = pure acc - | VU.unsafeIndex defs i == md = do - VUM.unsafeWrite mvDat (rowOff + i) (VU.unsafeIndex vals j) - VUM.unsafeWrite mvValid (rowOff + i) 1 - go (i + 1) (j + 1) acc - | otherwise = go (i + 1) j True - !newNull <- go 0 0 False - goChunks (rowOff + nDefs) (anyNull || newNull) rest -{-# INLINE foldNullableUnboxed #-} {- | Fold a stream of (values, def-levels, rep-levels) triples into a repeated (list) 'Column' using Dremel-style level stitching. @@ -335,8 +312,7 @@ -} foldRepeated :: forall m a. - ( RandomAccess m - , MonadIO m + ( MonadIO m , Columnable a , Columnable (Maybe [Maybe a]) , Columnable (Maybe [Maybe [Maybe a]]) @@ -344,10 +320,10 @@ ) => Int -> Int -> - Stream m (VB.Vector a, VU.Vector Int, VU.Vector Int) -> + PageFold m VB.Vector a -> m Column -foldRepeated maxRep maxDef stream = do - chunks <- Stream.toList stream +foldRepeated maxRep maxDef runPages = do + chunks <- collectPages runPages let allVals = VB.concat [vs | (vs, _, _) <- chunks] allDefs = VU.concat [ds | (_, ds, _) <- chunks] allReps = VU.concat [rs | (_, _, rs) <- chunks] @@ -359,8 +335,7 @@ foldRepeatedUnboxed :: forall m a. - ( RandomAccess m - , MonadIO m + ( MonadIO m , Columnable a , VU.Unbox a , Columnable (Maybe [Maybe a]) @@ -369,10 +344,10 @@ ) => Int -> Int -> - Stream m (VU.Vector a, VU.Vector Int, VU.Vector Int) -> + PageFold m VU.Vector a -> m Column -foldRepeatedUnboxed maxRep maxDef stream = do - chunks <- Stream.toList stream +foldRepeatedUnboxed maxRep maxDef runPages = do + chunks <- collectPages runPages let allVals = VB.convert $ VU.concat [vs | (vs, _, _) <- chunks] allDefs = VU.concat [ds | (_, ds, _) <- chunks] allReps = VU.concat [rs | (_, _, rs) <- chunks] @@ -381,3 +356,13 @@ 3 -> fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef allReps allDefs allVals) _ -> fromList (stitchList maxDef allReps allDefs allVals) + +{- | Collect all of a column's page triples into a list, in page order. Used by +the repeated/list folds, where Dremel level-stitching needs the full +concatenated rep/def/value arrays at once (so streaming gives no benefit here). +-} +collectPages :: + (Monad m) => + PageFold m v a -> + m [(v a, VU.Vector Int, VU.Vector Int)] +collectPages runPages = reverse <$> runPages (\acc triple -> return (triple : acc)) []