dataframe-parquet 1.5.1.0 → 1.6.0.0
raw patch · 20 files changed
+1319/−1094 lines, 20 filesdep ~dataframe-coredep ~dataframe-operations
Dependency ranges changed: dataframe-core, dataframe-operations
Files
- dataframe-parquet.cabal +7/−6
- src/DataFrame/IO/Parquet.hs +424/−178
- src/DataFrame/IO/Parquet/Binary.hs +2/−9
- src/DataFrame/IO/Parquet/Dictionary.hs +89/−88
- src/DataFrame/IO/Parquet/Encoding.hs +57/−20
- src/DataFrame/IO/Parquet/Levels.hs +10/−34
- src/DataFrame/IO/Parquet/Page.hs +186/−133
- src/DataFrame/IO/Parquet/Schema.hs +5/−18
- src/DataFrame/IO/Parquet/Seeking.hs +6/−36
- src/DataFrame/IO/Parquet/Thrift.hs +37/−45
- src/DataFrame/IO/Parquet/Time.hs +1/−2
- src/DataFrame/IO/Parquet/Utils.hs +7/−50
- src/DataFrame/IO/Parquet/Writer.hs +102/−77
- src/DataFrame/IO/Parquet/Writer/DefLevels.hs +21/−20
- src/DataFrame/IO/Parquet/Writer/Encoder.hs +102/−180
- src/DataFrame/IO/Parquet/Writer/Metadata.hs +3/−2
- src/DataFrame/IO/Utils/RandomAccess.hs +93/−113
- src/DataFrame/Typed/IO/Parquet.hs +6/−52
- stress/DataFrame10GB.hs +2/−2
- tests/Main.hs +159/−29
dataframe-parquet.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.4 name: dataframe-parquet-version: 1.5.1.0+version: 1.6.0.0 synopsis: Parquet reader and writer for the dataframe ecosystem. description: @DataFrame.IO.Parquet@ — pure-Haskell Parquet 2.0 reader and writer@@ -60,8 +60,8 @@ build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.14, containers >= 0.6.7 && < 0.10,- dataframe-core >= 2.5 && < 2.6,- dataframe-operations >= 2.5 && < 2.6,+ dataframe-core >= 2.6 && < 2.7,+ dataframe-operations >= 2.6 && < 2.7, dataframe-parsing >= 2.2 && < 2.3, primitive >= 0.7 && < 0.11, directory >= 1.3.0.0 && < 2,@@ -85,7 +85,8 @@ hs-source-dirs: tests build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.14,- dataframe-core >= 2.5 && < 2.6,+ dataframe-core >= 2.6 && < 2.7,+ dataframe-operations >= 2.6 && < 2.7, dataframe-parquet, directory >= 1.3.0.0 && < 2, filepath >= 1.4 && < 2,@@ -100,7 +101,7 @@ other-modules: DataFrame10GB hs-source-dirs: stress build-depends: base >= 4 && < 5,- dataframe-core >= 2.5 && < 2.6,+ dataframe-core >= 2.6 && < 2.7, dataframe-parquet, filepath >= 1.4 && < 2, temporary >= 1.3 && < 1.5,@@ -119,7 +120,7 @@ build-depends: base >= 4 && < 5, criterion >= 1 && < 2, deepseq >= 1.4 && < 2,- dataframe-core >= 2.5 && < 2.6,+ dataframe-core >= 2.6 && < 2.7, dataframe-parquet, directory >= 1.3 && < 2, filepath >= 1.4 && < 2,
src/DataFrame/IO/Parquet.hs view
@@ -62,8 +62,11 @@ ForceNonSeekable, withFileBufferedOrSeekable, _readParquetWithOpts,+ readParquetWithMetadata,+ FileMetadata, ) where +import Control.Concurrent (getNumCapabilities) import Control.Exception (throw) import Control.Monad import Control.Monad.IO.Class (MonadIO (..))@@ -72,7 +75,9 @@ import qualified Data.ByteString as BS import Data.Either (fromRight) import Data.Functor ((<&>))+import Data.IORef (atomicModifyIORef', newIORef) import Data.Int (Int32, Int64)+import Data.Maybe (fromJust, fromMaybe) #if !MIN_VERSION_base(4,20,0) import Data.List (foldl', transpose) #else@@ -83,9 +88,21 @@ import qualified Data.Text as T import Data.Time.Calendar (Day (ModifiedJulianDay)) import Data.Time.Clock (UTCTime (UTCTime), picosecondsToDiffTime)+import Data.Type.Equality (testEquality, (:~:) (Refl)) import qualified Data.Vector as Vector+import qualified Data.Vector.Mutable as MVector import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM import DataFrame.Errors (DataFrameException (ColumnsNotFoundException))+import DataFrame.IO.Parquet.Dictionary (+ DictVals,+ dictDouble,+ dictFloat,+ dictInt32,+ dictInt64,+ word32At,+ word64At,+ ) import DataFrame.IO.Parquet.Page ( PageDecoder, UnboxedPageDecoder,@@ -93,6 +110,7 @@ appendStringPageIO, boolDecoder, byteArrayDecoder,+ decodeRequiredNumericInto, doubleDecoder, fixedLenByteArrayDecoder, floatDecoder,@@ -109,11 +127,13 @@ ) import DataFrame.IO.Parquet.Thrift ( ColumnChunk (..),+ ColumnMetaData (..), DecimalType (..), FileMetadata (..), KeyValue (..), LogicalType (..), RowGroup (..),+ Statistics (..), ThriftType (..), TimeUnit (..), TimestampType (..),@@ -139,18 +159,26 @@ writeParquetWithOptions, ) import DataFrame.IO.Utils.RandomAccess (+ ChunkBuffer (..), RandomAccess (..),+ Range (Range), ReaderIO (runReaderIO), ) import DataFrame.Internal.Column (Column, Columnable) import qualified DataFrame.Internal.Column as DI import DataFrame.Internal.Column.Builder (+ concatColumnsIO, freezeTextChunk, mergeTextChunks, newTextBuilder, )+import DataFrame.Internal.Control.Concurrent (forkJoin_) import DataFrame.Internal.DataFrame (DataFrame (..))-import DataFrame.Internal.Expression (Expr, getColumns)+import DataFrame.Internal.Expression (+ BinaryOp (binaryName),+ Expr (..),+ getColumns,+ ) import DataFrame.Operations.Merge () import qualified DataFrame.Operations.Subset as DS import DataFrame.Schema (@@ -159,52 +187,32 @@ makeSchema, schemaType, )+import GHC.Float (castWord32ToFloat, castWord64ToDouble) import qualified Pinch import System.Directory (doesDirectoryExist) import System.FilePath ((</>)) import System.FilePath.Glob (glob) import System.IO (IOMode (ReadMode))---- Options -------------------------------------------------------------------{- | Options for reading Parquet data.--These options are applied in this order:--1. predicate filtering-2. column projection-3. row range-4. safe column promotion+import Type.Reflection (typeRep) -Column selection for @selectedColumns@ uses leaf column names only.--}+-- | Options for reading Parquet data. data ParquetReadOptions = ParquetReadOptions { selectedColumns :: Maybe [T.Text]- {- ^ Columns to keep in the final dataframe. If set, only these columns are returned.- Predicate-referenced columns are read automatically when needed and projected out after filtering.- -}+ -- ^ Columns to keep in the final dataframe. , predicate :: Maybe (Expr Bool) -- ^ Optional row filter expression applied before projection. , rowRange :: Maybe (Int, Int)- -- ^ Optional row slice @(start, end)@ with start-inclusive/end-exclusive semantics.+ -- ^ Optional row slice @(start, end)@, start-inclusive, end-exclusive. , safeColumns :: Bool- -- ^ When True, every column is promoted to OptionalColumn after read, regardless of nullability in the schema.+ -- ^ Promote every column to OptionalColumn after read.+ , widenIntegers :: Bool+ -- ^ Decode INT32 and INT64 columns as 'Int'.+ , parallel :: Bool+ -- ^ Decode row groups concurrently using every RTS capability. } deriving (Show) -{- | Default Parquet read options.--Equivalent to:--@-ParquetReadOptions- { selectedColumns = Nothing- , predicate = Nothing- , rowRange = Nothing- , safeColumns = False- }-@--}+-- | Default Parquet read options. defaultParquetReadOptions :: ParquetReadOptions defaultParquetReadOptions = ParquetReadOptions@@ -212,63 +220,287 @@ , predicate = Nothing , rowRange = Nothing , safeColumns = False+ , widenIntegers = False+ , parallel = True } --- Public API ----------------------------------------------------------------{- | Read a parquet file from path and load it into a dataframe.--==== __Example__-@-ghci> D.readParquet ".\/data\/mtcars.parquet"-@--}+-- | Read a parquet file from path and load it into a dataframe. readParquet :: FilePath -> IO DataFrame readParquet = readParquetWithOpts defaultParquetReadOptions -{- | Read a Parquet file using explicit read options.--==== __Example__-@-ghci> D.readParquetWithOpts-ghci| (D.defaultParquetReadOptions{D.selectedColumns = Just ["id"], D.rowRange = Just (0, 10)})-ghci| "./tests/data/alltypes_plain.parquet"-@--When @selectedColumns@ is set and @predicate@ references other columns, those predicate columns-are auto-included for decoding, then projected back to the requested output columns.--}+-- | Read a Parquet file using explicit read options. readParquetWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame readParquetWithOpts = _readParquetWithOpts Nothing -- | Internal entry point used by tests to force non-seekable mode. _readParquetWithOpts :: ForceNonSeekable -> ParquetReadOptions -> FilePath -> IO DataFrame-_readParquetWithOpts extraConfig opts path =- withFileBufferedOrSeekable extraConfig path ReadMode $ \file ->- runReaderIO (parseParquetWithOpts opts) file+_readParquetWithOpts extraConfig opts path = do+ metadata <-+ withFileBufferedOrSeekable+ extraConfig+ path+ ReadMode+ (runReaderIO parseFileMetadata)+ readParquetWithMetadata extraConfig metadata opts path -{- | Read Parquet files from a directory or glob path.+-- | Read a Parquet file whose footer has already been parsed.+readParquetWithMetadata ::+ ForceNonSeekable ->+ FileMetadata ->+ ParquetReadOptions ->+ FilePath ->+ IO DataFrame+readParquetWithMetadata extraConfig metadata opts path = do+ caps <- getNumCapabilities+ let nRowGroups = length (unField metadata.row_groups)+ if parallel opts && extraConfig /= Just True && caps > 1 && nRowGroups > 1+ then decodeRowGroupsParallel caps opts path metadata+ else withFileBufferedOrSeekable extraConfig path ReadMode $ \file ->+ runReaderIO (parseParquetWithOpts opts) file -This is equivalent to calling 'readParquetFilesWithOpts' with 'defaultParquetReadOptions'.--}-readParquetFiles :: FilePath -> IO DataFrame-readParquetFiles = readParquetFilesWithOpts defaultParquetReadOptions+-- | Decode row groups in parallel, one task per (row group, column).+decodeRowGroupsParallel ::+ Int -> ParquetReadOptions -> FilePath -> FileMetadata -> IO DataFrame+decodeRowGroupsParallel caps opts0 path metadata = do+ let plan = planColumns opts0 metadata+ descs = Vector.fromList (planDescs plan)+ allRowGroupRows =+ map (fromIntegral . unField . rg_num_rows) (unField metadata.row_groups)+ allOffsets = scanl (+) 0 allRowGroupRows+ -- Row groups whose statistics rule out the predicate are skipped.+ statsGroups =+ [ g+ | (g, rg) <- zip [0 ..] (unField metadata.row_groups)+ , maybe True (rowGroupMayMatch metadata rg) (predicate opts0)+ ]+ -- A row range only touches overlapping groups, then is rebased.+ (keptGroups, opts) = case rowRange opts0 of+ Nothing -> (statsGroups, opts0)+ Just (s, e) ->+ let groups =+ [ g+ | (g, off, n) <- zip3 [0 ..] allOffsets allRowGroupRows+ , off < e+ , off + n > s+ ]+ base = case groups of+ [] -> 0+ (g : _) -> allOffsets !! g+ in (groups, opts0{rowRange = Just (s - base, e - base)})+ chunks =+ Vector.fromList+ [Vector.fromList [cs !! g | g <- keptGroups] | cs <- planChunks plan]+ rowGroupRows = Vector.fromList [allRowGroupRows !! g | g <- keptGroups]+ rowOffsets = Vector.prescanl' (+) 0 rowGroupRows+ totalRows = Vector.sum rowGroupRows+ nCols = Vector.length descs+ nRg = Vector.length rowGroupRows+ nTasks = nCols * nRg+ taskChunk t = let (r, c) = t `divMod` nCols in chunks Vector.! c Vector.! r+ sinks <- Vector.mapM (newSink (widenIntegers opts) totalRows) descs+ results <- MVector.new nTasks+ counter <- newIORef (0 :: Int)+ -- Each worker owns a handle, so chunk reads run in parallel too.+ let worker = withFileBufferedOrSeekable Nothing path ReadMode $ \file ->+ let loop = do+ t <- atomicModifyIORef' counter (\n -> (n + 1, n))+ when (t < nTasks) $ do+ let (r, c) = t `divMod` nCols+ desc = descs Vector.! c+ chunk = taskChunk t+ (base, len) = chunkRange chunk+ bytes <- runReaderIO (readBytes (Range base len)) file+ let buffer = ChunkBuffer base bytes+ into ::+ forall a.+ (VU.Unbox a) =>+ (DictVals -> Maybe (VU.Vector a)) ->+ (BS.ByteString -> Int -> a) ->+ Int ->+ VUM.IOVector a ->+ IO ()+ into pick plainAt width out =+ void $+ runReaderIO+ ( decodeRequiredNumericInto+ pick+ plainAt+ width+ desc+ chunk+ out+ (rowOffsets Vector.! r)+ )+ buffer+ case sinks Vector.! c of+ SinkInt out width ->+ if width == 4+ then into (fmap (VU.map fromIntegral) . dictInt32) plainIntFrom32 4 out+ else into (fmap (VU.map fromIntegral) . dictInt64) plainIntFrom64 8 out+ SinkInt32 out -> into dictInt32 plainInt32 4 out+ SinkInt64 out -> into dictInt64 plainInt64 8 out+ SinkFloat out -> into dictFloat plainFloat 4 out+ SinkDouble out -> into dictDouble plainDouble 8 out+ SinkRows -> do+ col <-+ runReaderIO+ (parseColumnChunks (rowGroupRows Vector.! r) [chunk] desc)+ buffer+ MVector.write results (c * nRg + r) col+ loop+ in loop+ forkJoin_ (replicate (max 1 (min caps nTasks)) worker)+ cols <- forM [0 .. nCols - 1] $ \c -> case sinks Vector.! c of+ SinkInt out _ -> DI.UnboxedColumn Nothing <$> VU.unsafeFreeze out+ SinkInt32 out -> DI.UnboxedColumn Nothing <$> VU.unsafeFreeze out+ SinkInt64 out -> DI.UnboxedColumn Nothing <$> VU.unsafeFreeze out+ SinkFloat out -> DI.UnboxedColumn Nothing <$> VU.unsafeFreeze out+ SinkDouble out -> DI.UnboxedColumn Nothing <$> VU.unsafeFreeze out+ SinkRows ->+ concatColumnsIO+ =<< forM [0 .. nRg - 1] (\r -> MVector.read results (c * nRg + r))+ pure (assembleDataFrame opts totalRows metadata plan cols) -{- | Read multiple Parquet files (directory or glob) using explicit options.+-- | Where a column's decoded values go.+data ColumnSink+ = -- | Widened integers; the width is the source's bytes per value.+ SinkInt !(VUM.IOVector Int) !Int+ | SinkInt32 !(VUM.IOVector Int32)+ | SinkInt64 !(VUM.IOVector Int64)+ | SinkFloat !(VUM.IOVector Float)+ | SinkDouble !(VUM.IOVector Double)+ | SinkRows -If @path@ is a directory, all non-directory entries are read.-If @path@ is a glob, matching files are read.+newSink :: Bool -> Int -> ColumnDescription -> IO ColumnSink+newSink widen totalRows description+ | description.maxRepetitionLevel /= 0 || description.maxDefinitionLevel /= 0 =+ pure SinkRows+ | otherwise = case description.colElementType of+ Just (INT32 _)+ | widen -> (`SinkInt` 4) <$> VUM.unsafeNew totalRows+ | otherwise -> SinkInt32 <$> VUM.unsafeNew totalRows+ Just (INT64 _)+ | widen -> (`SinkInt` 8) <$> VUM.unsafeNew totalRows+ | otherwise -> SinkInt64 <$> VUM.unsafeNew totalRows+ Just (FLOAT _) -> SinkFloat <$> VUM.unsafeNew totalRows+ Just (DOUBLE _) -> SinkDouble <$> VUM.unsafeNew totalRows+ _ -> pure SinkRows -For multi-file reads, @rowRange@ is applied once after concatenation (global range semantics).+plainIntFrom32 :: BS.ByteString -> Int -> Int+plainIntFrom32 bs o = fromIntegral (fromIntegral (word32At bs o) :: Int32)+{-# INLINE plainIntFrom32 #-} -==== __Example__-@-ghci> D.readParquetFilesWithOpts-ghci| (D.defaultParquetReadOptions{D.selectedColumns = Just ["id"], D.rowRange = Just (0, 5)})-ghci| "./tests/data/alltypes_plain*.parquet"-@--}+plainIntFrom64 :: BS.ByteString -> Int -> Int+plainIntFrom64 bs o = fromIntegral (word64At bs o)+{-# INLINE plainIntFrom64 #-}++plainInt32 :: BS.ByteString -> Int -> Int32+plainInt32 bs o = fromIntegral (word32At bs o)+{-# INLINE plainInt32 #-}++plainInt64 :: BS.ByteString -> Int -> Int64+plainInt64 bs o = fromIntegral (word64At bs o)+{-# INLINE plainInt64 #-}++plainFloat :: BS.ByteString -> Int -> Float+plainFloat bs o = castWord32ToFloat (word32At bs o)+{-# INLINE plainFloat #-}++plainDouble :: BS.ByteString -> Int -> Double+plainDouble bs o = castWord64ToDouble (word64At bs o)+{-# INLINE plainDouble #-}++-- | Whether a row group can hold a row satisfying the predicate.+rowGroupMayMatch :: FileMetadata -> RowGroup -> Expr Bool -> Bool+rowGroupMayMatch metadata rg = go+ where+ allNames = getColumnNames (drop 1 (unField metadata.schema))+ chunks = unField (rg_columns rg)+ go :: Expr Bool -> Bool+ go (Binary op l r)+ | binaryName op == "and" = case (l, r) of+ (l' :: Expr lt, r' :: Expr rt)+ | Just Refl <- testEquality (typeRep @lt) (typeRep @Bool)+ , Just Refl <- testEquality (typeRep @rt) (typeRep @Bool) ->+ go l' && go r'+ _ -> True+ | otherwise = case (l, r) of+ (Col name :: Expr lt, Lit (v :: rt))+ | Just Refl <- testEquality (typeRep @lt) (typeRep @Int)+ , Just Refl <- testEquality (typeRep @rt) (typeRep @Int) ->+ decide name (binaryName op) v+ (Lit (v :: lt), Col name :: Expr rt)+ | Just Refl <- testEquality (typeRep @lt) (typeRep @Int)+ , Just Refl <- testEquality (typeRep @rt) (typeRep @Int) ->+ decide name (flipped (binaryName op)) v+ _ -> True+ go _ = True+ flipped :: T.Text -> T.Text+ flipped n = case n of+ "lt" -> "gt"+ "gt" -> "lt"+ "leq" -> "geq"+ "geq" -> "leq"+ other -> other+ -- @column <cmp> v@ can hold for some row in [lo, hi].+ decide :: T.Text -> T.Text -> Int -> Bool+ decide name cmp v = case columnRange name of+ Nothing -> True+ Just (lo, hi) -> case cmp of+ "eq" -> lo <= v && v <= hi+ "neq" -> not (lo == v && hi == v)+ "lt" -> lo < v+ "leq" -> lo <= v+ "gt" -> hi > v+ "geq" -> hi >= v+ _ -> True+ columnRange :: T.Text -> Maybe (Int, Int)+ columnRange name = do+ i <- L.findIndex (columnMatches allNames name) allNames+ cc <- chunks `atMay` i+ meta <- unField (cc_meta_data cc)+ stats <- unField (cmd_statistics meta)+ let decodeInt bs = case unField (cmd_type meta) of+ INT32 _+ | BS.length bs == 4 ->+ Just (fromIntegral (fromIntegral (word32At bs 0) :: Int32))+ INT64 _ | BS.length bs == 8 -> Just (fromIntegral (word64At bs 0))+ _ -> Nothing+ lo <-+ decodeInt+ =<< (unField (stats_min_value stats) `orElse` unField (stats_min stats))+ hi <-+ decodeInt+ =<< (unField (stats_max_value stats) `orElse` unField (stats_max stats))+ pure (lo, hi)+ orElse (Just x) _ = Just x+ orElse Nothing y = y+ atMay xs i = if i >= 0 && i < length xs then Just (xs !! i) else Nothing++-- | File offset and byte length of a column chunk (dictionary page included).+chunkRange :: ColumnChunk -> (Integer, Int)+chunkRange columnChunk =+ let meta = fromJust . unField $ columnChunk.cc_meta_data+ dataOffset = fromIntegral . unField $ meta.cmd_data_page_offset+ dictOffset = fromIntegral <$> unField meta.cmd_dictionary_page_offset+ in ( fromMaybe dataOffset dictOffset+ , fromIntegral . unField $ meta.cmd_total_compressed_size+ )++-- | The columns a read decodes, after projection, in schema order.+data ColumnPlan = ColumnPlan+ { planNames :: [T.Text]+ , planDescs :: [ColumnDescription]+ , planChunks :: [[ColumnChunk]]+ -- ^ Per column, one chunk per row group.+ }++-- | Read Parquet files from a directory or glob path.+readParquetFiles :: FilePath -> IO DataFrame+readParquetFiles = readParquetFilesWithOpts defaultParquetReadOptions++-- | Read multiple Parquet files (directory or glob) using options. readParquetFilesWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame readParquetFilesWithOpts opts path = do isDir <- doesDirectoryExist path@@ -288,12 +520,7 @@ dfs <- mapM (readParquetWithOpts optsWithoutRowRange) files pure (applyRowRange opts (mconcat dfs)) --- Core parsing pipeline -----------------------------------------------------{- | Parse a Parquet file via the 'RandomAccess' handle, applying all-read options. This is the central parsing entry point used by-'_readParquetWithOpts'.--}+-- | Parse a Parquet file via the 'RandomAccess' handle. parseParquetWithOpts :: (RandomAccess m, MonadIO m) => ParquetReadOptions ->@@ -304,92 +531,120 @@ {-# INLINEABLE parseParquetWithOpts #-} parseParquetWithOpts opts = do metadata <- parseFileMetadata+ let plan = planColumns opts metadata+ vectorLength = fileRowCount metadata+ rawCols <-+ zipWithM (parseColumnChunks vectorLength) (planChunks plan) (planDescs plan)+ return (assembleDataFrame opts vectorLength metadata plan rawCols) +-- | Row count from the footer, falling back to the sum of row-group counts.+fileRowCount :: FileMetadata -> Int+fileRowCount metadata =+ let topLevelRows = fromIntegral . unField $ metadata.num_rows :: Int+ rgRows =+ sum $ map (fromIntegral . unField . rg_num_rows) (unField metadata.row_groups) ::+ Int+ in if topLevelRows > 0 then topLevelRows else rgRows++columnMatches :: [T.Text] -> T.Text -> T.Text -> Bool+columnMatches allNames requested name =+ name == requested+ || (leafOf name == requested && requested `notElem` allNames)++-- | The last dot-separated segment of a column path.+leafOf :: T.Text -> T.Text+leafOf = last . T.splitOn "."++-- | Validate the requested columns against the schema and pick what to decode.+planColumns :: ParquetReadOptions -> FileMetadata -> ColumnPlan+planColumns opts metadata = let schemaElems = unField metadata.schema allNames = getColumnNames (drop 1 schemaElems)- leafNames = L.nub (map (last . T.splitOn ".") allNames)+ knownNames = L.nub (allNames ++ map leafOf allNames) predicateColumns = maybe [] (L.nub . getColumns) (predicate opts) selectedColumnsForRead = case selectedColumns opts of Nothing -> Nothing Just selected -> Just (L.nub (selected ++ predicateColumns)) - -- 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.+ -- TODO: With rowRange set, skip row groups the range does not overlap. - -- TODO: When predicate is set, inspect cmd_statistics min/max values for- -- predicate-referenced columns in each RowGroup and skip groups where- -- statistics prove the predicate cannot be satisfied.+ -- TODO: Skip row groups whose statistics rule out the predicate. - -- Validate selected columns- case selectedColumnsForRead of- Nothing -> pure ()- Just requested ->- let missing = requested L.\\ leafNames- in unless (L.null missing) $- liftIO $- throw- ( ColumnsNotFoundException- missing- "readParquetWithOpts"- leafNames- )+ -- Validate selected columns+ validated = case selectedColumnsForRead of+ Nothing -> ()+ Just requested ->+ let missing = [r | r <- requested, not (any (columnMatches allNames r) allNames)]+ in if L.null missing+ then ()+ else+ throw+ ( ColumnsNotFoundException+ missing+ "readParquetWithOpts"+ knownNames+ ) - let descriptions = generateColumnDescriptions schemaElems+ descriptions = generateColumnDescriptions schemaElems chunks = columnChunksForAll metadata nCols = length chunks nDescs = length descriptions - unless (nCols == nDescs) $- error $- "Column count mismatch: got "- <> show nCols- <> " columns but schema implied "- <> show nDescs- <> " columns"-- -- Some files omit the top-level num_rows field; fall back to summing row-group counts.- let topLevelRows = fromIntegral . unField $ metadata.num_rows :: Int- rgRows =- sum $ map (fromIntegral . unField . rg_num_rows) (unField metadata.row_groups) ::- Int- vectorLength = if topLevelRows > 0 then topLevelRows else rgRows+ counted =+ if nCols == nDescs+ then ()+ else+ error $+ "Column count mismatch: got "+ <> show nCols+ <> " columns but schema implied "+ <> show nDescs+ <> " columns" - -- 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+ -- Decode only the columns needed for the output plus any predicate.+ keep name = case selectedColumnsForRead of Nothing -> True- Just req -> last (T.splitOn "." name) `elem` req+ Just req -> any (\r -> columnMatches allNames r name) req kept = filter (\(n, _, _) -> keep n) (zip3 allNames chunks descriptions)- keptNames = [n | (n, _, _) <- kept]- keptChunks = [c | (_, c, _) <- kept]- keptDescs = [d | (_, _, d) <- kept]-- rawCols <- zipWithM (parseColumnChunks vectorLength) keptChunks keptDescs+ in validated `seq`+ counted `seq`+ ColumnPlan+ { planNames = [n | (n, _, _) <- kept]+ , planDescs = [d | (_, _, d) <- kept]+ , planChunks = [c | (_, c, _) <- kept]+ } +-- | Apply logical types and read options to the decoded columns.+assembleDataFrame ::+ ParquetReadOptions -> Int -> FileMetadata -> ColumnPlan -> [Column] -> DataFrame+assembleDataFrame opts totalRows metadata plan rawCols = let hints = nativeTypeHints metadata+ keptNames = planNames plan+ widen col+ | not (widenIntegers opts) = col+ | otherwise = case col of+ DI.UnboxedColumn Nothing (v :: VU.Vector t)+ | Just Refl <- testEquality (typeRep @t) (typeRep @Int32) ->+ DI.UnboxedColumn Nothing (VU.map (fromIntegral @Int32 @Int) v)+ | Just Refl <- testEquality (typeRep @t) (typeRep @Int64) ->+ DI.UnboxedColumn Nothing (VU.map (fromIntegral @Int64 @Int) v)+ _ -> col finalCols = zipWith (restoreNativeType hints) keptNames- (zipWith applyDescLogicalType keptDescs rawCols)+ (map widen (zipWith applyDescLogicalType (planDescs plan) rawCols)) indices = Map.fromList $ zip keptNames [0 ..]- dimensions = (vectorLength, length finalCols)-- let df =+ dimensions = (totalRows, length finalCols)+ df = DataFrame (Vector.fromListN (length finalCols) finalCols) indices dimensions Map.empty-- return (applyReadOptions opts df)+ in applyReadOptions opts df -{- | Parse the file-level Thrift metadata from the Parquet file footer.-Validates the trailing 4-byte magic marker (\"PAR1\") before decoding.--}+-- | Parse the file-level Thrift metadata from the Parquet footer. parseFileMetadata :: (RandomAccess m) => m FileMetadata parseFileMetadata = do footerBytes <- readSuffix 8@@ -438,6 +693,12 @@ ColumnDescription -> ReaderIO FileBufferedOrSeekable Column #-}+{-# SPECIALIZE parseColumnChunks ::+ Int ->+ [ColumnChunk] ->+ ColumnDescription ->+ ReaderIO ChunkBuffer Column+ #-} {-# INLINEABLE parseColumnChunks #-} parseColumnChunks totalRows chunks description | description.maxRepetitionLevel == 0 && description.maxDefinitionLevel == 0 =@@ -478,16 +739,11 @@ go decoder = foldNonNullable totalRows (foldColumnPagesM description decoder chunks) - -- Decode a non-nullable BYTE_ARRAY (UTF-8) column straight into a single- -- shared byte buffer + offsets ('PackedText'), instead of a boxed vector- -- of per-row 'Text'. Each page's decoded 'Text' values (which share the- -- chunk dictionary for dictionary-encoded pages) are appended by memcpy- -- into one builder across all pages/chunks, then frozen once. This is the- -- same representation the fast CSV reader uses and matches Arrow's string- -- layout: no retained per-row 'Text' headers, no eager UTF-8 validation.+ -- Decode non-nullable UTF-8 BYTE_ARRAY straight into one 'PackedText'. goPackedText :: m Column goPackedText = do- builder <- liftIO $ stToIO (newTextBuilder totalRows (totalRows * 8))+ builder <-+ liftIO $ stToIO (newTextBuilder totalRows (textByteHint totalRows chunks)) _ <- foldColumnDataPagesM description@@ -507,6 +763,16 @@ unboxedGo decoder = foldNonNullableUnboxed totalRows (foldColumnPagesM description decoder chunks) +-- | Initial byte capacity for a text column.+textByteHint :: Int -> [ColumnChunk] -> Int+textByteHint totalRows chunks =+ max (totalRows * 8) $+ sum+ [ fromIntegral (unField (cmd_total_uncompressed_size meta))+ | cc <- chunks+ , Just meta <- [unField (cc_meta_data cc)]+ ]+ -- | Decode an optional (nullable) column. {-# INLINEABLE getNullableColumn #-} getNullableColumn ::@@ -541,13 +807,11 @@ go decoder = foldNullable maxDef totalRows (foldColumnPagesM description decoder chunks) - -- Nullable BYTE_ARRAY (UTF-8): decode straight into a 'PackedText' (shared- -- byte buffer + offsets + validity bitmap) via the text builder, walking- -- def-levels to interleave nulls. Avoids the boxed @Vector Text@ the- -- generic 'foldNullable' path would build.+ -- Nullable UTF-8 BYTE_ARRAY: decode straight into a 'PackedText'. goPackedTextNullable :: m Column goPackedTextNullable = do- builder <- liftIO $ stToIO (newTextBuilder totalRows (totalRows * 8))+ builder <-+ liftIO $ stToIO (newTextBuilder totalRows (textByteHint totalRows chunks)) _ <- foldColumnDataPagesM description@@ -622,8 +886,6 @@ unboxedGo decoder = foldRepeatedUnboxed maxRep maxDef (foldColumnPagesM description decoder chunks) --- Options application ------------------------------------------------------ applyRowRange :: ParquetReadOptions -> DataFrame -> DataFrame applyRowRange opts df = maybe df (`DS.range` df) (rowRange opts)@@ -648,11 +910,7 @@ . applySelectedColumns opts . applyPredicate opts --- Logical type conversion ---------------------------------------------------{- | Apply a column-description's logical type annotation to convert raw-decoded values (e.g. millisecond integers → 'UTCTime').--}+-- | Apply a column's logical type annotation to its decoded values. applyDescLogicalType :: ColumnDescription -> DI.Column -> DI.Column applyDescLogicalType desc = applyLogicalType (colLogicalType desc) @@ -660,10 +918,7 @@ applyLogicalType (Just (LT_TIMESTAMP f)) col = let ts = unField f unit = unField ts.timestamp_unit- -- (ticks per second, picoseconds per tick) for each unit. Convert at- -- native precision: a millisecond grid keeps ms, a nanosecond grid- -- keeps ns. (The old code multiplied by @1_000_000 `div` divisor@,- -- which truncated NANOS to 0 and collapsed every value to the epoch.)+ -- Convert at native precision (ticks per second, ps per tick). conv = case unit of MILLIS _ -> epochToUTCTime 1_000 1_000_000_000 MICROS _ -> epochToUTCTime 1_000_000 1_000_000@@ -674,18 +929,14 @@ scale = unField dt.decimal_scale precision = unField dt.decimal_precision in if precision <= 9- then case DI.toVector @Int32 @VU.Vector col of- Right xs ->- DI.fromUnboxedVector $- VU.map (\raw -> fromIntegral @Int32 @Double raw / 10 ^ scale) xs- Left _ -> col+ then+ fromRight col $+ DI.mapColumn (\raw -> fromIntegral @Int32 @Double raw / 10 ^ scale) col else if precision <= 18- then case DI.toVector @Int64 @VU.Vector col of- Right xs ->- DI.fromUnboxedVector $- VU.map (\raw -> fromIntegral @Int64 @Double raw / 10 ^ scale) xs- Left _ -> col+ then+ fromRight col $+ DI.mapColumn (\raw -> fromIntegral @Int64 @Double raw / 10 ^ scale) col else col applyLogicalType _ col = col @@ -725,28 +976,23 @@ ] restoreNativeType :: Schema -> T.Text -> DI.Column -> DI.Column-restoreNativeType hints name col = case Map.lookup leaf (elements hints) of+restoreNativeType hints name col = case lookupHint of Just ty | stampedAs @Int ty -> narrow (fromIntegral @Int64 @Int) | stampedAs @Integer ty -> narrow (fromIntegral @Int64 @Integer) _ -> col where- leaf = last (T.splitOn "." name)+ lookupHint =+ case Map.lookup name (elements hints) of+ Just ty -> Just ty+ Nothing -> Map.lookup (leafOf name) (elements hints) stampedAs :: forall a. (Columnable a, Read a, Columnable (Maybe a)) => SchemaType -> Bool stampedAs ty = ty == schemaType @a || ty == schemaType @(Maybe a) narrow :: (Columnable a) => (Int64 -> a) -> DI.Column narrow f = fromRight col (DI.mapColumn f col) -{- | Convert an epoch timestamp expressed as @ticksPerSecond@ ticks/second-(each tick = @psPerTick@ picoseconds) to 'UTCTime', at full precision.--Splits into days + within-day picoseconds with integer 'divMod' (which floors,-so the split is correct for pre-epoch negative values too), and never forms-picoseconds-since-epoch — that would overflow 'Int64' for modern dates — only-the bounded within-day picosecond count. 40587 is the Modified Julian Day of-the Unix epoch (1970-01-01).--}+-- | Convert an epoch timestamp with the given tick rate to 'UTCTime'. epochToUTCTime :: Int64 -> Integer -> Int64 -> UTCTime epochToUTCTime ticksPerSecond psPerTick v = let (s, subTicks) = v `divMod` ticksPerSecond
src/DataFrame/IO/Parquet/Binary.hs view
@@ -18,12 +18,7 @@ readUVarInt :: BS.ByteString -> (Word64, BS.ByteString) readUVarInt xs = loop xs 0 0 0 where- {-- Each input byte contributes:- - lower 7 payload bits- - The high bit (0x80) is the continuation flag: 1 = more bytes follow, 0 = last byte- Why the magic number 10: For a 64‑bit integer we need at most ceil(64 / 7) = 10 bytes- -}+ -- The high bit flags continuation; 10 bytes cover a 64-bit integer. loop :: BS.ByteString -> Word64 -> Int -> Int -> (Word64, BS.ByteString) loop bs result _ 10 = (result, bs) loop xs' result shiftAmt i = case BS.uncons xs' of@@ -113,9 +108,7 @@ (fromIntegral size) ((`readSingleByte` buf) . fromIntegral) -{- | Allocate a fixed-size buffer, repeat the action on each index.-Fill it into the buffer to get a ByteString.--}+-- | Allocate a fixed-size buffer and fill it by repeating an action. fillByteStringByWord8 :: Int -> (Int -> IO Word8) -> IO BS.ByteString fillByteStringByWord8 size getByte = do bracketOnError
src/DataFrame/IO/Parquet/Dictionary.hs view
@@ -1,6 +1,17 @@ {-# LANGUAGE BangPatterns #-} -module DataFrame.IO.Parquet.Dictionary (DictVals (..), readDictVals, decodeRLEBitPackedHybrid) where+module DataFrame.IO.Parquet.Dictionary (+ DictVals (..),+ readDictVals,+ dictText,+ dictInt32,+ dictInt64,+ dictFloat,+ dictDouble,+ word32At,+ word64At,+ decodeRLEBitPackedHybrid,+) where import Data.Bits import qualified Data.ByteString as BS@@ -10,109 +21,101 @@ import Data.Text.Encoding import Data.Time (UTCTime) import qualified Data.Vector as V+import qualified Data.Vector.Unboxed as VU import Data.Word import DataFrame.IO.Parquet.Binary (readUVarInt) import DataFrame.IO.Parquet.Thrift (ThriftType (..)) import DataFrame.IO.Parquet.Time (int96ToUTCTime)-import DataFrame.Internal.Binary (- littleEndianInt32,- littleEndianWord32,- littleEndianWord64,- )+import DataFrame.Internal.Binary (littleEndianWord32) import GHC.Float +-- | A decoded dictionary page with unboxed numeric entries. data DictVals- = DBool (V.Vector Bool)- | DInt32 (V.Vector Int32)- | DInt64 (V.Vector Int64)- | DInt96 (V.Vector UTCTime)- | DFloat (V.Vector Float)- | DDouble (V.Vector Double)- | DText (V.Vector T.Text)+ = DBool !(VU.Vector Bool)+ | DInt32 !(VU.Vector Int32)+ | DInt64 !(VU.Vector Int64)+ | DInt96 !(V.Vector UTCTime)+ | DFloat !(VU.Vector Float)+ | DDouble !(VU.Vector Double)+ | DText !(VU.Vector Int) !(VU.Vector Int) !BS.ByteString deriving (Show, Eq) -{- | Decode the values from a dictionary page.--The @numVals@ argument is the entry count declared in the dictionary page-header. It is used to limit BOOLEAN decoding (1-bit-per-value encoding has-no natural delimiter).--The @typeLength@ argument is only meaningful for FIXED_LEN_BYTE_ARRAY: it is-the byte-width of each individual dictionary entry, NOT the total number of-entries. Passing @numVals@ here (the old behaviour) would cause it to be-misread as an element size, yielding a dictionary that is far too small.--}+-- | Decode the values from a dictionary page. readDictVals :: ThriftType -> BS.ByteString -> Int32 -> Maybe Int32 -> DictVals-readDictVals (BOOLEAN _) bs count _ = DBool (V.fromList (take (fromIntegral count) $ readPageBool bs))-readDictVals (INT32 _) bs _ _ = DInt32 (V.fromList (readPageInt32 bs))-readDictVals (INT64 _) bs _ _ = DInt64 (V.fromList (readPageInt64 bs))-readDictVals (INT96 _) bs _ _ = DInt96 (V.fromList (readPageInt96Times bs))-readDictVals (FLOAT _) bs _ _ = DFloat (V.fromList (readPageFloat bs))-readDictVals (DOUBLE _) bs _ _ = DDouble (V.fromList (readPageWord64 bs))-readDictVals (BYTE_ARRAY _) bs _ _ = DText (V.fromList (readPageBytes bs))+readDictVals (BOOLEAN _) bs count _ =+ let n = min (fromIntegral count) (8 * BS.length bs)+ in DBool $ VU.generate n $ \i ->+ (BSU.unsafeIndex bs (i `shiftR` 3) `shiftR` (i .&. 7)) .&. 1 == 1+readDictVals (INT32 _) bs _ _ =+ DInt32 $ VU.generate (BS.length bs `div` 4) $ \i -> fromIntegral (word32At bs (4 * i))+readDictVals (INT64 _) bs _ _ =+ DInt64 $ VU.generate (BS.length bs `div` 8) $ \i -> fromIntegral (word64At bs (8 * i))+readDictVals (INT96 _) bs _ _ =+ DInt96 $ V.generate (BS.length bs `div` 12) $ \i ->+ int96ToUTCTime (BSU.unsafeTake 12 (BSU.unsafeDrop (12 * i) bs))+readDictVals (FLOAT _) bs _ _ =+ DFloat $ VU.generate (BS.length bs `div` 4) $ \i -> castWord32ToFloat (word32At bs (4 * i))+readDictVals (DOUBLE _) bs _ _ =+ DDouble $ VU.generate (BS.length bs `div` 8) $ \i -> castWord64ToDouble (word64At bs (8 * i))+readDictVals (BYTE_ARRAY _) bs _ _ = textDict bs readDictVals (FIXED_LEN_BYTE_ARRAY _) bs _ (Just len) =- DText (V.fromList (readPageFixedBytes bs (fromIntegral len)))+ let w = fromIntegral len+ n = if w <= 0 then 0 else BS.length bs `div` w+ in DText (VU.generate n (* w)) (VU.replicate n w) bs readDictVals t _ _ _ = error $ "Unsupported dictionary type: " ++ show t -readPageInt32 :: BS.ByteString -> [Int32]-readPageInt32 xs- | BS.null xs = []- | otherwise = littleEndianInt32 (BS.take 4 xs) : readPageInt32 (BS.drop 4 xs)--readPageWord64 :: BS.ByteString -> [Double]-readPageWord64 xs- | BS.null xs = []- | otherwise =- castWord64ToDouble (littleEndianWord64 (BS.take 8 xs))- : readPageWord64 (BS.drop 8 xs)+-- | Length-prefixed byte-array entries, in two passes.+textDict :: BS.ByteString -> DictVals+textDict bs = DText starts lens bs+ where+ !len = BS.length bs+ count !off !n+ | off + 4 > len = n+ | otherwise = count (off + 4 + fromIntegral (word32At bs off)) (n + 1)+ !n = count 0 0+ positions = VU.unfoldrExactN n step 0+ step off = let l = fromIntegral (word32At bs off) in ((off + 4, l), off + 4 + l)+ (starts, lens) = VU.unzip positions -readPageBytes :: BS.ByteString -> [T.Text]-readPageBytes xs- | BS.null xs = []- | otherwise =- let lenBytes = fromIntegral (littleEndianInt32 $ BS.take 4 xs)- totalBytesRead = lenBytes + 4- in decodeUtf8Lenient (BS.take lenBytes (BS.drop 4 xs))- : readPageBytes (BS.drop totalBytesRead xs)+-- | Dictionary entry @i@ as 'T.Text' (lenient decode, for the boxed paths).+dictText :: DictVals -> Int -> T.Text+dictText (DText starts lens bs) i =+ decodeUtf8Lenient+ ( BSU.unsafeTake+ (VU.unsafeIndex lens i)+ (BSU.unsafeDrop (VU.unsafeIndex starts i) bs)+ )+dictText d _ = error ("dictText: wrong dict type, got " ++ show d) -readPageBool :: BS.ByteString -> [Bool]-readPageBool bs =- concatMap (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7]) (BS.unpack bs)+dictInt32 :: DictVals -> Maybe (VU.Vector Int32)+dictInt32 (DInt32 ds) = Just ds+dictInt32 _ = Nothing -readPageInt64 :: BS.ByteString -> [Int64]-readPageInt64 xs- | BS.null xs = []- | otherwise =- fromIntegral (littleEndianWord64 (BS.take 8 xs)) : readPageInt64 (BS.drop 8 xs)+dictInt64 :: DictVals -> Maybe (VU.Vector Int64)+dictInt64 (DInt64 ds) = Just ds+dictInt64 _ = Nothing -readPageFloat :: BS.ByteString -> [Float]-readPageFloat xs- | BS.null xs = []- | otherwise =- castWord32ToFloat (littleEndianWord32 (BS.take 4 xs))- : readPageFloat (BS.drop 4 xs)+dictFloat :: DictVals -> Maybe (VU.Vector Float)+dictFloat (DFloat ds) = Just ds+dictFloat _ = Nothing -readNInt96Times :: Int -> BS.ByteString -> ([UTCTime], BS.ByteString)-readNInt96Times 0 bs = ([], bs)-readNInt96Times k bs =- let timestamp96 = BS.take 12 bs- utcTime = int96ToUTCTime timestamp96- bs' = BS.drop 12 bs- (times, rest) = readNInt96Times (k - 1) bs'- in (utcTime : times, rest)+dictDouble :: DictVals -> Maybe (VU.Vector Double)+dictDouble (DDouble ds) = Just ds+dictDouble _ = Nothing -readPageInt96Times :: BS.ByteString -> [UTCTime]-readPageInt96Times bs- | BS.null bs = []- | otherwise =- let (times, _) = readNInt96Times (BS.length bs `div` 12) bs- in times+word32At :: BS.ByteString -> Int -> Word32+word32At bs o =+ fromIntegral (BSU.unsafeIndex bs o)+ .|. (fromIntegral (BSU.unsafeIndex bs (o + 1)) `unsafeShiftL` 8)+ .|. (fromIntegral (BSU.unsafeIndex bs (o + 2)) `unsafeShiftL` 16)+ .|. (fromIntegral (BSU.unsafeIndex bs (o + 3)) `unsafeShiftL` 24)+{-# INLINE word32At #-} -readPageFixedBytes :: BS.ByteString -> Int -> [T.Text]-readPageFixedBytes xs len- | BS.null xs = []- | otherwise =- decodeUtf8Lenient (BS.take len xs) : readPageFixedBytes (BS.drop len xs) len+word64At :: BS.ByteString -> Int -> Word64+word64At bs o =+ fromIntegral (word32At bs o)+ .|. (fromIntegral (word32At bs (o + 4)) `unsafeShiftL` 32)+{-# INLINE word64At #-} unpackBitPacked :: Int -> Int -> BS.ByteString -> ([Word32], BS.ByteString) unpackBitPacked bw count bs@@ -124,7 +127,7 @@ rest = BS.drop totalBytes bs in (extractBits bw count chunk, rest) --- | LSB-first bit accumulator: reads each byte once with no intermediate ByteString allocation.+-- | LSB-first bit accumulator that reads each byte once. extractBits :: Int -> Int -> BS.ByteString -> [Word32] extractBits bw count bs = go 0 (0 :: Word64) 0 count where@@ -145,9 +148,7 @@ | bitWidth == 0 = ([0], bs) | BS.null bs = ([], bs) | otherwise =- -- readUVarInt is evaluated here, inside the guard that has already- -- confirmed bs is non-empty. Keeping it in a where clause would cause- -- it to be forced before the BS.null guard under {-# LANGUAGE Strict #-}.+ -- Evaluated inside the guard: a where clause would force it too early. let (hdr64, afterHdr) = readUVarInt bs isPacked = (hdr64 .&. 1) == 1 in if isPacked
src/DataFrame/IO/Parquet/Encoding.hs view
@@ -7,11 +7,13 @@ bitWidthForMaxLevel, -- Vector-based RLE/bit-packed decoder (from new parser) decodeRLEBitPackedHybrid,+ rleHybridFold, extractBitsInto, fillRun, decodeDictIndices, ) where +import Control.Monad (when) import Control.Monad.ST (ST, runST) import Data.Bits import qualified Data.ByteString as BS@@ -22,10 +24,6 @@ import DataFrame.IO.Parquet.Binary (readUVarInt) import DataFrame.Internal.Binary (littleEndianWord32) --- ------------------------------------------------------------------------------ Level-width helpers (used by Levels.hs)--- ---------------------------------------------------------------------------- ceilLog2 :: Int -> Int ceilLog2 x | x <= 1 = 0@@ -34,10 +32,6 @@ bitWidthForMaxLevel :: Int -> Int bitWidthForMaxLevel maxLevel = ceilLog2 (maxLevel + 1) --- ------------------------------------------------------------------------------ Vector-based RLE / bit-packed hybrid decoder--- ---------------------------------------------------------------------------- decodeRLEBitPackedHybrid :: -- | Bit width per value (0 = all zeros, use 'VU.replicate') Int ->@@ -66,8 +60,7 @@ let groups = fromIntegral (hdr64 `shiftR` 1) :: Int totalVals = groups * 8 takeN = min (need - filled) totalVals- -- Consume all the bytes for this group even if we- -- only need a subset of the values.+ -- Consume the whole group even if we need less. bytesN = (bw * totalVals + 7) `div` 8 (chunk, rest) = BS.splitAt bytesN afterHdr extractBitsInto bw takeN chunk mv filled@@ -87,10 +80,7 @@ fillRun mv i end = VUM.set (VUM.unsafeSlice i (end - i) mv) {-# INLINE fillRun #-} -{- | Write @count@ bit-width-@bw@ values from @bs@ into @mv@ starting at-@offset@, reading the byte buffer with a single-pass LSB-first accumulator.-No intermediate list or ByteString allocation.--}+-- | Write bit-packed values into a mutable vector in a single pass. extractBitsInto :: -- | Bit width Int ->@@ -116,12 +106,7 @@ in go (byteIdx + 1) (acc .|. (b `unsafeShiftL` accBits)) (accBits + 8) done {-# INLINE extractBitsInto #-} -{- | Decode @need@ dictionary indices from a DATA_PAGE bit-width-prefixed-stream (the first byte encodes the bit-width of all subsequent RLE\/bitpacked-values).--Returns the index vector (as 'Int') and the unconsumed bytes.--}+-- | Decode dictionary indices from a bit-width-prefixed DATA_PAGE stream. decodeDictIndices :: Int -> BS.ByteString -> (VU.Vector Int, BS.ByteString) decodeDictIndices need bs = case BS.uncons bs of Nothing -> error "decodeDictIndices: empty stream"@@ -130,3 +115,55 @@ (raw, rest1) = decodeRLEBitPackedHybrid bw need rest0 in (VU.map fromIntegral raw, rest1) {-# INLINE decodeDictIndices #-}++-- | Stream RLE\/bit-packed hybrid values to a sink.+rleHybridFold ::+ Int -> Int -> BS.ByteString -> (Int -> Word32 -> ST s ()) -> ST s ()+rleHybridFold bw need bs sink+ | bw == 0 = zeros 0+ | otherwise = go 0 bs+ where+ !mask = if bw == 32 then maxBound else (1 `unsafeShiftL` bw) - 1 :: Word32+ zeros !i = when (i < need) (sink i 0 >> zeros (i + 1))+ go !filled !buf+ | filled >= need = pure ()+ | BS.null buf = pure ()+ | otherwise =+ let (hdr64, afterHdr) = readUVarInt buf+ isPacked = (hdr64 .&. 1) == 1+ in if isPacked+ then do+ let groups = fromIntegral (hdr64 `shiftR` 1) :: Int+ totalVals = groups * 8+ takeN = min (need - filled) totalVals+ bytesN = (bw * totalVals + 7) `div` 8+ (chunk, rest) = BS.splitAt bytesN afterHdr+ extractBitsFold bw takeN chunk (\j w -> sink (filled + j) w)+ go (filled + takeN) rest+ else do+ let runLen = fromIntegral (hdr64 `shiftR` 1) :: Int+ nbytes = (bw + 7) `div` 8+ val = littleEndianWord32 (BS.take 4 afterHdr) .&. mask+ takeN = min (need - filled) runLen+ run !j = when (j < takeN) (sink (filled + j) val >> run (j + 1))+ run 0+ go (filled + takeN) (BS.drop nbytes afterHdr)+{-# INLINE rleHybridFold #-}++-- | 'extractBitsInto' with a sink instead of a destination vector.+extractBitsFold ::+ Int -> Int -> BS.ByteString -> (Int -> Word32 -> ST s ()) -> ST s ()+extractBitsFold bw count bs sink = go 0 (0 :: Word64) 0 0+ where+ !mask = if bw == 32 then maxBound else (1 `unsafeShiftL` bw) - 1 :: Word64+ !len = BS.length bs+ go !byteIdx !acc !accBits !done+ | done >= count = return ()+ | accBits >= bw = do+ sink done (fromIntegral (acc .&. mask))+ go byteIdx (acc `unsafeShiftR` bw) (accBits - bw) (done + 1)+ | byteIdx >= len = return ()+ | otherwise =+ let b = fromIntegral (BSU.unsafeIndex bs byteIdx) :: Word64+ in go (byteIdx + 1) (acc .|. (b `unsafeShiftL` accBits)) (accBits + 8) done+{-# INLINE extractBitsFold #-}
src/DataFrame/IO/Parquet/Levels.hs view
@@ -23,14 +23,7 @@ ) import DataFrame.Internal.Binary (littleEndianWord32) --- ------------------------------------------------------------------------------ Level readers--- -----------------------------------------------------------------------------{- | Convert a 'Word32' level vector to an 'Int' level vector while counting-how many entries equal @maxDef@. Single pass; allocates a single-'VU.Vector Int' of length @VU.length raw@.--}+-- | Convert a 'Word32' level vector to 'Int', counting @maxDef@ entries. convertAndCount :: Int -> VU.Vector Word32 -> (VU.Vector Int, Int) convertAndCount maxDef raw = runST $ do let !n = VU.length raw@@ -65,8 +58,8 @@ (defVec, nPresent, afterDef) = decodeLevelBlock bwDef n afterRep in (defVec, repVec, nPresent, afterDef) where- -- For rep block we don't need nPresent; we still get one cheaply.- decodeLevelBlock 0 n' buf = (VU.replicate n' 0, n' * fromEnum (maxDef == 0), buf)+ -- No levels to read: non-nullable, non-repeated paths never use them.+ decodeLevelBlock 0 n' buf = (VU.empty, n' * fromEnum (maxDef == 0), buf) decodeLevelBlock bw n' buf = let blockLen = fromIntegral (littleEndianWord32 (BS.take 4 buf)) :: Int blockData = BS.take blockLen (BS.drop 4 buf)@@ -94,22 +87,19 @@ bwRep = bitWidthForMaxLevel maxRep bwDef = bitWidthForMaxLevel maxDef repVec- | bwRep == 0 = VU.replicate n 0+ | bwRep == 0 = VU.empty | otherwise = let (raw, _) = decodeRLEBitPackedHybrid bwRep n repBytes (out, _) = convertAndCount maxDef raw in out (defVec, nPresent)- | bwDef == 0 = (VU.replicate n 0, n * fromEnum (maxDef == 0))+ | bwDef == 0 = (VU.empty, n * fromEnum (maxDef == 0)) | otherwise = let (raw, _) = decodeRLEBitPackedHybrid bwDef n defBytes in convertAndCount maxDef raw in (defVec, repVec, nPresent, afterDefBytes) -{- | Stitch a singly-nested list column (@maxRep == 1@) from vector-format-definition and repetition levels plus a compact present-values vector.-Returns one @Maybe [Maybe a]@ per top-level row.--}+-- | Stitch a singly-nested list column from levels and present values. stitchList :: Int -> VU.Vector Int ->@@ -123,9 +113,7 @@ toRow ((_, d, _) : _) | d == 0 = Nothing toRow grp = Just [v | (_, _, v) <- grp] -{- | Stitch a doubly-nested list column (@maxRep == 2@).-@defT1@ is the def threshold at which the depth-1 element is present.--}+-- | Stitch a doubly-nested list column (@maxRep == 2@). stitchList2 :: Int -> Int ->@@ -146,10 +134,7 @@ toLeaf [] = Nothing toLeaf ((_, _, v) : _) = v -{- | Stitch a triply-nested list column (@maxRep == 3@).-@defT1@ and @defT2@ are the def thresholds for depth-1 and depth-2-elements respectively.--}+-- | Stitch a triply-nested list column (@maxRep == 3@). stitchList3 :: Int -> Int ->@@ -174,14 +159,7 @@ toLeaf [] = Nothing toLeaf ((_, _, v) : _) = v --- ------------------------------------------------------------------------------ Internal helpers--- -----------------------------------------------------------------------------{- | Zip rep and def level vectors with a present-values vector, tagging each-position as @Just value@ (when @def == maxDef@) or @Nothing@.-Returns a flat list of @(rep, def, Maybe a)@ triplets for row-splitting.--}+-- | Zip rep and def levels with present values into triplets. pairWithValsV :: Int -> VU.Vector Int ->@@ -200,9 +178,7 @@ then (r, d, Just (VB.unsafeIndex values j)) : go (i + 1) (j + 1) else (r, d, Nothing) : go (i + 1) j -{- | Group a flat triplet list into rows.-A new group begins whenever @rep <= bound@.--}+-- | Group a flat triplet list into rows. splitAtRepBound :: Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]] splitAtRepBound _ [] = [] splitAtRepBound bound (t : ts) =
src/DataFrame/IO/Parquet/Page.hs view
@@ -19,6 +19,7 @@ foldColumnPagesM, foldColumnDataPagesM, RawPage,+ decodeRequiredNumericInto, appendStringPageIO, appendNullableStringPageIO, ) where@@ -36,13 +37,21 @@ import qualified Data.Vector as VB import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM import Data.Word (Word8) import DataFrame.IO.Parquet.Decompress (decompressData) import DataFrame.IO.Parquet.Dictionary ( DictVals (..),+ dictText, readDictVals,+ word32At,+ word64At, )-import DataFrame.IO.Parquet.Encoding (decodeDictIndices)+import DataFrame.IO.Parquet.Encoding (+ decodeDictIndices,+ decodeRLEBitPackedHybrid,+ rleHybridFold,+ ) import DataFrame.IO.Parquet.Levels (readLevelsV1, readLevelsV2) import DataFrame.IO.Parquet.Thrift ( ColumnChunk (..),@@ -60,15 +69,10 @@ import DataFrame.IO.Parquet.Time (int96ToUTCTime) import DataFrame.IO.Parquet.Utils (ColumnDescription (..)) import DataFrame.IO.Utils.RandomAccess (RandomAccess (..), Range (Range))-import DataFrame.Internal.Binary (- littleEndianInt32,- littleEndianWord32,- littleEndianWord64,- )+import DataFrame.Internal.Binary (littleEndianInt32) import DataFrame.Internal.Column.Builder ( TextBuilder, appendNull,- appendText, appendTextSliceFromPtr, ) import Foreign.Ptr (Ptr, castPtr, plusPtr)@@ -77,58 +81,45 @@ import Pinch (decodeWithLeftovers) import qualified Pinch --- ------------------------------------------------------------------------------ Types--- -----------------------------------------------------------------------------{- | A type-specific page decoder.-Given the optional dictionary, the page encoding, the number of present-values, and the decompressed value bytes, returns exactly @nPresent@ values.--}+-- | A type-specific page decoder. type PageDecoder a = Maybe DictVals -> Encoding -> Int -> BS.ByteString -> VB.Vector a type UnboxedPageDecoder a = Maybe DictVals -> Encoding -> Int -> BS.ByteString -> VU.Vector a --- ------------------------------------------------------------------------------ Per-type decoders--- ---------------------------------------------------------------------------- boolDecoder :: UnboxedPageDecoder Bool boolDecoder mDict enc nPresent bs = case enc of- -- PLAIN bools are bit-packed (1 bit/value, LSB-first). Generate the- -- unboxed vector directly by indexing the bit for each row, avoiding an- -- intermediate @[Bool]@ list.+ -- PLAIN bools are bit-packed; index the bits directly, avoiding a list. PLAIN _ -> VU.generate nPresent $ \i -> (BSU.unsafeIndex bs (i `shiftR` 3) `shiftR` (i .&. 7)) .&. 1 == 1- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getBool- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getBool+ RLE_DICTIONARY _ -> gatherDict "boolDecoder" getBool mDict nPresent bs+ PLAIN_DICTIONARY _ -> gatherDict "boolDecoder" getBool mDict nPresent bs _ -> error ("boolDecoder: unsupported encoding " ++ show enc) where- getBool (DBool ds) i = ds VB.! i- getBool d _ = error ("boolDecoder: wrong dict type, got " ++ show d)+ getBool (DBool ds) = Just ds+ getBool _ = Nothing int32Decoder :: UnboxedPageDecoder Int32 int32Decoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.convert (readNInt32 nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getInt32- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getInt32+ PLAIN _ -> readNInt32 nPresent bs+ RLE_DICTIONARY _ -> gatherDict "int32Decoder" getInt32 mDict nPresent bs+ PLAIN_DICTIONARY _ -> gatherDict "int32Decoder" getInt32 mDict nPresent bs _ -> error ("int32Decoder: unsupported encoding " ++ show enc) where- getInt32 (DInt32 ds) i = ds VB.! i- getInt32 d _ = error ("int32Decoder: wrong dict type, got " ++ show d)+ getInt32 (DInt32 ds) = Just ds+ getInt32 _ = Nothing int64Decoder :: UnboxedPageDecoder Int64 int64Decoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.convert (readNInt64 nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getInt64- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getInt64+ PLAIN _ -> readNInt64 nPresent bs+ RLE_DICTIONARY _ -> gatherDict "int64Decoder" getInt64 mDict nPresent bs+ PLAIN_DICTIONARY _ -> gatherDict "int64Decoder" getInt64 mDict nPresent bs _ -> error ("int64Decoder: unsupported encoding " ++ show enc) where- getInt64 (DInt64 ds) i = ds VB.! i- getInt64 d _ = error ("int64Decoder: wrong dict type, got " ++ show d)+ getInt64 (DInt64 ds) = Just ds+ getInt64 _ = Nothing int96Decoder :: PageDecoder UTCTime int96Decoder mDict enc nPresent bs = case enc of@@ -137,28 +128,28 @@ PLAIN_DICTIONARY _ -> lookupDict mDict nPresent bs getInt96 _ -> error ("int96Decoder: unsupported encoding " ++ show enc) where- getInt96 (DInt96 ds) i = ds VB.! i+ getInt96 (DInt96 ds) i = VB.unsafeIndex ds i getInt96 d _ = error ("int96Decoder: wrong dict type, got " ++ show d) floatDecoder :: UnboxedPageDecoder Float floatDecoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.convert (readNFloat nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getFloat- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getFloat+ PLAIN _ -> readNFloat nPresent bs+ RLE_DICTIONARY _ -> gatherDict "floatDecoder" getFloat mDict nPresent bs+ PLAIN_DICTIONARY _ -> gatherDict "floatDecoder" getFloat mDict nPresent bs _ -> error ("floatDecoder: unsupported encoding " ++ show enc) where- getFloat (DFloat ds) i = ds VB.! i- getFloat d _ = error ("floatDecoder: wrong dict type, got " ++ show d)+ getFloat (DFloat ds) = Just ds+ getFloat _ = Nothing doubleDecoder :: UnboxedPageDecoder Double doubleDecoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.convert (readNDouble nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getDouble- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getDouble+ PLAIN _ -> readNDouble nPresent bs+ RLE_DICTIONARY _ -> gatherDict "doubleDecoder" getDouble mDict nPresent bs+ PLAIN_DICTIONARY _ -> gatherDict "doubleDecoder" getDouble mDict nPresent bs _ -> error ("doubleDecoder: unsupported encoding " ++ show enc) where- getDouble (DDouble ds) i = ds VB.! i- getDouble d _ = error ("doubleDecoder: wrong dict type, got " ++ show d)+ getDouble (DDouble ds) = Just ds+ getDouble _ = Nothing byteArrayDecoder :: PageDecoder T.Text byteArrayDecoder mDict enc nPresent bs = case enc of@@ -167,8 +158,7 @@ PLAIN_DICTIONARY _ -> lookupDict mDict nPresent bs getText _ -> error ("byteArrayDecoder: unsupported encoding " ++ show enc) where- getText (DText ds) i = ds VB.! i- getText d _ = error ("byteArrayDecoder: wrong dict type, got " ++ show d)+ getText = dictText fixedLenByteArrayDecoder :: Int -> PageDecoder T.Text fixedLenByteArrayDecoder len mDict enc nPresent bs = case enc of@@ -177,12 +167,9 @@ PLAIN_DICTIONARY _ -> lookupDict mDict nPresent bs getText _ -> error ("fixedLenByteArrayDecoder: unsupported encoding " ++ show enc) where- getText (DText ds) i = ds VB.! i- getText d _ = error ("fixedLenByteArrayDecoder: wrong dict type, got " ++ show d)+ getText = dictText -{- | Shared dictionary-path helper: decode @nPresent@ RLE/bit-packed indices-and look each one up in the dictionary.--}+-- | Decode RLE\/bit-packed indices and look each up in the dictionary. lookupDict :: Maybe DictVals -> Int ->@@ -195,22 +182,25 @@ let (idxs, _) = decodeDictIndices nPresent bs in VB.generate nPresent (f dict . VU.unsafeIndex idxs) -unboxedLookupDict ::+-- | Decode dictionary indices and gather from an unboxed dictionary.+gatherDict :: (VU.Unbox a) =>+ String ->+ (DictVals -> Maybe (VU.Vector a)) -> Maybe DictVals -> Int -> BS.ByteString ->- (DictVals -> Int -> a) -> VU.Vector a-unboxedLookupDict mDict nPresent bs f = case mDict of- Nothing -> error "Dictionary-encoded page but no dictionary page seen"- Just dict ->- let (idxs, _) = decodeDictIndices nPresent bs- in VU.generate nPresent (f dict . VU.unsafeIndex idxs)---- ------------------------------------------------------------------------------ Core page-iteration loop--- ---------------------------------------------------------------------------+gatherDict who pick mDict nPresent bs = case mDict of+ Nothing -> error (who ++ ": dictionary-encoded page but no dictionary page seen")+ Just dict -> case pick dict of+ Nothing -> error (who ++ ": wrong dict type, got " ++ show dict)+ Just ds -> case BS.uncons bs of+ Nothing -> error (who ++ ": empty index stream")+ Just (w0, rest) ->+ let (idxs, _) = decodeRLEBitPackedHybrid (fromIntegral w0) nPresent rest+ in VU.map (VU.unsafeIndex ds . fromIntegral) idxs+{-# INLINE gatherDict #-} -- | Read the raw (compressed) byte range for a column chunk. readChunkBytes ::@@ -228,30 +218,11 @@ rawBytes <- readBytes (Range offset compLen) return (codec, pType, rawBytes) -{- | A decoded DATA page handed to a fold step: the running dictionary, the-page encoding, the present-value count, the (decompressed) value bytes, and the-definition/repetition level vectors.--}+-- | A decoded DATA page handed to a fold step. type RawPage = (Maybe DictVals, Encoding, Int, BS.ByteString, VU.Vector Int, VU.Vector Int) -{- | Left-fold a monadic step over every DATA page (V1 or V2) of every column-chunk, in order, threading the running dictionary internally and handing each-page to @step@ as a 'RawPage' (no value decoding — the step decides how to-materialize, e.g. into a typed vector or straight into a text buffer).--Dictionary pages update the running dictionary; @INDEX_PAGE@s are skipped. Only-one page's bytes are live at a time.---- 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--- overlap the requested range, avoiding decompression of irrelevant pages--- entirely.--}+-- | Left-fold a monadic step over every DATA page of every column chunk. {-# INLINEABLE foldColumnDataPagesM #-} foldColumnDataPagesM :: forall m acc.@@ -311,8 +282,7 @@ enc = unField dph2.dph2_encoding defLen = unField dph2.dph2_definition_levels_byte_length repLen = unField dph2.dph2_repetition_levels_byte_length- -- V2: levels are never compressed; only the value- -- payload is (optionally) compressed.+ -- V2 levels are never compressed, only the values. isCompressed = fromMaybe True (unField dph2.dph2_is_compressed) (defLvls, repLvls, nPresent, compValBytes) = readLevelsV2 n maxDef maxRep repLen defLen pageData@@ -324,10 +294,7 @@ goPages dict codec pType rest' acc' INDEX_PAGE _ -> goPages dict codec pType rest' acc -{- | Left-fold over per-page value triples, decoding each page with @decoder@.-A thin wrapper over 'foldColumnDataPagesM' for the typed (numeric/boxed) column-paths.--}+-- | Left-fold over per-page value triples with a given decoder. {-# INLINEABLE foldColumnPagesM #-} foldColumnPagesM :: forall m v a acc.@@ -343,14 +310,7 @@ \acc (dict, enc, nPresent, valBytes, defLvls, repLvls) -> step acc (decoder dict enc nPresent valBytes, defLvls, repLvls) -{- | Append a non-nullable BYTE_ARRAY page's strings straight into a-'TextBuilder', no intermediate boxed 'Text' vector.--PLAIN pages copy each value's UTF-8 bytes directly from the page buffer pointer-('appendTextSliceFromPtr'), skipping the per-value 'decodeUtf8Lenient' +-'ByteString' slicing + 'Text' allocation that @readNTexts@ would do.-Dictionary pages append the shared dictionary 'Text's by reference.--}+-- | Append a non-nullable BYTE_ARRAY page's strings into a 'TextBuilder'. appendStringPageIO :: TextBuilder RealWorld -> Maybe DictVals ->@@ -372,15 +332,21 @@ _ -> error ("appendStringPageIO: unsupported encoding " ++ show enc) where dictAppend = case mDict of- Just (DText ds) ->- let (idxs, _) = decodeDictIndices nPresent bs- in stToIO (VU.mapM_ (\i -> appendText builder (ds VB.! i)) idxs)+ Just (DText starts lens dictBytes) ->+ BSU.unsafeUseAsCStringLen dictBytes $ \(cptr, _) -> do+ let p = castPtr cptr :: Ptr Word8+ (idxs, _) = decodeDictIndices nPresent bs+ VU.forM_ idxs $ \i ->+ stToIO+ ( appendTextSliceFromPtr+ builder+ (p `plusPtr` VU.unsafeIndex starts i)+ (VU.unsafeIndex lens i)+ ) Just d -> error ("appendStringPageIO: wrong dict type, got " ++ show d) Nothing -> error "appendStringPageIO: dictionary-encoded page but no dictionary seen" -{- | Read a little-endian 4-byte length prefix at byte @off@ from a raw page-pointer.--}+-- | Read a little-endian 4-byte length prefix from a raw page pointer. readLenLE :: Ptr Word8 -> Int -> IO Int readLenLE p off = do b0 <- peekByteOff p off :: IO Word8@@ -394,13 +360,7 @@ .|. (fromIntegral b3 `shiftL` 24) {-# INLINE readLenLE #-} -{- | Append a NULLABLE BYTE_ARRAY page into a 'TextBuilder', interleaving nulls-by walking the definition levels: a present row (@def == maxDef@) consumes the-next stored value (dictionary entry by reference, or the next length-prefixed-PLAIN slice by memcpy); a null row ('appendNull') writes an offset and clears a-validity bit. Builds 'PackedText' + validity bitmap with no per-row 'Text' or-@[Maybe a]@ allocation. Only present values are stored in the page payload.--}+-- | Append a nullable BYTE_ARRAY page into a 'TextBuilder'. appendNullableStringPageIO :: TextBuilder RealWorld -> Int ->@@ -427,13 +387,20 @@ where nDefs = VU.length defs dictAppend = case mDict of- Just (DText ds) -> do- let (idxs, _) = decodeDictIndices nPresent bs+ Just (DText starts lens dictBytes) -> BSU.unsafeUseAsCStringLen dictBytes $ \(cptr, _) -> do+ let p = castPtr cptr :: Ptr Word8+ (idxs, _) = decodeDictIndices nPresent bs go !_ !i | i >= nDefs = pure () go !j !i- | VU.unsafeIndex defs i == maxDef =- stToIO (appendText builder (ds VB.! VU.unsafeIndex idxs j))- >> go (j + 1) (i + 1)+ | VU.unsafeIndex defs i == maxDef = do+ let d = VU.unsafeIndex idxs j+ stToIO+ ( appendTextSliceFromPtr+ builder+ (p `plusPtr` VU.unsafeIndex starts d)+ (VU.unsafeIndex lens d)+ )+ go (j + 1) (i + 1) | otherwise = stToIO (appendNull builder) >> go j (i + 1) go 0 0 Just d -> error ("appendNullableStringPageIO: wrong dict type, got " ++ show d)@@ -441,35 +408,24 @@ error "appendNullableStringPageIO: dictionary-encoded page but no dictionary seen" --- ------------------------------------------------------------------------------ Page header parsing--- ---------------------------------------------------------------------------- parsePageHeader :: BS.ByteString -> Either String (BS.ByteString, PageHeader) parsePageHeader = decodeWithLeftovers Pinch.compactProtocol --- ------------------------------------------------------------------------------ Batch value readers--- ---------------------------------------------------------------------------- readNInt32 :: Int -> BS.ByteString -> VU.Vector Int32-readNInt32 n bs = VU.generate n $ \i -> littleEndianInt32 (BS.drop (4 * i) bs)+readNInt32 n bs = VU.generate n $ \i -> fromIntegral (word32At bs (4 * i)) readNInt64 :: Int -> BS.ByteString -> VU.Vector Int64-readNInt64 n bs = VU.generate n $ \i ->- fromIntegral (littleEndianWord64 (BS.drop (8 * i) bs))+readNInt64 n bs = VU.generate n $ \i -> fromIntegral (word64At bs (8 * i)) readNInt96 :: Int -> BS.ByteString -> [UTCTime] readNInt96 0 _ = [] readNInt96 n bs = int96ToUTCTime (BS.take 12 bs) : readNInt96 (n - 1) (BS.drop 12 bs) readNFloat :: Int -> BS.ByteString -> VU.Vector Float-readNFloat n bs = VU.generate n $ \i ->- castWord32ToFloat (littleEndianWord32 (BS.drop (4 * i) bs))+readNFloat n bs = VU.generate n $ \i -> castWord32ToFloat (word32At bs (4 * i)) readNDouble :: Int -> BS.ByteString -> VU.Vector Double-readNDouble n bs = VU.generate n $ \i ->- castWord64ToDouble (littleEndianWord64 (BS.drop (8 * i) bs))+readNDouble n bs = VU.generate n $ \i -> castWord64ToDouble (word64At bs (8 * i)) readNTexts :: Int -> BS.ByteString -> [T.Text] readNTexts 0 _ = []@@ -483,3 +439,100 @@ readNFixedTexts len n bs = decodeUtf8Lenient (BS.take len bs) : readNFixedTexts len (n - 1) (BS.drop len bs)++-- | Decode a required numeric column chunk straight into @out[off ..]@.+decodeRequiredNumericInto ::+ forall a m.+ (VU.Unbox a, RandomAccess m, MonadIO m) =>+ (DictVals -> Maybe (VU.Vector a)) ->+ (BS.ByteString -> Int -> a) ->+ Int ->+ ColumnDescription ->+ ColumnChunk ->+ VUM.IOVector a ->+ Int ->+ m Int+decodeRequiredNumericInto pick plainAt width description chunk out off0 = do+ (codec, pType, rawBytes) <- readChunkBytes chunk+ goPages Nothing codec pType rawBytes off0+ where+ goPages ::+ Maybe (VU.Vector a) ->+ CompressionCodec ->+ ThriftType ->+ BS.ByteString ->+ Int ->+ m Int+ goPages dict codec pType bs !off+ | BS.null bs = return off+ | otherwise = case parsePageHeader bs of+ Left e -> error ("decodeRequiredNumericInto: 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+ (pageData, rest') = BS.splitAt compSz rest+ case unField hdr.ph_type of+ DICTIONARY_PAGE _ -> do+ let dictHdr =+ fromMaybe+ (error "DICTIONARY_PAGE: missing dictionary page header")+ (unField hdr.ph_dictionary_page_header)+ numVals = unField dictHdr.diph_num_values+ decompressed <- liftIO $ decompressData uncmpSz codec pageData+ let d = readDictVals pType decompressed numVals description.typeLength+ case pick d of+ Nothing -> error "decodeRequiredNumericInto: dictionary page mistyped"+ Just !ds -> goPages (Just ds) codec pType rest' off+ DATA_PAGE _ -> do+ let dph =+ fromMaybe+ (error "DATA_PAGE: missing data page header")+ (unField hdr.ph_data_page_header)+ n = fromIntegral . unField $ dph.dph_num_values+ enc = unField dph.dph_encoding+ decompressed <- liftIO $ decompressData uncmpSz codec pageData+ liftIO (writePage dict enc n decompressed off)+ goPages dict codec pType rest' (off + n)+ DATA_PAGE_V2 _ -> do+ let dph2 =+ fromMaybe+ (error "DATA_PAGE_V2: missing data page header v2")+ (unField hdr.ph_data_page_header_v2)+ n = fromIntegral . unField $ dph2.dph2_num_values+ enc = unField dph2.dph2_encoding+ levelBytes =+ fromIntegral (unField dph2.dph2_definition_levels_byte_length)+ + fromIntegral (unField dph2.dph2_repetition_levels_byte_length)+ isCompressed = fromMaybe True (unField dph2.dph2_is_compressed)+ compValBytes = BS.drop levelBytes pageData+ valBytes <-+ if isCompressed+ then liftIO $ decompressData (uncmpSz - levelBytes) codec compValBytes+ else pure compValBytes+ liftIO (writePage dict enc n valBytes off)+ goPages dict codec pType rest' (off + n)+ INDEX_PAGE _ -> goPages dict codec pType rest' off++ writePage ::+ Maybe (VU.Vector a) -> Encoding -> Int -> BS.ByteString -> Int -> IO ()+ writePage dict enc n vals off = case enc of+ PLAIN _ ->+ let go !i+ | i >= n = pure ()+ | otherwise = do+ VUM.unsafeWrite out (off + i) (plainAt vals (width * i))+ go (i + 1)+ in go 0+ RLE_DICTIONARY _ -> gather+ PLAIN_DICTIONARY _ -> gather+ _ -> error ("decodeRequiredNumericInto: unsupported encoding " ++ show enc)+ where+ gather = case dict of+ Nothing -> error "decodeRequiredNumericInto: dictionary page missing"+ Just ds -> case BS.uncons vals of+ Nothing -> error "decodeRequiredNumericInto: empty index stream"+ Just (w0, rest) ->+ stToIO $+ rleHybridFold (fromIntegral w0) n rest $ \i w ->+ VUM.unsafeWrite out (off + i) (VU.unsafeIndex ds (fromIntegral w))+{-# INLINE decodeRequiredNumericInto #-}
src/DataFrame/IO/Parquet/Schema.hs view
@@ -3,15 +3,7 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TypeApplications #-} -{- |-Module : DataFrame.IO.Parquet.Schema-License : MIT--Helpers for converting a parquet schema (a list of 'SchemaElement') into an-empty 'DataFrame' whose columns have the right types but no rows. Used by-'DataFrame.TH.declareColumnsFromParquetFile' and by any tooling that wants-to inspect a parquet schema as a 'DataFrame'.--}+-- | Converts a parquet schema into an empty 'DataFrame'. module DataFrame.IO.Parquet.Schema ( schemaToEmptyDataFrame, schemaElemToColumn,@@ -33,21 +25,16 @@ unField, ) import DataFrame.Internal.Column (Column, fromList)-import DataFrame.Internal.DataFrame (DataFrame, fromNamedColumns)+import DataFrame.Internal.DataFrame (DataFrame, fromColumns) -{- | Build an empty 'DataFrame' from a flat list of parquet 'SchemaElement's.-Only leaf elements (those with no children) become columns. Columns whose-name is in @nullableCols@ are typed as @Maybe a@; the rest are typed as @a@.--}+-- | Build an empty 'DataFrame' from a flat list of 'SchemaElement's. schemaToEmptyDataFrame :: S.Set T.Text -> [SchemaElement] -> DataFrame schemaToEmptyDataFrame nullableCols elems = let leafElems = filter (\e -> Maybe.fromMaybe 0 (unField e.num_children) == 0) elems- in fromNamedColumns (map (schemaElemToColumn nullableCols) leafElems)+ in fromColumns (map (schemaElemToColumn nullableCols) leafElems) -{- | Convert a single parquet 'SchemaElement' into a named empty 'Column',-picking a nullable or non-nullable representation based on @nullableCols@.--}+-- | Convert a single 'SchemaElement' into a named empty 'Column'. schemaElemToColumn :: S.Set T.Text -> SchemaElement -> (T.Text, Column) schemaElemToColumn nullableCols element = let colName = unField element.name
src/DataFrame/IO/Parquet/Seeking.hs view
@@ -1,7 +1,4 @@-{- | This module contains low-level utilities around file seeking.--later this module can be renamed / moved to an internal module.--}+-- | Low-level utilities around file seeking. module DataFrame.IO.Parquet.Seeking ( SeekableHandle (getSeekableHandle), SeekMode (..),@@ -22,28 +19,10 @@ import Data.Int import System.IO -{- | This handle carries a proof that it must be seekable.-Note: Handle and SeekableHandle are not thread safe, should not be-shared across threads, beaware when running parallel/concurrent code.--Not seekable:- - stdin / stdout- - pipes / FIFOs--But regular files are always seekable. Parquet fundamentally wants random-access, a non-seekable source will not support effecient access without-buffering the entire file.--}+-- | A handle that carries a proof that it is seekable. newtype SeekableHandle = SeekableHandle {getSeekableHandle :: Handle} -{- | If we truely want to support non-seekable files, we need to also consider the case-to buffer the entire file in memory.--Not thread safe, contains mutable reference (as Handle already is).--If we need concurrent / parallel parsing or something, we need to read into ByteString-first, not sharing the same handle.--}+-- | Handle that is either seekable or buffered in memory. data FileBufferedOrSeekable = FileBuffered !(IORef Int64) !BS.ByteString | FileSeekable !SeekableHandle@@ -57,9 +36,7 @@ -- | For testing only type ForceNonSeekable = Maybe Bool -{- | Smart constructor for FileBufferedOrSeekable, tries to keep in the seekable case-if possible.--}+-- | Smart constructor that stays in the seekable case when possible. mkFileBufferedOrSeekable :: ForceNonSeekable -> Handle -> IO FileBufferedOrSeekable mkFileBufferedOrSeekable forceNonSeek h = do@@ -68,11 +45,7 @@ then FileBuffered <$> newIORef 0 <*> BS.hGetContents h else pure $ FileSeekable $ SeekableHandle h -{- | With / bracket pattern for FileBufferedOrSeekable--Warning: do not return the FileBufferedOrSeekable outside the scope of the action as-it will be closed.--}+-- | Bracket-style access to a 'FileBufferedOrSeekable'. withFileBufferedOrSeekable :: ForceNonSeekable -> FilePath ->@@ -83,10 +56,7 @@ fbos <- mkFileBufferedOrSeekable forceNonSeek h action fbos -{- | 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.--}+-- | Read the last @n@ bytes, e.g. to get metadata without the whole file. readLastBytes :: Integer -> FileBufferedOrSeekable -> IO BS.ByteString readLastBytes n (FileSeekable sh) = do let h = getSeekableHandle sh
src/DataFrame/IO/Parquet/Thrift.hs view
@@ -12,8 +12,7 @@ import Pinch (Enumeration, Field, Pinchable (..)) import qualified Pinch --- Primitive Parquet Types--- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L32+-- Primitive Parquet Types (parquet.thrift#L32) data ThriftType = BOOLEAN (Enumeration 0) | INT32 (Enumeration 1)@@ -27,7 +26,7 @@ instance Pinchable ThriftType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L183+-- parquet.thrift#L183 data FieldRepetitionType = REQUIRED (Enumeration 0) | OPTIONAL (Enumeration 1)@@ -36,11 +35,11 @@ instance Pinchable FieldRepetitionType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L203+-- parquet.thrift#L203 data Encoding = PLAIN (Enumeration 0) | -- GROUP_VAR_INT Encoding was never used- -- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L578+ -- parquet.thrift#L578 PLAIN_DICTIONARY (Enumeration 2) | RLE (Enumeration 3) | BIT_PACKED (Enumeration 4)@@ -53,7 +52,7 @@ instance Pinchable Encoding --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L244+-- parquet.thrift#L244 data CompressionCodec = UNCOMPRESSED (Enumeration 0) | SNAPPY (Enumeration 1)@@ -67,7 +66,7 @@ instance Pinchable CompressionCodec --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L261+-- parquet.thrift#L261 data PageType = DATA_PAGE (Enumeration 0) | INDEX_PAGE (Enumeration 1)@@ -77,7 +76,7 @@ instance Pinchable PageType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L271+-- parquet.thrift#L271 data BoundaryOrder = UNORDERED (Enumeration 0) | ASCENDING (Enumeration 1)@@ -86,12 +85,9 @@ instance Pinchable BoundaryOrder --- Logical type annotations--- Empty structs can't use deriving Generic with Pinch, so we use a unit-like workaround.--- We represent empty structs as a newtype over () with a manual Pinchable instance.+-- Logical type annotations; empty structs use manual Pinchable instances. --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L283--- struct StringType {}+-- parquet.thrift#L283: struct StringType {} data StringType = StringType deriving (Eq, Show) instance Pinchable StringType where type Tag StringType = Pinch.TStruct@@ -158,7 +154,7 @@ pinch _ = Pinch.struct [] unpinch _ = pure VariantType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L290+-- parquet.thrift#L290 data TimeUnit = MILLIS (Field 1 MilliSeconds) | MICROS (Field 2 MicroSeconds)@@ -185,7 +181,7 @@ pinch _ = Pinch.struct [] unpinch _ = pure NanoSeconds --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L317+-- parquet.thrift#L317 data DecimalType = DecimalType { decimal_scale :: Field 1 Int32@@ -195,7 +191,7 @@ instance Pinchable DecimalType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L328+-- parquet.thrift#L328 data IntType = IntType { int_bitWidth :: Field 1 Int8@@ -205,7 +201,7 @@ instance Pinchable IntType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L338+-- parquet.thrift#L338 data TimeType = TimeType { time_isAdjustedToUTC :: Field 1 Bool@@ -215,7 +211,7 @@ instance Pinchable TimeType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L349+-- parquet.thrift#L349 data TimestampType = TimestampType { timestamp_isAdjustedToUTC :: Field 1 Bool@@ -225,8 +221,7 @@ instance Pinchable TimestampType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L360--- union LogicalType+-- parquet.thrift#L360: union LogicalType data LogicalType = LT_STRING (Field 1 StringType) | LT_MAP (Field 2 MapType)@@ -247,7 +242,7 @@ instance Pinchable LogicalType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L270+-- parquet.thrift#L270 data ConvertedType = UTF8 (Enumeration 0) | MAP (Enumeration 1)@@ -275,10 +270,10 @@ instance Pinchable ConvertedType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L505+-- parquet.thrift#L505 data SchemaElement = SchemaElement- { schematype :: Field 1 (Maybe ThriftType) -- called just type in parquet.thrift+ { schematype :: Field 1 (Maybe ThriftType) -- named type upstream , type_length :: Field 2 (Maybe Int32) , repetition_type :: Field 3 (Maybe FieldRepetitionType) , name :: Field 4 Text@@ -293,7 +288,7 @@ instance Pinchable SchemaElement --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L560+-- parquet.thrift#L560 data Statistics = Statistics { stats_max :: Field 1 (Maybe ByteString)@@ -309,7 +304,7 @@ instance Pinchable Statistics --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L600+-- parquet.thrift#L600 data PageEncodingStats = PageEncodingStats { pes_page_type :: Field 1 PageType@@ -320,7 +315,7 @@ instance Pinchable PageEncodingStats --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L614+-- parquet.thrift#L614 data ColumnMetaData = ColumnMetaData { cmd_type :: Field 1 ThriftType@@ -343,14 +338,14 @@ instance Pinchable ColumnMetaData --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L875+-- parquet.thrift#L875 data EncryptionWithFooterKey = EncryptionWithFooterKey deriving (Eq, Show) instance Pinchable EncryptionWithFooterKey where type Tag EncryptionWithFooterKey = Pinch.TStruct pinch _ = Pinch.struct [] unpinch _ = pure EncryptionWithFooterKey --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L883+-- parquet.thrift#L883 data EncryptionWithColumnKey = EncryptionWithColumnKey { ewck_path_in_schema :: Field 1 [Text]@@ -360,8 +355,7 @@ instance Pinchable EncryptionWithColumnKey --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L893--- union ColumnCryptoMetaData+-- parquet.thrift#L893: union ColumnCryptoMetaData data ColumnCryptoMetaData = CCM_ENCRYPTION_WITH_FOOTER_KEY (Field 1 EncryptionWithFooterKey) | CCM_ENCRYPTION_WITH_COLUMN_KEY (Field 2 EncryptionWithColumnKey)@@ -369,7 +363,7 @@ instance Pinchable ColumnCryptoMetaData --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L899+-- parquet.thrift#L899 data ColumnChunk = ColumnChunk { cc_file_path :: Field 1 (Maybe Text)@@ -386,7 +380,7 @@ instance Pinchable ColumnChunk --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L940+-- parquet.thrift#L940 data SortingColumn = SortingColumn { sc_column_idx :: Field 1 Int32@@ -397,7 +391,7 @@ instance Pinchable SortingColumn --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L958+-- parquet.thrift#L958 data RowGroup = RowGroup { rg_columns :: Field 1 [ColumnChunk]@@ -412,7 +406,7 @@ instance Pinchable RowGroup --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L980+-- parquet.thrift#L980 data KeyValue = KeyValue { kv_key :: Field 1 Text@@ -422,8 +416,7 @@ instance Pinchable KeyValue --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L990--- union ColumnOrder+-- parquet.thrift#L990: union ColumnOrder newtype ColumnOrder = TYPE_ORDER (Field 1 TypeDefinedOrder) deriving (Eq, Show, Generic)@@ -437,7 +430,7 @@ pinch _ = Pinch.struct [] unpinch _ = pure TypeDefinedOrder --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1094+-- parquet.thrift#L1094 data AesGcmV1 = AesGcmV1 { aes_gcm_v1_aad_prefix :: Field 1 (Maybe ByteString)@@ -448,7 +441,7 @@ instance Pinchable AesGcmV1 --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1107+-- parquet.thrift#L1107 data AesGcmCtrV1 = AesGcmCtrV1 { aes_gcm_ctr_v1_aad_prefix :: Field 1 (Maybe ByteString)@@ -459,8 +452,7 @@ instance Pinchable AesGcmCtrV1 --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1118--- union EncryptionAlgorithm+-- parquet.thrift#L1118: union EncryptionAlgorithm data EncryptionAlgorithm = AES_GCM_V1 (Field 1 AesGcmV1) | AES_GCM_CTR_V1 (Field 2 AesGcmCtrV1)@@ -468,7 +460,7 @@ instance Pinchable EncryptionAlgorithm --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1001+-- parquet.thrift#L1001 data PageLocation = PageLocation { pl_offset :: Field 1 Int64@@ -479,7 +471,7 @@ instance Pinchable PageLocation --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1017+-- parquet.thrift#L1017 data OffsetIndex = OffsetIndex { oi_page_locations :: Field 1 [PageLocation]@@ -489,7 +481,7 @@ instance Pinchable OffsetIndex --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1033+-- parquet.thrift#L1033 data ColumnIndex = ColumnIndex { ci_null_pages :: Field 1 [Bool]@@ -504,7 +496,7 @@ instance Pinchable ColumnIndex --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1248+-- parquet.thrift#L1248 data DataPageHeader = DataPageHeader { dph_num_values :: Field 1 Int32@@ -563,7 +555,7 @@ instance Pinchable PageHeader --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1277+-- parquet.thrift#L1277 data FileMetadata = FileMetadata { version :: Field 1 Int32
src/DataFrame/IO/Parquet/Time.hs view
@@ -42,8 +42,7 @@ year = 100 * b + d - 4800 + m `div` 10 in fromGregorian year (fromIntegral month) (fromIntegral day) --- I include this here even though it's unused because we'll likely use--- it for the writer. Since int96 is deprecated this is only included for completeness anyway.+-- Unused for now; kept for the writer since int96 is deprecated. utcTimeToInt96 :: UTCTime -> BS.ByteString utcTimeToInt96 (UTCTime day diffTime) = let julianDay = dayToJulianDay day
src/DataFrame/IO/Parquet/Utils.hs view
@@ -48,11 +48,7 @@ ) import DataFrame.Internal.Column.Bitmap (buildBitmapFromValid) -{- | 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.--}+-- | A left-fold driver over a column's per-page triples. type PageFold m v a = forall acc. (acc -> (v a, VU.Vector Int, VU.Vector Int) -> m acc) -> acc -> m acc@@ -161,12 +157,7 @@ childLeaves = go children subPath False in childLeaves ++ go rest path skipThis -{- | Fold a column's value pages into a non-nullable 'Column'.--Pre-allocates a mutable vector of @totalRows@ and fills it page-by-page via a-single streaming left fold ('PageFold'), avoiding any intermediate list or-concatenation allocation. Only one page's values are live at a time.--}+-- | Fold a column's value pages into a non-nullable 'Column'. {-# INLINEABLE foldNonNullable #-} foldNonNullable :: forall m a.@@ -200,29 +191,14 @@ 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+ VU.unsafeCopy (VUM.unsafeSlice off n mv) chunk return (off + n) ) (0 :: Int) dat <- liftIO $ VU.unsafeFreeze mv return (UnboxedColumn Nothing dat) -{- | Fold a column's (values, def-levels) pages into a nullable 'Column'.--Pre-allocates the output buffer and a valid-mask vector of @totalRows@, then-scatters values inline during a single streaming left fold ('PageFold').--A 'hasNull' flag is accumulated during the scatter so the-'buildBitmapFromValid' call is skipped entirely when all values are present.--}+-- | Fold a column's value and def-level pages into a nullable 'Column'. {-# INLINEABLE foldNullable #-} foldNullable :: forall m a.@@ -232,13 +208,7 @@ PageFold m VB.Vector a -> m Column foldNullable maxDef totalRows runPages = do- -- null slots hold an error thunk, guarded by bitmap.- --- -- IMPORTANT: 'VBM.unsafeWrite' for boxed vectors stores a *pointer* to- -- the value without evaluating it, so unsupported-encoding error thunks- -- would be silently swallowed into the column data and only fire lazily- -- when user code reads a cell. The '!v' bang pattern forces each value- -- to WHNF before the write, surfacing decoder errors immediately.+ -- The '!v' bang forces each value so decoder errors surface immediately. mvDat <- liftIO $ VBM.replicate totalRows (error "parquet: null slot accessed") mvValid <- liftIO (VUM.new totalRows :: IO (VUM.IOVector Word8))@@ -303,17 +273,7 @@ else return Nothing return (UnboxedColumn maybeBm dat) -{- | Fold a stream of (values, def-levels, rep-levels) triples into a-repeated (list) 'Column' using Dremel-style level stitching.--The stitching function is selected by @maxRep@:-- * @maxRep == 1@ → 'stitchList' → @[Maybe [Maybe a]]@- * @maxRep == 2@ → 'stitchList2' → @[Maybe [Maybe [Maybe a]]]@- * @maxRep >= 3@ → 'stitchList3' → @[Maybe [Maybe [Maybe [Maybe a]]]]@--Threshold formula: @defT_r = maxDef - 2 * (maxRep - r)@.--}+-- | Fold page triples into a repeated (list) 'Column' by level stitching. {-# INLINEABLE foldRepeated #-} foldRepeated :: forall m a.@@ -363,10 +323,7 @@ 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).--}+-- | Collect all of a column's page triples into a list, in page order. collectPages :: (Monad m) => PageFold m v a ->
src/DataFrame/IO/Parquet/Writer.hs view
@@ -12,12 +12,15 @@ nativeTypeKeyValues, ) where +import Control.Exception (evaluate) import Control.Monad (forM_, unless, when)+import Control.Monad.ST (RealWorld, ST, stToIO) import qualified Data.ByteString as BS import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef) import Data.Int (Int64) import Data.Maybe (fromJust) import Data.Primitive.ByteArray (getSizeofMutableByteArray)+import Data.STRef (STRef, modifySTRef', newSTRef, readSTRef, writeSTRef) import qualified Data.Text as T import qualified Data.Vector as VB import DataFrame.IO.Parquet.Thrift hiding (schema)@@ -47,12 +50,12 @@ WritableBinaryHandle, atomicallyWriteFile, bufferResidency,- bufferToByteString, ensureCapacity, flushBufferToBuffer, flushBufferToFile, mallocBuffer, resetPosition,+ unsafeBufferView, withWritableBinaryFile, writeByteString, writeByteStringToFile,@@ -73,27 +76,27 @@ data ParquetWriterState = ParquetWriterState { outputFileHandle :: !WritableBinaryHandle- , columnChunks :: !(VB.Vector ColumnChunkState)+ , columnChunks :: !(VB.Vector (ColumnChunkState RealWorld)) , currentFileOffsetRef :: !(IORef Int64)- , scratchBuffer :: !MemoryBuffer+ , scratchBuffer :: !(MemoryBuffer RealWorld) , rowGroupMetadataRef :: !(IORef [RowGroup]) , rowNumberRef :: !(IORef Int) } -data ColumnChunkState = ColumnChunkState+data ColumnChunkState s = ColumnChunkState { columnName :: !T.Text , nullable :: !Bool , schema :: !SchemaElement- , encoder :: !Encoder- , buffer :: !MemoryBuffer- , uncompressedBufferSize :: !(IORef Int64)- , pageState :: !PageState+ , encoder :: !(Encoder s)+ , buffer :: !(MemoryBuffer s)+ , uncompressedBufferSize :: !(STRef s Int64)+ , pageState :: !(PageState s) } -data PageState = PageState- { pageBuffer :: !MemoryBuffer- , definitionLevels :: !DefLevels- , currentRowCount :: !(IORef Int)+data PageState s = PageState+ { pageBuffer :: !(MemoryBuffer s)+ , definitionLevels :: !(DefLevels s)+ , currentRowCount :: !(STRef s Int) } writeParquet :: FilePath -> DataFrame -> IO ()@@ -144,21 +147,27 @@ -- | Write rows @[startRow, endRow)@ of the frame to a single Parquet file. writeShard ::- ParquetWriteOptions -> FilePath -> DataFrame -> Int -> Int -> IO ()+ ParquetWriteOptions ->+ FilePath ->+ DataFrame ->+ Int ->+ Int ->+ IO () writeShard options path_ df startRow endRow = do let names = columnNames df shardRows = max 0 (endRow - startRow) columnChunks_ <-- VB.fromList- <$> mapM- ( \columnName_ ->- initColumnChunkState- options- columnName_- (fromJust (getColumn columnName_ df))- )- names- scratchBuffer_ <- mallocBuffer (max 1 options.pageSize)+ stToIO $+ VB.fromList+ <$> mapM+ ( \columnName_ ->+ initColumnChunkState+ options+ columnName_+ (fromJust (getColumn columnName_ df))+ )+ names+ scratchBuffer_ <- stToIO (mallocBuffer (max 1 options.pageSize)) atomicallyWriteFile path_ $ \path -> withWritableBinaryFile path $ \output -> do writeByteStringToFile output magic currentFileOffsetRef_ <- newIORef 4@@ -179,7 +188,9 @@ | rowNum >= batchEnd = pure () | otherwise = do let count = min subBatch (batchEnd - rowNum)- VB.forM_ columnChunks_ (writeRows options scratchBuffer_ rowNum count)+ VB.forM_ columnChunks_ $ \ccs -> do+ pageFull <- stToIO (writeRows options rowNum count ccs)+ when pageFull (flushPage options scratchBuffer_ ccs) modifyIORef' rowNumberRef_ (+ count) writeBatch (rowNum + count) batchEnd loop :: Int -> IO ()@@ -188,8 +199,9 @@ | otherwise = do let batchEnd = rowNum + min interval (endRow - rowNum) writeBatch rowNum batchEnd- size <- bufferedSize columnChunks_- when (size >= options.rowGroupSize) (flushRowGroup options writerState)+ size <- stToIO (bufferedSize columnChunks_)+ when (size >= options.rowGroupSize) $+ flushRowGroup options writerState loop batchEnd loop startRow flushRowGroup options writerState@@ -215,26 +227,30 @@ , Just col <- [getColumn name df] ] +-- | Encode @count@ rows from @firstRow@ into the column's page. writeRows ::- ParquetWriteOptions -> MemoryBuffer -> Int -> Int -> ColumnChunkState -> IO ()-writeRows options scratch firstRow count ccs = do+ ParquetWriteOptions -> Int -> Int -> ColumnChunkState s -> ST s Bool+writeRows options firstRow count ccs = do let page = ccs.pageState buf = page.pageBuffer encode = ccs.encoder.encodeValue dl = page.definitionLevels end = firstRow + count - pos0 <- readIORef buf.positionRef+ pos0 <- readSTRef buf.positionRef let margin = options.pageSize arr0 <- ensureCapacity buf (pos0 + max margin (count * 64)) size0 <- getSizeofMutableByteArray arr0 let go !size !pos !row- | row >= end = writeIORef buf.positionRef pos+ | row >= end = writeSTRef buf.positionRef pos | pos + margin > size = do -- Rare: buffer nearly full, grow it- writeIORef buf.positionRef pos- arr' <- ensureCapacity buf (pos + max margin ((end - row) * 64))+ writeSTRef buf.positionRef pos+ arr' <-+ ensureCapacity+ buf+ (pos + max margin ((end - row) * 64)) size' <- getSizeofMutableByteArray arr' go size' pos row | otherwise = do@@ -246,30 +262,35 @@ go size0 pos0 firstRow -- Batch bookkeeping: once per sub-batch instead of per value- modifyIORef' page.currentRowCount (+ count)+ modifySTRef' page.currentRowCount (+ count) flushDef dl pageRes <- bufferResidency buf defRes <- bufferResidency dl.dlBuf- when- (pageRes + defRes >= options.pageSize)- (flushPage options scratch ccs)+ pure (pageRes + defRes >= options.pageSize) -flushPage :: ParquetWriteOptions -> MemoryBuffer -> ColumnChunkState -> IO ()+flushPage ::+ ParquetWriteOptions ->+ MemoryBuffer RealWorld ->+ ColumnChunkState RealWorld ->+ IO () flushPage options scratch columnChunkState = do let page = columnChunkState.pageState- numPageRows <- readIORef page.currentRowCount+ numPageRows <- stToIO (readSTRef page.currentRowCount) when (numPageRows > 0) $ do- pos <- readIORef page.pageBuffer.positionRef- pos' <- columnChunkState.encoder.finishValues page.pageBuffer pos- writeIORef page.pageBuffer.positionRef pos'- body <- assemblePageBody scratch columnChunkState+ body <- stToIO $ do+ pos <- readSTRef page.pageBuffer.positionRef+ pos' <- columnChunkState.encoder.finishValues page.pageBuffer pos+ writeSTRef page.pageBuffer.positionRef pos'+ assemblePageBody scratch columnChunkState writeDataPage options.compressionCodec numPageRows body columnChunkState- resetPosition page.pageBuffer- resetPosition page.definitionLevels.dlBuf- resetPosition scratch- writeIORef page.currentRowCount 0+ stToIO $ do+ resetPosition page.pageBuffer+ resetPosition page.definitionLevels.dlBuf+ resetPosition scratch+ writeSTRef page.currentRowCount 0 -assemblePageBody :: MemoryBuffer -> ColumnChunkState -> IO MemoryBuffer+assemblePageBody ::+ MemoryBuffer s -> ColumnChunkState s -> ST s (MemoryBuffer s) assemblePageBody scratch columnChunkState | not columnChunkState.nullable = pure columnChunkState.pageState.pageBuffer | otherwise = do@@ -283,12 +304,19 @@ pure scratch writeDataPage ::- CompressionCodec -> Int -> MemoryBuffer -> ColumnChunkState -> IO ()+ CompressionCodec ->+ Int ->+ MemoryBuffer RealWorld ->+ ColumnChunkState RealWorld ->+ IO () writeDataPage codec numPageRows body columnChunkState = do- uncompressedPageSize <- bufferResidency body+ uncompressedPageSize <- stToIO (bufferResidency body) compressedBody <- case codec of UNCOMPRESSED _ -> pure Nothing- SNAPPY _ -> Just . Snappy.compress <$> bufferToByteString body+ SNAPPY _ -> do+ -- Compress now: the page buffer is reused for the next page.+ page <- unsafeBufferView body+ Just <$> evaluate (Snappy.compress page) other -> error ("writeParquet: unsupported codec " <> show other) let compressedPageSize = maybe uncompressedPageSize BS.length compressedBody headerBytes =@@ -297,11 +325,12 @@ (mkDataPageHeader numPageRows uncompressedPageSize compressedPageSize) writeByteString columnChunkState.buffer headerBytes case compressedBody of- Nothing -> flushBufferToBuffer body columnChunkState.buffer+ Nothing -> stToIO (flushBufferToBuffer body columnChunkState.buffer) Just bytes -> writeByteString columnChunkState.buffer bytes- modifyIORef'- columnChunkState.uncompressedBufferSize- (+ fromIntegral (BS.length headerBytes + uncompressedPageSize))+ stToIO $+ modifySTRef'+ columnChunkState.uncompressedBufferSize+ (+ fromIntegral (BS.length headerBytes + uncompressedPageSize)) flushRowGroup :: ParquetWriteOptions -> ParquetWriterState -> IO () flushRowGroup options writerState = do@@ -314,13 +343,18 @@ VB.foldM' ( \(acc, totalCompressedSize, totalUncompressedSize) columnChunkState -> do offset <- readIORef writerState.currentFileOffsetRef- compressedSize <- bufferResidency columnChunkState.buffer- uncompressedSize <- readIORef columnChunkState.uncompressedBufferSize- flushBufferToFile writerState.outputFileHandle columnChunkState.buffer+ (compressedSize, uncompressedSize) <-+ stToIO $+ (,)+ <$> bufferResidency columnChunkState.buffer+ <*> readSTRef columnChunkState.uncompressedBufferSize+ flushBufferToFile+ writerState.outputFileHandle+ columnChunkState.buffer writeIORef writerState.currentFileOffsetRef (offset + fromIntegral compressedSize)- writeIORef columnChunkState.uncompressedBufferSize 0+ stToIO (writeSTRef columnChunkState.uncompressedBufferSize 0) let columnChunk = mkColumnChunk options.compressionCodec@@ -349,20 +383,22 @@ ) writeIORef writerState.rowNumberRef 0 -bufferedSize :: VB.Vector ColumnChunkState -> IO Int+bufferedSize :: VB.Vector (ColumnChunkState s) -> ST s Int bufferedSize = VB.foldM' ( \total columnChunkState -> do chunkSize <- bufferResidency columnChunkState.buffer- valuesSize <- bufferResidency columnChunkState.pageState.pageBuffer+ valuesSize <-+ bufferResidency columnChunkState.pageState.pageBuffer defLevelsSize <-- bufferResidency columnChunkState.pageState.definitionLevels.dlBuf+ bufferResidency+ columnChunkState.pageState.definitionLevels.dlBuf pure (total + chunkSize + valuesSize + defLevelsSize) ) 0 initColumnChunkState ::- ParquetWriteOptions -> T.Text -> Column -> IO ColumnChunkState+ ParquetWriteOptions -> T.Text -> Column -> ST s (ColumnChunkState s) initColumnChunkState options columnName_ column = do encoder_ <- buildEncoder column let nullable_ = hasMissing column@@ -374,20 +410,9 @@ encoder_.convertedType encoder_.logicalType bufferSize = max 1 options.pageSize- -- ColumnChunk Buffers start at page size and grow to their- -- actual size over the course of building out the first row- -- group.- -- Each column chunk in a row group must have the same number- -- of rows, but each column chunk is liable to fit the same- -- number of rows in varying amounts of data depending on the- -- encoding and the compression characteristics of the data.- -- So the optimal buffer size of each column chunk is liable- -- to vary- -- As a result while one specific column chunk in a row group- -- is likely to hit the page limit, the others are liable to be- -- much smaller than the limit.+ -- Buffers start at page size and grow, as chunk sizes vary by encoding. buffer_ <- mallocBuffer bufferSize- uncompressedBufferSize_ <- newIORef 0+ uncompressedBufferSize_ <- newSTRef 0 pageState_ <- initPageState bufferSize pure ColumnChunkState@@ -400,11 +425,11 @@ , pageState = pageState_ } -initPageState :: Int -> IO PageState+initPageState :: Int -> ST s (PageState s) initPageState bufferSize = do pageBuffer_ <- mallocBuffer bufferSize definitionLevels_ <- newDefLevels- currentRowCount_ <- newIORef 0+ currentRowCount_ <- newSTRef 0 pure PageState { pageBuffer = pageBuffer_
src/DataFrame/IO/Parquet/Writer/DefLevels.hs view
@@ -8,51 +8,52 @@ ) where import Control.Monad (when)+import Control.Monad.ST (ST) import Data.Bits (shiftL, shiftR, (.&.), (.|.))-import Data.IORef (IORef, newIORef, readIORef, writeIORef)+import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef) import Data.Word (Word64) import DataFrame.IO.Utils.RandomAccess (MemoryBuffer, mallocBuffer, writeWord8) -data DefLevels = DefLevels- { dlBuf :: !MemoryBuffer- , dlValue :: !(IORef Int)- , dlCount :: !(IORef Int)+data DefLevels s = DefLevels+ { dlBuf :: !(MemoryBuffer s)+ , dlValue :: !(STRef s Int)+ , dlCount :: !(STRef s Int) } -newDefLevels :: IO DefLevels-newDefLevels = DefLevels <$> mallocBuffer 64 <*> newIORef 0 <*> newIORef 0+newDefLevels :: ST s (DefLevels s)+newDefLevels = DefLevels <$> mallocBuffer 64 <*> newSTRef 0 <*> newSTRef 0 -pushDef :: DefLevels -> Int -> IO ()+pushDef :: DefLevels s -> Int -> ST s () pushDef dl value = do- count <- readIORef dl.dlCount+ count <- readSTRef dl.dlCount if count == 0- then writeIORef dl.dlValue value >> writeIORef dl.dlCount 1+ then writeSTRef dl.dlValue value >> writeSTRef dl.dlCount 1 else do- current <- readIORef dl.dlValue+ current <- readSTRef dl.dlValue if current == value- then writeIORef dl.dlCount (count + 1)+ then writeSTRef dl.dlCount (count + 1) else do writeDefRun dl current count- writeIORef dl.dlValue value- writeIORef dl.dlCount 1+ writeSTRef dl.dlValue value+ writeSTRef dl.dlCount 1 {-# INLINE pushDef #-} -flushDef :: DefLevels -> IO ()+flushDef :: DefLevels s -> ST s () flushDef dl = do- count <- readIORef dl.dlCount+ count <- readSTRef dl.dlCount when (count > 0) $ do- value <- readIORef dl.dlValue+ value <- readSTRef dl.dlValue writeDefRun dl value count- writeIORef dl.dlCount 0+ writeSTRef dl.dlCount 0 {-# INLINE flushDef #-} -writeDefRun :: DefLevels -> Int -> Int -> IO ()+writeDefRun :: DefLevels s -> Int -> Int -> ST s () writeDefRun dl value count = do writeLeb128 dl.dlBuf (fromIntegral (count `shiftL` 1)) writeWord8 dl.dlBuf (fromIntegral value) {-# INLINE writeDefRun #-} -writeLeb128 :: MemoryBuffer -> Word64 -> IO ()+writeLeb128 :: MemoryBuffer s -> Word64 -> ST s () writeLeb128 buffer value | value < 0x80 = writeWord8 buffer (fromIntegral value) | otherwise = do
src/DataFrame/IO/Parquet/Writer/Encoder.hs view
@@ -10,14 +10,11 @@ buildEncoder, ) where -import Control.Monad.ST (stToIO)+import Control.Monad.ST (ST) import Data.Bits (shiftL, (.|.))-import Data.IORef (newIORef, readIORef, writeIORef) import Data.Int (Int32, Int64)-import Data.Primitive.ByteArray (- withMutableByteArrayContents,- writeByteArray,- )+import Data.Primitive.ByteArray (withMutableByteArrayContents, writeByteArray)+import Data.STRef (newSTRef, readSTRef, writeSTRef) import qualified Data.Text as T import qualified Data.Text.Array as TA import Data.Text.Internal (Text (Text))@@ -55,153 +52,136 @@ import Pinch (enum, putField) import Type.Reflection (typeRep) -data Encoder = Encoder+data Encoder s = Encoder { encType :: !ThriftType , convertedType :: !(Maybe ConvertedType) , logicalType :: !(Maybe LogicalType)- , encodeValue :: !(MemoryBuffer -> Int -> Int -> IO (Int, Bool))- , finishValues :: !(MemoryBuffer -> Int -> IO Int)+ , encodeValue :: !(MemoryBuffer s -> Int -> Int -> ST s (Int, Bool))+ , finishValues :: !(MemoryBuffer s -> Int -> ST s Int) } -buildEncoder :: Column -> IO Encoder+buildEncoder :: Column -> ST s (Encoder s) buildEncoder col | hasElemType @Int32 col = pure $- scalarEncoder @Int32+ scalarEncoder (INT32 enum) Nothing Nothing- (\buffer pos v -> writeWord32At buffer pos (fromIntegral v) >> pure (pos + 4))- col+ ( unboxedColumnWriter @Int32+ col+ (\buffer pos v -> write32 buffer pos (fromIntegral v))+ ) | hasElemType @Int64 col = pure $- scalarEncoder @Int64+ scalarEncoder (INT64 enum) Nothing Nothing- (\buffer pos v -> writeWord64At buffer pos (fromIntegral v) >> pure (pos + 8))- col- -- Ints in GHC can be 32 bit or 64 bit integers depending on the- -- underlying computers architecture. So we'll do 64bit integers- -- to cover all our bases+ ( unboxedColumnWriter @Int64+ col+ (\buffer pos v -> write64 buffer pos (fromIntegral v))+ )+ -- Int is 32 or 64 bit by platform, so always write 64-bit integers. | hasElemType @Int col = pure $- scalarEncoder @Int+ scalarEncoder (INT64 enum) Nothing Nothing- (\buffer pos v -> writeWord64At buffer pos (fromIntegral v) >> pure (pos + 8))- col+ ( unboxedColumnWriter @Int+ col+ (\buffer pos v -> write64 buffer pos (fromIntegral v))+ ) | hasElemType @Integer col = pure $- scalarEncoder @Integer+ scalarEncoder (INT64 enum) Nothing Nothing- writeInteger64At- col+ (columnWriter @Integer col writeInteger64At) | hasElemType @Float col = pure $- scalarEncoder @Float+ scalarEncoder (FLOAT enum) Nothing Nothing- ( \buffer pos v -> writeWord32At buffer pos (castFloatToWord32 v) >> pure (pos + 4)+ ( unboxedColumnWriter @Float+ col+ (\buffer pos v -> write32 buffer pos (castFloatToWord32 v)) )- col | hasElemType @Double col = pure $- scalarEncoder @Double+ scalarEncoder (DOUBLE enum) Nothing Nothing- ( \buffer pos v -> writeWord64At buffer pos (castDoubleToWord64 v) >> pure (pos + 8)+ ( unboxedColumnWriter @Double+ col+ (\buffer pos v -> write64 buffer pos (castDoubleToWord64 v)) )- col | hasElemType @Bool col = boolEncoder col | hasElemType @T.Text col = pure (textEncoder col) | hasElemType @UTCTime col = pure (timestampEncoder col) | otherwise = error ("writeParquet: unsupported column type " <> columnTypeString col)+ where+ write32 buffer pos value = do+ writeWord32At buffer pos value+ pure (pos + 4)+ write64 buffer pos value = do+ writeWord64At buffer pos value+ pure (pos + 8) +type ColumnWriter s = MemoryBuffer s -> Int -> Int -> ST s (Int, Bool)+ scalarEncoder ::- forall a.- (Columnable a) => ThriftType -> Maybe ConvertedType -> Maybe LogicalType ->- (MemoryBuffer -> Int -> a -> IO Int) ->- Column ->- Encoder-scalarEncoder tt conv logical writePrim col =- Encoder tt conv logical (columnWriter @a col writePrim) (\_ pos -> pure pos)-{-# INLINEABLE scalarEncoder #-}-{-# SPECIALIZE scalarEncoder ::- ThriftType ->- Maybe ConvertedType ->- Maybe LogicalType ->- (MemoryBuffer -> Int -> Int32 -> IO Int) ->- Column ->- Encoder- #-}-{-# SPECIALIZE scalarEncoder ::- ThriftType ->- Maybe ConvertedType ->- Maybe LogicalType ->- (MemoryBuffer -> Int -> Int64 -> IO Int) ->- Column ->- Encoder- #-}-{-# SPECIALIZE scalarEncoder ::- ThriftType ->- Maybe ConvertedType ->- Maybe LogicalType ->- (MemoryBuffer -> Int -> Float -> IO Int) ->- Column ->- Encoder- #-}-{-# SPECIALIZE scalarEncoder ::- ThriftType ->- Maybe ConvertedType ->- Maybe LogicalType ->- (MemoryBuffer -> Int -> Double -> IO Int) ->- Column ->- Encoder- #-}-{-# SPECIALIZE scalarEncoder ::- ThriftType ->- Maybe ConvertedType ->- Maybe LogicalType ->- (MemoryBuffer -> Int -> Int -> IO Int) ->+ ColumnWriter s ->+ Encoder s+scalarEncoder tt conv logical encode =+ Encoder tt conv logical encode (\_ pos -> pure pos)++-- | Encode a column whose element type has no 'VU.Unbox' instance.+columnWriter ::+ forall a s.+ (Columnable a) => Column ->- Encoder- #-}-{-# SPECIALIZE scalarEncoder ::- ThriftType ->- Maybe ConvertedType ->- Maybe LogicalType ->- (MemoryBuffer -> Int -> Integer -> IO Int) ->+ (MemoryBuffer s -> Int -> a -> ST s Int) ->+ ColumnWriter s+columnWriter = columnWriterWith Nothing+{-# INLINE columnWriter #-}++-- | Encode a column whose element type has a 'VU.Unbox' instance.+unboxedColumnWriter ::+ forall a s.+ (Columnable a, VU.Unbox a) => Column ->- Encoder- #-}+ (MemoryBuffer s -> Int -> a -> ST s Int) ->+ ColumnWriter s+unboxedColumnWriter = columnWriterWith (Just VU.unsafeIndex)+{-# INLINE unboxedColumnWriter #-} -columnWriter ::- forall a.+columnWriterWith ::+ forall a s. (Columnable a) =>+ Maybe (VU.Vector a -> Int -> a) -> Column ->- (MemoryBuffer -> Int -> a -> IO Int) ->- MemoryBuffer ->- Int ->- Int ->- IO (Int, Bool)-columnWriter col writePrim = case col of+ (MemoryBuffer s -> Int -> a -> ST s Int) ->+ ColumnWriter s+columnWriterWith unboxedIndex col writePrim = case col of BoxedColumn bitmap (values :: VB.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of Just Refl -> writeFrom bitmap (VB.unsafeIndex values) Nothing -> mismatch UnboxedColumn bitmap (values :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of- Just Refl -> writeFrom bitmap (VU.unsafeIndex values)+ Just Refl ->+ writeFrom+ bitmap+ (maybe (VU.unsafeIndex values) ($ values) unboxedIndex) Nothing -> mismatch _ -> mismatch where@@ -213,114 +193,56 @@ mismatch = error ("writeParquet: incompatible column representation for " <> columnTypeString col)-{-# INLINEABLE columnWriter #-}-{-# SPECIALIZE columnWriter ::- Column ->- (MemoryBuffer -> Int -> Int32 -> IO Int) ->- MemoryBuffer ->- Int ->- Int ->- IO (Int, Bool)- #-}-{-# SPECIALIZE columnWriter ::- Column ->- (MemoryBuffer -> Int -> Int64 -> IO Int) ->- MemoryBuffer ->- Int ->- Int ->- IO (Int, Bool)- #-}-{-# SPECIALIZE columnWriter ::- Column ->- (MemoryBuffer -> Int -> Float -> IO Int) ->- MemoryBuffer ->- Int ->- Int ->- IO (Int, Bool)- #-}-{-# SPECIALIZE columnWriter ::- Column ->- (MemoryBuffer -> Int -> Double -> IO Int) ->- MemoryBuffer ->- Int ->- Int ->- IO (Int, Bool)- #-}-{-# SPECIALIZE columnWriter ::- Column ->- (MemoryBuffer -> Int -> Bool -> IO Int) ->- MemoryBuffer ->- Int ->- Int ->- IO (Int, Bool)- #-}-{-# SPECIALIZE columnWriter ::- Column ->- (MemoryBuffer -> Int -> UTCTime -> IO Int) ->- MemoryBuffer ->- Int ->- Int ->- IO (Int, Bool)- #-}-{-# SPECIALIZE columnWriter ::- Column ->- (MemoryBuffer -> Int -> Int -> IO Int) ->- MemoryBuffer ->- Int ->- Int ->- IO (Int, Bool)- #-}-{-# SPECIALIZE columnWriter ::- Column ->- (MemoryBuffer -> Int -> Integer -> IO Int) ->- MemoryBuffer ->- Int ->- Int ->- IO (Int, Bool)- #-}+{-# INLINE columnWriterWith #-} isPresent :: Maybe Bitmap -> Int -> Bool isPresent Nothing _ = True isPresent (Just bitmap) row = bitmapTestBit bitmap row {-# INLINE isPresent #-} -boolEncoder :: Column -> IO Encoder+boolEncoder :: Column -> ST s (Encoder s) boolEncoder col = do- bitsRef <- newIORef (0 :: Word8)- countRef <- newIORef (0 :: Int)+ bitsRef <- newSTRef (0 :: Word8)+ countRef <- newSTRef (0 :: Int) let addBit buffer pos value = do- bits <- readIORef bitsRef- count <- readIORef countRef+ bits <- readSTRef bitsRef+ count <- readSTRef countRef let bits' = if value then bits .|. ((1 :: Word8) `shiftL` count) else bits count' = count + 1 if count' == 8 then do- arr <- readIORef buffer.arrayRef+ arr <- readSTRef buffer.arrayRef writeByteArray arr pos bits'- writeIORef bitsRef 0- writeIORef countRef 0+ writeSTRef bitsRef 0+ writeSTRef countRef 0 pure (pos + 1) else do- writeIORef bitsRef bits'- writeIORef countRef count'+ writeSTRef bitsRef bits'+ writeSTRef countRef count' pure pos finish buffer pos = do- count <- readIORef countRef+ count <- readSTRef countRef pos' <- if count > 0 then do- bits <- readIORef bitsRef- arr <- readIORef buffer.arrayRef+ bits <- readSTRef bitsRef+ arr <- readSTRef buffer.arrayRef writeByteArray arr pos bits pure (pos + 1) else pure pos- writeIORef bitsRef 0- writeIORef countRef 0+ writeSTRef bitsRef 0+ writeSTRef countRef 0 pure pos' pure- (Encoder (BOOLEAN enum) Nothing Nothing (columnWriter @Bool col addBit) finish)+ ( Encoder+ (BOOLEAN enum)+ Nothing+ Nothing+ (unboxedColumnWriter @Bool col addBit)+ finish+ ) -textEncoder :: Column -> Encoder+textEncoder :: Column -> Encoder s textEncoder col = Encoder (BYTE_ARRAY enum)@@ -351,18 +273,18 @@ pure (pos', True) | otherwise = pure (pos, False) writeTextSlice buffer pos bytes offset count = do- writeIORef buffer.positionRef pos+ writeSTRef buffer.positionRef pos _ <- ensureCapacity buffer (pos + 4 + count) writeWord32At buffer pos (fromIntegral count)- arr <- readIORef buffer.arrayRef+ arr <- readSTRef buffer.arrayRef withMutableByteArrayContents arr $ \ptr ->- stToIO (TA.copyToPointer bytes offset (ptr `plusPtr` (pos + 4)) count)+ TA.copyToPointer bytes offset (ptr `plusPtr` (pos + 4)) count pure (pos + 4 + count) mismatch = error ("writeParquet: incompatible text representation for " <> columnTypeString col) -timestampEncoder :: Column -> Encoder+timestampEncoder :: Column -> Encoder s timestampEncoder col = Encoder (INT64 enum)
src/DataFrame/IO/Parquet/Writer/Metadata.hs view
@@ -10,6 +10,7 @@ magic, ) where +import Control.Monad.ST (stToIO) import qualified Data.ByteString as BS import Data.Int (Int64) import qualified Data.Text as T@@ -165,9 +166,9 @@ , footer_signing_key_metadata = putField Nothing } footer = Pinch.encode Pinch.compactProtocol metadata- buffer <- mallocBuffer (BS.length footer + 8)+ buffer <- stToIO (mallocBuffer (BS.length footer + 8)) writeByteString buffer footer- writeWord32LE buffer (fromIntegral (BS.length footer))+ stToIO (writeWord32LE buffer (fromIntegral (BS.length footer))) writeByteString buffer magic flushBufferToFile output buffer
src/DataFrame/IO/Utils/RandomAccess.hs view
@@ -1,4 +1,5 @@ {-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE MagicHash #-} {-# LANGUAGE OverloadedRecordDot #-} module DataFrame.IO.Utils.RandomAccess (@@ -8,6 +9,7 @@ ReaderIO (runReaderIO), LocalFile, MMappedFile,+ ChunkBuffer (..), unsafeToByteString, WritableBinaryHandle, openWritableBinaryFile,@@ -29,31 +31,31 @@ writeDoubleLE, bufferResidency, bufferToByteString,+ unsafeBufferView, flushBufferToBuffer, resetPosition, flushBufferToFile, writeByteStringToFile, ) where -import Control.Exception (bracket, bracketOnError, finally)+import Control.Exception (bracket, bracketOnError, finally, throw) import Control.Monad (when) import Control.Monad.IO.Class (MonadIO (..))-import Control.Monad.Primitive (RealWorld)-import Control.Monad.ST (stToIO)+import Control.Monad.ST (RealWorld, ST, stToIO) import Data.Bits (shiftR)-import qualified Data.ByteString as BS-import Data.ByteString.Internal (ByteString (PS), create)+import Data.ByteString.Internal (ByteString (BS, PS), create) import qualified Data.ByteString.Unsafe as BU-import Data.IORef (IORef, newIORef, readIORef, writeIORef) import Data.Int (Int64) import Data.Primitive.ByteArray (- MutableByteArray,+ MutableByteArray (..), copyMutableByteArray, getSizeofMutableByteArray,+ mutableByteArrayContents, newPinnedByteArray, withMutableByteArrayContents, writeByteArray, )+import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef) import qualified Data.Text.Array as TA import qualified Data.Vector.Storable as VS import Data.Word (Word32, Word64, Word8)@@ -65,6 +67,8 @@ ) import Foreign (castForeignPtr, castPtr, copyBytes, plusPtr) import GHC.Float (castDoubleToWord64, castFloatToWord32)+import GHC.ForeignPtr (ForeignPtr (..), ForeignPtrContents (PlainPtr))+import GHC.Ptr (Ptr (..)) import System.Directory (copyPermissions, doesFileExist, removeFile, renameFile) import System.FilePath (takeDirectory) import System.IO (@@ -120,6 +124,14 @@ fGet handle length' readSuffix n = ReaderIO (readLastBytes $ fromIntegral n) +-- | Bytes of one column chunk, already read; @base@ is their file offset.+data ChunkBuffer = ChunkBuffer !Integer !ByteString++instance RandomAccess (ReaderIO ChunkBuffer) where+ readBytes (Range offset' length') = ReaderIO $ \(ChunkBuffer base bs) ->+ pure (BU.unsafeTake length' (BU.unsafeDrop (fromInteger (offset' - base)) bs))+ readSuffix _ = ReaderIO $ \_ -> error "ChunkBuffer: no file footer"+ type MMappedFile = ReaderIO (VS.Vector Word8) -- The instance exists but we don't have the means to mmap the file currently@@ -139,20 +151,7 @@ where (ptr, offset', len) = VS.unsafeToForeignPtr v --- Writer Buffer --------------------------------------------------------------------- Refer to DataFrame.IO.Parquet.Writer for a justification of what we're doing here--- There's some overlap here with what's going on in Seeking.hs, so, if this bothers--- us, eventually someone will have to come back and reconcile the writer buffer--- approach with the reader oriented patterns in Seeking.hs.------ We're using MutableByteArrays here for convenience and because we don't need--- the more powerful abstractions vector provides (which uses ByteArrays internally)------ since we want to use hPutBuf, we're going to need a Ptr, which means are ByteArrya--- must be pinned. Now growing pinned arrays can be problematic, but in the vast majority--- of cases we shouldn't be growing more than once, if that. See the docs for--- Data.Primitive.ByteArray.byteArrayContents.+-- Writer buffers are pinned MutableByteArrays so hPutBuf can take a Ptr. newtype WritableBinaryHandle = WritableBinaryHandle {unHandle :: Handle} @@ -163,7 +162,10 @@ hSetBuffering h NoBuffering pure . WritableBinaryHandle $ h -atomicallyWriteFile :: FilePath -> (FilePath -> IO a) -> IO a+atomicallyWriteFile ::+ FilePath ->+ (FilePath -> IO a) ->+ IO a atomicallyWriteFile path action = bracketOnError openAction@@ -188,86 +190,76 @@ pure tmpFile ) -withWritableBinaryFile :: FilePath -> (WritableBinaryHandle -> IO a) -> IO a+withWritableBinaryFile ::+ FilePath ->+ (WritableBinaryHandle -> IO a) ->+ IO a withWritableBinaryFile filepath = bracket (openWritableBinaryFile filepath) (hClose . unHandle) -data MemoryBuffer = MemoryBuffer- { arrayRef :: !(IORef (MutableByteArray RealWorld))- , positionRef :: !(IORef Int)+data MemoryBuffer s = MemoryBuffer+ { arrayRef :: !(STRef s (MutableByteArray s))+ , positionRef :: !(STRef s Int) } -mallocBuffer :: Int -> IO MemoryBuffer+mallocBuffer :: Int -> ST s (MemoryBuffer s) mallocBuffer capacity- | capacity < 0 = ioError $ userError "mallocBuffer: negative capacity"+ | capacity < 0 = throw (userError "mallocBuffer: negative capacity") | otherwise = do array <- newPinnedByteArray capacity- MemoryBuffer <$> newIORef array <*> newIORef 0+ MemoryBuffer <$> newSTRef array <*> newSTRef 0 --- We're using pinned ByteArrays so we must--- not use the grow function brovided by Data.Primitive--- instead we must alloocate a new pinned ByteArray.--- We might have been worried about heap fragmentation--- because a single pinned object in a 4KB GHC block can--- keep the whole plock alive but oyr buffers will tend to--- be much larger than that.--- But the memory usage will temporarily spike to 2.5x the size of--- the buffer, but it should be fine since the current writer is single threaded--- and grows *should* be rare.--- If it becomes an issue we should start tracking an array of pointers--- to buffers intsead of replacing them wholesale so grwoing a buffer--- is just a matter of adding a new buffer to the array (which we can--- pre-allocate to three elements to begin with and grow it only on the--- off chance that a buffer required more than three grows).-ensureCapacity :: MemoryBuffer -> Int -> IO (MutableByteArray RealWorld)+-- Pinned arrays cannot grow in place, so allocate a new one and copy.+ensureCapacity :: MemoryBuffer s -> Int -> ST s (MutableByteArray s) ensureCapacity buffer needed = do- array <- readIORef buffer.arrayRef+ array <- readSTRef buffer.arrayRef maxSize <- getSizeofMutableByteArray array if needed <= maxSize then pure array else do- position <- readIORef buffer.positionRef+ position <- readSTRef buffer.positionRef grown <- newPinnedByteArray (needed + (needed `div` 2)) copyMutableByteArray grown 0 array 0 position- writeIORef buffer.arrayRef grown+ writeSTRef buffer.arrayRef grown pure grown {-# INLINE ensureCapacity #-} -writeWord8 :: MemoryBuffer -> Word8 -> IO ()+writeWord8 :: MemoryBuffer s -> Word8 -> ST s () writeWord8 buffer b = do- position <- readIORef buffer.positionRef+ position <- readSTRef buffer.positionRef array <- ensureCapacity buffer (position + 1) writeByteArray array position b- writeIORef buffer.positionRef (position + 1)+ writeSTRef buffer.positionRef (position + 1) {-# INLINE writeWord8 #-} -writeByteString :: MemoryBuffer -> ByteString -> IO ()+-- Reading the ByteString's pointer needs IO, so this stays out of ST.+writeByteString :: MemoryBuffer RealWorld -> ByteString -> IO () writeByteString buffer bs = BU.unsafeUseAsCStringLen bs $ \(source, len) -> do- position <- readIORef buffer.positionRef- array <- ensureCapacity buffer (position + len)+ position <- stToIO (readSTRef buffer.positionRef)+ array <- stToIO (ensureCapacity buffer (position + len)) withMutableByteArrayContents array $ \dst -> copyBytes (dst `plusPtr` position) (castPtr source) len- writeIORef buffer.positionRef (position + len)+ stToIO (writeSTRef buffer.positionRef (position + len)) {-# INLINE writeByteString #-} -writeWord32LE :: MemoryBuffer -> Word32 -> IO ()+writeWord32LE :: MemoryBuffer s -> Word32 -> ST s () writeWord32LE buffer w = do- position <- readIORef buffer.positionRef+ position <- readSTRef buffer.positionRef writeWord32At buffer position w- writeIORef buffer.positionRef (position + 4)+ writeSTRef buffer.positionRef (position + 4) {-# INLINE writeWord32LE #-} -writeWord64LE :: MemoryBuffer -> Word64 -> IO ()+writeWord64LE :: MemoryBuffer s -> Word64 -> ST s () writeWord64LE buffer w = do- position <- readIORef buffer.positionRef+ position <- readSTRef buffer.positionRef writeWord64At buffer position w- writeIORef buffer.positionRef (position + 8)+ writeSTRef buffer.positionRef (position + 8) {-# INLINE writeWord64LE #-} -writeWord32At :: MemoryBuffer -> Int -> Word32 -> IO ()+writeWord32At :: MemoryBuffer s -> Int -> Word32 -> ST s () writeWord32At buffer position w = do array <- ensureCapacity buffer (position + 4) writeByteArray array position (fromIntegral w :: Word8)@@ -276,7 +268,7 @@ writeByteArray array (position + 3) (fromIntegral (w `shiftR` 24) :: Word8) {-# INLINE writeWord32At #-} -writeWord64At :: MemoryBuffer -> Int -> Word64 -> IO ()+writeWord64At :: MemoryBuffer s -> Int -> Word64 -> ST s () writeWord64At buffer position w = do array <- ensureCapacity buffer (position + 8) writeByteArray array position (fromIntegral w :: Word8)@@ -289,14 +281,14 @@ writeByteArray array (position + 7) (fromIntegral (w `shiftR` 56) :: Word8) {-# INLINE writeWord64At #-} -writeInteger64 :: MemoryBuffer -> Integer -> IO ()+writeInteger64 :: MemoryBuffer s -> Integer -> ST s () writeInteger64 buffer value = do- position <- readIORef buffer.positionRef+ position <- readSTRef buffer.positionRef newPosition <- writeInteger64At buffer position value- writeIORef buffer.positionRef newPosition+ writeSTRef buffer.positionRef newPosition {-# INLINE writeInteger64 #-} -writeInteger64At :: MemoryBuffer -> Int -> Integer -> IO Int+writeInteger64At :: MemoryBuffer s -> Int -> Integer -> ST s Int writeInteger64At buffer position value | value < toInteger (minBound :: Int64) = outOfRange | value > toInteger (maxBound :: Int64) = outOfRange@@ -305,24 +297,24 @@ pure (position + 8) where outOfRange =- ioError (userError "writeParquet: Integer value is outside the INT64 range")+ throw (userError "writeParquet: Integer value is outside the INT64 range") {-# INLINE writeInteger64At #-} -writeFloatLE :: MemoryBuffer -> Float -> IO ()+writeFloatLE :: MemoryBuffer s -> Float -> ST s () writeFloatLE buffer = writeWord32LE buffer . castFloatToWord32 {-# INLINE writeFloatLE #-} -writeDoubleLE :: MemoryBuffer -> Double -> IO ()+writeDoubleLE :: MemoryBuffer s -> Double -> ST s () writeDoubleLE buffer = writeWord64LE buffer . castDoubleToWord64 {-# INLINE writeDoubleLE #-} -flushBufferToBuffer :: MemoryBuffer -> MemoryBuffer -> IO ()+flushBufferToBuffer :: MemoryBuffer s -> MemoryBuffer s -> ST s () flushBufferToBuffer source destination | source.arrayRef == destination.arrayRef = pure () | otherwise = do- sourceArray <- readIORef source.arrayRef- sourcePosition <- readIORef source.positionRef- destinationPosition <- readIORef destination.positionRef+ sourceArray <- readSTRef source.arrayRef+ sourcePosition <- readSTRef source.positionRef+ destinationPosition <- readSTRef destination.positionRef destinationArray <- ensureCapacity destination (destinationPosition + sourcePosition) copyMutableByteArray@@ -331,51 +323,39 @@ sourceArray 0 sourcePosition- writeIORef destination.positionRef (destinationPosition + sourcePosition)- writeIORef source.positionRef 0+ writeSTRef destination.positionRef (destinationPosition + sourcePosition)+ writeSTRef source.positionRef 0 {-# INLINE flushBufferToBuffer #-} -bufferToByteString :: MemoryBuffer -> IO ByteString+bufferToByteString :: MemoryBuffer RealWorld -> IO ByteString bufferToByteString buffer = do- array <- readIORef buffer.arrayRef- position <- readIORef buffer.positionRef+ array <- stToIO (readSTRef buffer.arrayRef)+ position <- stToIO (readSTRef buffer.positionRef) create position $ \dst -> withMutableByteArrayContents array $ \src -> copyBytes dst (castPtr src) position -bufferResidency :: MemoryBuffer -> IO Int-bufferResidency buffer = readIORef buffer.positionRef+-- | The buffer's contents as a ByteString, without copying.+unsafeBufferView :: MemoryBuffer RealWorld -> IO ByteString+unsafeBufferView buffer = do+ array@(MutableByteArray array#) <- stToIO (readSTRef buffer.arrayRef)+ position <- stToIO (readSTRef buffer.positionRef)+ let Ptr address# = mutableByteArrayContents array+ pure (BS (ForeignPtr address# (PlainPtr array#)) position)++bufferResidency :: MemoryBuffer s -> ST s Int+bufferResidency buffer = readSTRef buffer.positionRef {-# INLINE bufferResidency #-} -resetPosition :: MemoryBuffer -> IO ()-resetPosition buffer = writeIORef buffer.positionRef 0+resetPosition :: MemoryBuffer s -> ST s ()+resetPosition buffer = writeSTRef buffer.positionRef 0 {-# INLINE resetPosition #-} --- I tested write speeds by doing (on Apple Silicon)--- `dd if=/dev/zero of=test bs={$n}k oflag=direct conv=fdatasync--- Results:------ ```--- | block size | data (GiB) | time (s) | GiB/s |--- |------------|------------|-----------|-------|--- | 4k | 4.00 | 2.371 | 1.69 |--- | 8k | 4.00 | 1.486 | 2.69 |--- | 16k | 4.00 | 1.045 | 3.83 |--- | 32k | 4.00 | 0.740 | 5.40 |--- | 64k | 4.00 | 0.675 | 5.92 |--- | 128k | 4.00 | 0.669 | 5.98 |--- | 256k | 4.00 | 0.664 | 6.03 |--- | 512k | 4.00 | 0.670 | 5.97 |--- | 1024k | 4.00 | 0.664 | 6.02 |--- | 4096k | 4.00 | 0.668 | 5.99 |--- ```--- So when writing to a file to minimize syscall overhead while--- trying not to create dirty pages in the kernel page cache, we'll--- be flushing in 256 KiB chunks.-flushBufferToFile :: WritableBinaryHandle -> MemoryBuffer -> IO ()+-- Flush in 256 KiB chunks; larger blocks gave no write speed in dd tests.+flushBufferToFile :: WritableBinaryHandle -> MemoryBuffer RealWorld -> IO () flushBufferToFile (WritableBinaryHandle h) buffer = do- array <- readIORef buffer.arrayRef- position <- readIORef buffer.positionRef+ array <- stToIO (readSTRef buffer.arrayRef)+ position <- stToIO (readSTRef buffer.positionRef) withMutableByteArrayContents array $ \ptr -> do let chunkSize = 262144 go offset@@ -385,7 +365,7 @@ hPutBuf h (ptr `plusPtr` offset) n go (offset + n) go 0- writeIORef buffer.positionRef 0+ stToIO (resetPosition buffer) writeByteStringToFile :: WritableBinaryHandle -> ByteString -> IO () writeByteStringToFile (WritableBinaryHandle h) bs =@@ -399,13 +379,13 @@ go (offset + n) go 0 -appendTextArraySlice :: MemoryBuffer -> TA.Array -> Int -> Int -> IO ()+appendTextArraySlice :: MemoryBuffer s -> TA.Array -> Int -> Int -> ST s () appendTextArraySlice buffer source offset count- | count < 0 = ioError $ userError "appendTextArraySlice: negative length"+ | count < 0 = throw (userError "appendTextArraySlice: negative length") | otherwise = do- position <- readIORef buffer.positionRef+ position <- readSTRef buffer.positionRef array <- ensureCapacity buffer (position + count) withMutableByteArrayContents array $ \destination ->- stToIO (TA.copyToPointer source offset (destination `plusPtr` position) count)- writeIORef buffer.positionRef (position + count)+ TA.copyToPointer source offset (destination `plusPtr` position) count+ writeSTRef buffer.positionRef (position + count) {-# INLINE appendTextArraySlice #-}
src/DataFrame/Typed/IO/Parquet.hs view
@@ -2,21 +2,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{- | Typed Parquet reading.--The reader validates the file against a type-level schema as it loads, supplied-by type application:--@-type Trips = '[ '(\"id\", Int), '(\"fare\", Double)]--trips <- readParquet \@Trips \"trips.parquet\" -- IO (TypedDataFrame Trips)-@--'readParquet' (and 'readParquetWithOpts'\/'readParquetFiles') throw a-'DataFrameException' on schema mismatch; 'readParquetWithError' returns the-mismatch as an 'Either' instead.--}+-- | Typed Parquet reading, validated against a type-level schema. module DataFrame.Typed.IO.Parquet ( readParquet, readParquetWithError,@@ -34,10 +20,7 @@ import DataFrame.Typed.Schema (KnownSchema, schemaColumnNames) import DataFrame.Typed.Types (TypedDataFrame) -{- | Read only the columns @cols@ names, unless the caller selected columns-themselves. Parquet supplies the element types, so the schema contributes the-projection and the freeze checks the types.--}+-- | Read only the columns the schema names, unless the caller chose. schemaOptions :: forall cols. (KnownSchema cols) => ParquetReadOptions -> ParquetReadOptions schemaOptions opts =@@ -45,27 +28,12 @@ { selectedColumns = selectedColumns opts <|> Just (schemaColumnNames @cols) } -{- | Read a Parquet file into a typed DataFrame, throwing on schema mismatch.-Reads only the columns @cols@ names.--==== __Example__-@-ghci> trips <- readParquet \@Trips \"trips.parquet\"-@--}+-- | Read a Parquet file into a typed DataFrame, throwing on mismatch. readParquet :: forall cols. (KnownSchema cols) => FilePath -> IO (TypedDataFrame cols) readParquet = readParquetWithOpts @cols defaultParquetReadOptions -{- | Read a Parquet file, returning a descriptive error on schema mismatch or-a missing column instead of throwing.--==== __Example__-@-ghci> readParquetWithError \@Trips \"trips.parquet\"-Right (TDF ...)-@--}+-- | Read a Parquet file, returning an error on mismatch. readParquetWithError :: forall cols. (KnownSchema cols) =>@@ -78,15 +46,7 @@ Left (e :: SomeException) -> Left (T.pack (show e)) Right df -> freezeWithError df -{- | Read a Parquet file with custom options, throwing on schema mismatch. The-schema still supplies the column selection; an explicit 'selectedColumns'-takes precedence.--==== __Example__-@-ghci> trips <- readParquetWithOpts \@Trips defaultParquetReadOptions{rowRange = Just (0, 10)} \"trips.parquet\"-@--}+-- | Read a Parquet file with custom options, throwing on mismatch. readParquetWithOpts :: forall cols. (KnownSchema cols) =>@@ -95,13 +55,7 @@ Parquet.readParquetWithOpts (schemaOptions @cols opts) path >>= freezeOrThrow @cols -{- | Read a directory\/glob of Parquet files into a typed DataFrame.--==== __Example__-@-ghci> trips <- readParquetFiles \@Trips \".\/data\/trips\/*.parquet\"-@--}+-- | Read a directory\/glob of Parquet files into a typed DataFrame. readParquetFiles :: forall cols. (KnownSchema cols) => FilePath -> IO (TypedDataFrame cols) readParquetFiles path =
stress/DataFrame10GB.hs view
@@ -19,7 +19,7 @@ import DataFrame.Internal.Column (Column (..)) import DataFrame.Internal.Column.Bitmap (Bitmap) import DataFrame.Internal.Data.PackedText (mkPackedContiguous32)-import DataFrame.Internal.DataFrame (DataFrame, fromNamedColumns)+import DataFrame.Internal.DataFrame (DataFrame, fromColumns) stressRows :: Int stressRows = 1_000_000@@ -40,7 +40,7 @@ * fromIntegral (2 * textBytesPerRow + 2 * (4 + 8 + 4 + 8)) stressDataFrame :: DataFrame-stressDataFrame = fromNamedColumns (concatMap columnGroup [0 .. stressGroups - 1])+stressDataFrame = fromColumns (concatMap columnGroup [0 .. stressGroups - 1]) columnGroup :: Int -> [(T.Text, Column)] columnGroup group =
tests/Main.hs view
@@ -1,4 +1,6 @@+{-# LANGUAGE DataKinds #-} {-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-} -- | Tests for the writer-buffer logic in "DataFrame.IO.Utils.RandomAccess". module Main where@@ -12,10 +14,18 @@ import Test.HUnit import Control.Monad (unless)+import Control.Monad.ST (stToIO) import Data.Int (Int32, Int64) import Data.Maybe (fromJust) import qualified Data.Text as T-import DataFrame.IO.Parquet (readParquet, readParquetFiles)+import DataFrame.IO.Parquet (+ ParquetReadOptions (..),+ defaultParquetReadOptions,+ readParquet,+ readParquetFiles,+ readParquetWithOpts,+ _readParquetWithOpts,+ ) import DataFrame.IO.Parquet.Writer ( ParquetWriteOptions (..), defaultParquetWriteOptions,@@ -27,21 +37,24 @@ import DataFrame.Internal.DataFrame ( DataFrame, columnNames,- fromNamedColumns,+ fromColumns, getColumn, )+import DataFrame.Operations.Subset (range, select)+import qualified DataFrame.Typed.IO.Parquet as Typed+import qualified DataFrame.Typed.Operations as DT import System.Directory (listDirectory) directWrites :: Test directWrites = TestCase $ do- buffer <- mallocBuffer 1- writeWord8 buffer 0xaa- writeWord32LE buffer 0x78563412- writeWord64LE buffer 0x0807060504030201- writeFloatLE buffer 1- writeDoubleLE buffer 1+ buffer <- stToIO (mallocBuffer 1)+ stToIO (writeWord8 buffer 0xaa)+ stToIO (writeWord32LE buffer 0x78563412)+ stToIO (writeWord64LE buffer 0x0807060504030201)+ stToIO (writeFloatLE buffer 1)+ stToIO (writeDoubleLE buffer 1) writeByteString buffer (BS.pack [0xfe, 0xff])- residency <- bufferResidency buffer+ residency <- stToIO (bufferResidency buffer) bytes <- bufferToByteString buffer assertEqual "direct write residency" 27 residency assertEqual@@ -80,20 +93,20 @@ directBufferFlush :: Test directBufferFlush = TestCase $ do- source <- mallocBuffer 0- destination <- mallocBuffer 0+ source <- stToIO (mallocBuffer 0)+ destination <- stToIO (mallocBuffer 0) writeByteString destination (BS.pack [1, 2]) writeByteString source (BS.pack [3, 4, 5])- flushBufferToBuffer source destination- sourceResidency <- bufferResidency source+ stToIO (flushBufferToBuffer source destination)+ sourceResidency <- stToIO (bufferResidency source) destinationBytes <- bufferToByteString destination assertEqual "source cleared" 0 sourceResidency assertEqual "destination appended" (BS.pack [1, 2, 3, 4, 5]) destinationBytes- flushBufferToBuffer destination destination+ stToIO (flushBufferToBuffer destination destination) selfFlushedBytes <- bufferToByteString destination assertEqual "self flush is a no-op" destinationBytes selfFlushedBytes- resetPosition destination- destinationResidency <- bufferResidency destination+ stToIO (resetPosition destination)+ destinationResidency <- stToIO (bufferResidency destination) assertEqual "reset position" 0 destinationResidency directFileFlush :: Test@@ -101,11 +114,11 @@ withSystemTempDirectory "dfpq-buffer" $ \dir -> do let outPath = dir </> "out.bin" payload = BS.pack (take 300000 (cycle [0 .. 255]))- buffer <- mallocBuffer 1+ buffer <- stToIO (mallocBuffer 1) writeByteString buffer payload withWritableBinaryFile outPath $ \output -> flushBufferToFile output buffer- residency <- bufferResidency buffer+ residency <- stToIO (bufferResidency buffer) contents <- BS.readFile outPath assertEqual "source cleared after file flush" 0 residency assertEqual "large payload round-trips" payload contents@@ -128,11 +141,25 @@ df' <- readParquet out assertEqual label df df' +-- Row groups of 512 rows give the parallel reader several tasks per column.+parallelMatchesSequential :: String -> FilePath -> Test+parallelMatchesSequential label path = TestCase $+ withSystemTempDirectory "dfpq-parallel" $ \dir -> do+ df <- readParquet path+ let out = dir </> "out.parquet"+ writeParquetWithOptions tinyWriteOpts out df+ sequential <-+ readParquetWithOpts defaultParquetReadOptions{parallel = False} out+ parallelRead <-+ readParquetWithOpts defaultParquetReadOptions{parallel = True} out+ assertEqual (label ++ ": sequential") df sequential+ assertEqual (label ++ ": parallel") df parallelRead+ writerRoundTripLargeText :: Test writerRoundTripLargeText = TestCase $ withSystemTempDirectory "dfpq-writer" $ \dir -> do let payload = T.replicate 4096 "abcdefgh"- df = fromNamedColumns [("text", fromList [payload, "short"])]+ df = fromColumns [("text", fromList [payload, "short"])] firstOut = dir </> "large-text-1.parquet" secondOut = dir </> "large-text-2.parquet" writeParquetWithOptions tinyWriteOpts firstOut df@@ -142,9 +169,7 @@ assertEqual "large text first round-trip" df firstRoundTrip assertEqual "large text second round-trip" df secondRoundTrip -{- | Sharded writes: @maxRowsPerFile@ splits the frame across a glob pattern,-and reading the shards back reproduces the original frame.--}+-- | Sharded writes split the frame and reading them back matches. writerRoundTripSharded :: String -> FilePath -> Int -> Int -> Test writerRoundTripSharded label path rowsPerFile expectedShards = TestCase $ withSystemTempDirectory "dfpq-writer" $ \dir -> do@@ -180,16 +205,12 @@ `catch` (\e -> True <$ evaluate (Prelude.length (show (e :: SomeException)))) unless threw (assertFailure "expected an error for a path without '*'") -{- | Parquet has one 64-bit integer type, so @Int@, @Int64@ and @Integer@-columns all land in the file as @INT64@. The writer stamps the original-Haskell type in the footer so the reader can put it back; without that, a-CSV-inferred @Int@ column silently widens to @Int64@ on a round trip.--}+-- | Int, Int64 and Integer columns round-trip with their Haskell type. writerRoundTripNativeIntTypes :: Test writerRoundTripNativeIntTypes = TestCase $ withSystemTempDirectory "dfpq-writer" $ \dir -> do let df =- fromNamedColumns+ fromColumns [ ("int", fromList [1 :: Int, 2, 3]) , ("int64", fromList [1 :: Int64, 2, 3]) , ("int32", fromList [1 :: Int32, 2, 3])@@ -221,10 +242,45 @@ , subBatchRows = 3 } +dottedColumnNames :: Test+dottedColumnNames = TestCase $+ withSystemTempDirectory "dfpq-dotted" $ \dir -> do+ let df =+ fromColumns+ [ ("sepal.length", fromList [5.1, 4.9, 4.7 :: Double])+ , ("length", fromList [1, 2, 3 :: Int])+ , ("variety", fromList ["a", "b", "c" :: T.Text])+ ]+ out = dir </> "dotted.parquet"+ writeParquet out df+ whole <- readParquet out+ assertEqual "names survive the round trip" (columnNames df) (columnNames whole)+ projected <-+ readParquetWithOpts+ defaultParquetReadOptions{selectedColumns = Just ["sepal.length", "variety"]}+ out+ assertEqual+ "projection keeps the dotted column"+ ["sepal.length", "variety"]+ (columnNames projected)+ assertEqual+ "projected rows match"+ (select ["sepal.length", "variety"] df)+ projected+ leaf <-+ readParquetWithOpts+ defaultParquetReadOptions{selectedColumns = Just ["length"]}+ out+ assertEqual "a full name wins over a leaf alias" ["length"] (columnNames leaf)+ typed <-+ Typed.readParquet @'[ '("sepal.length", Double), '("variety", T.Text)] out+ assertEqual "the typed reader projects a dotted schema" 3 (DT.nRows typed)+ tests :: Test tests = TestList- [ TestLabel "direct buffer writes" directWrites+ [ TestLabel "dotted column names project by full name" dottedColumnNames+ , TestLabel "direct buffer writes" directWrites , TestLabel "direct buffer-to-buffer flush" directBufferFlush , TestLabel "direct buffer-to-file flush" directFileFlush , TestLabel@@ -243,6 +299,12 @@ "writer roundtrip: alltypes_tiny_pages" (writerRoundTrip "alltypes_tiny_pages" "tests/data/alltypes_tiny_pages.parquet") , TestLabel+ "parallel read matches sequential: transactions"+ (parallelMatchesSequential "transactions" "tests/data/transactions.parquet")+ , TestLabel+ "parallel read matches sequential: mtcars"+ (parallelMatchesSequential "mtcars" "tests/data/mtcars.parquet")+ , TestLabel "writer roundtrip: transactions" (writerRoundTrip "transactions" "tests/data/transactions.parquet") , TestLabel@@ -273,7 +335,75 @@ "tests/data/alltypes_plain.parquet" ) , TestLabel "writer roundtrip: large text" writerRoundTripLargeText+ , TestLabel "parallel read: many row groups" parallelReadRowGroups+ , TestLabel "parallel read: row range skips row groups" parallelReadRowRange+ , TestLabel "parallel read: widened integers" parallelReadWidenIntegers ]++-- | A frame big enough for 'tinyWriteOpts' to split into many row groups.+manyRowGroupsFrame :: DataFrame+manyRowGroupsFrame =+ fromColumns+ [ ("n", fromList [i * 7 - 3000 | i <- [0 .. 2999 :: Int]])+ , ("d", fromList [fromIntegral i / 8 :: Double | i <- [0 .. 2999 :: Int]])+ ,+ ( "t"+ , fromList+ [T.pack (show i) <> T.replicate (i `mod` 5) "é" | i <- [0 .. 2999 :: Int]]+ )+ ]++withManyRowGroups :: (FilePath -> IO ()) -> IO ()+withManyRowGroups body = withSystemTempDirectory "dfpq-rg" $ \dir -> do+ let out = dir </> "rg.parquet"+ writeParquetWithOptions tinyWriteOpts out manyRowGroupsFrame+ body out++-- | The parallel reader matches the sequential one and the source frame.+parallelReadRowGroups :: Test+parallelReadRowGroups = TestCase $ withManyRowGroups $ \out -> do+ parallel <- readParquet out+ sequential <- _readParquetWithOpts (Just True) defaultParquetReadOptions out+ assertEqual "parallel vs source" manyRowGroupsFrame parallel+ assertEqual "parallel vs sequential" sequential parallel+ projected <-+ readParquetWithOpts+ defaultParquetReadOptions{selectedColumns = Just ["t", "n"]}+ out+ assertEqual "projection" (select ["t", "n"] manyRowGroupsFrame) projected++-- | A row range only decodes the row groups it overlaps, and slices exactly.+parallelReadRowRange :: Test+parallelReadRowRange = TestCase $ withManyRowGroups $ \out -> do+ let ranges = [(0, 5), (7, 1300), (1299, 1301), (2990, 3000), (2995, 4000), (100, 100)]+ mapM_+ ( \r -> do+ got <- readParquetWithOpts defaultParquetReadOptions{rowRange = Just r} out+ assertEqual ("range " ++ show r) (range r manyRowGroupsFrame) got+ )+ ranges++-- | 'widenIntegers' turns Int32 and Int64 columns into Int on both read paths.+parallelReadWidenIntegers :: Test+parallelReadWidenIntegers = TestCase $ withSystemTempDirectory "dfpq-widen" $ \dir -> do+ let n = 3000 :: Int+ df =+ fromColumns+ [ ("i32", fromList [fromIntegral (i - 1500) :: Int32 | i <- [0 .. n - 1]])+ , ("i64", fromList [fromIntegral i * 4000000000 :: Int64 | i <- [0 .. n - 1]])+ ]+ expected =+ fromColumns+ [ ("i32", fromList [i - 1500 | i <- [0 .. n - 1]])+ , ("i64", fromList [i * 4000000000 | i <- [0 .. n - 1]])+ ]+ out = dir </> "ints.parquet"+ opts = defaultParquetReadOptions{widenIntegers = True}+ writeParquetWithOptions tinyWriteOpts out df+ parallel <- readParquetWithOpts opts out+ sequential <- _readParquetWithOpts (Just True) opts out+ assertEqual "parallel" expected parallel+ assertEqual "sequential" expected sequential main :: IO () main = do