diff --git a/benchmark/Writer10GB.hs b/benchmark/Writer10GB.hs new file mode 100644 --- /dev/null +++ b/benchmark/Writer10GB.hs @@ -0,0 +1,49 @@ +module Main (main) where + +import Control.DeepSeq (NFData (rnf)) +import Criterion.Main (bench, defaultMain, envWithCleanup, whnfIO) +import DataFrame.IO.Parquet.Writer (writeParquet) +import DataFrame.Internal.DataFrame (DataFrame, forceDataFrame) +import DataFrame10GB (stressDataFrame) +import System.Directory (removeDirectoryRecursive) +import System.FilePath (()) +import System.IO.Temp (createTempDirectory, getCanonicalTemporaryDirectory) + +data BenchmarkEnvironment = BenchmarkEnvironment + { benchmarkDataFrame :: DataFrame + , benchmarkDirectory :: FilePath + , benchmarkOutput :: FilePath + } + +instance NFData BenchmarkEnvironment where + rnf environment = + forceDataFrame (benchmarkDataFrame environment) `seq` + rnf (benchmarkDirectory environment) `seq` + rnf (benchmarkOutput environment) + +prepareEnvironment :: IO BenchmarkEnvironment +prepareEnvironment = do + temporary <- getCanonicalTemporaryDirectory + directory <- createTempDirectory temporary "dataframe-parquet-writer-10gb" + pure + BenchmarkEnvironment + { benchmarkDataFrame = stressDataFrame + , benchmarkDirectory = directory + , benchmarkOutput = directory "benchmark.parquet" + } + +cleanupEnvironment :: BenchmarkEnvironment -> IO () +cleanupEnvironment = removeDirectoryRecursive . benchmarkDirectory + +main :: IO () +main = + defaultMain + [ envWithCleanup prepareEnvironment cleanupEnvironment $ \environment -> + -- Memory usage for this benchmark will be north of 20 GB. + bench "write 10 GiB dataframe" $ + whnfIO + ( writeParquet + (benchmarkOutput environment) + (benchmarkDataFrame environment) + ) + ] diff --git a/dataframe-parquet.cabal b/dataframe-parquet.cabal --- a/dataframe-parquet.cabal +++ b/dataframe-parquet.cabal @@ -1,7 +1,6 @@ -cabal-version: 2.4 +cabal-version: 3.4 name: dataframe-parquet -version: 1.1.0.1 - +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 @@ -20,6 +19,7 @@ copyright: (c) 2024-2026 Michael Chavinda category: Data tested-with: GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2 +extra-source-files: tests/data/*.parquet common warnings ghc-options: @@ -29,6 +29,11 @@ -Wunused-local-binds -Wunused-packages +flag stress-tests + description: Build and run the opt-in 10 GiB Parquet roundtrip stress test. + default: False + manual: True + library import: warnings ghc-options: -O2 @@ -45,22 +50,83 @@ DataFrame.IO.Parquet.Thrift DataFrame.IO.Parquet.Time DataFrame.IO.Parquet.Utils + DataFrame.IO.Parquet.Writer + DataFrame.IO.Parquet.Writer.DefLevels + DataFrame.IO.Parquet.Writer.Encoder + DataFrame.IO.Parquet.Writer.Metadata + DataFrame.IO.Parquet.Writer.Options DataFrame.IO.Utils.RandomAccess + DataFrame.Typed.IO.Parquet build-depends: base >= 4 && < 5, - bytestring >= 0.11 && < 0.13, - containers >= 0.6.7 && < 0.9, - dataframe-core ^>= 1.1, - dataframe-operations ^>= 1.1.1, - dataframe-parsing ^>= 1.0.2, + bytestring >= 0.11 && < 0.14, + containers >= 0.6.7 && < 0.10, + 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, filepath >= 1.4 && < 2, Glob >= 0.10 && < 1, pinch >= 0.5 && < 1, - snappy-hs ^>= 0.1, + snappy-hs >= 0.1 && < 0.3, text >= 2.1 && < 3, time >= 1.12 && < 2, - vector ^>= 0.13, + vector >= 0.13 && < 0.15, zlib >= 0.5 && < 1, - zstd >= 0.1.2.0 && < 0.2 + zstd >= 0.1.2.0 && < 0.3 hs-source-dirs: src default-language: Haskell2010 + + +test-suite dataframe-parquet-tests + import: warnings + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: tests + build-depends: base >= 4 && < 5, + bytestring >= 0.11 && < 0.14, + dataframe-core >= 2.6 && < 2.7, + dataframe-operations >= 2.6 && < 2.7, + dataframe-parquet, + directory >= 1.3.0.0 && < 2, + filepath >= 1.4 && < 2, + temporary >= 1.3 && < 1.5, + text >= 2.1 && < 3, + HUnit >= 1.6 && < 1.8 + default-language: Haskell2010 + +executable dataframe-parquet-10gb-stress + import: warnings + main-is: StressMain.hs + other-modules: DataFrame10GB + hs-source-dirs: stress + build-depends: base >= 4 && < 5, + dataframe-core >= 2.6 && < 2.7, + dataframe-parquet, + filepath >= 1.4 && < 2, + temporary >= 1.3 && < 1.5, + text >= 2.1 && < 3, + time >= 1.12 && < 2, + vector >= 0.13 && < 0.15 + default-language: Haskell2010 + -- ghc-options: -O2 -threaded -rtsopts -with-rtsopts=-N + +benchmark dataframe-parquet-writer-10gb + import: warnings + type: exitcode-stdio-1.0 + main-is: Writer10GB.hs + other-modules: DataFrame10GB + hs-source-dirs: benchmark, stress + build-depends: base >= 4 && < 5, + criterion >= 1 && < 2, + deepseq >= 1.4 && < 2, + dataframe-core >= 2.6 && < 2.7, + dataframe-parquet, + directory >= 1.3 && < 2, + filepath >= 1.4 && < 2, + temporary >= 1.3 && < 1.5, + text >= 2.1 && < 3, + time >= 1.12 && < 2, + vector >= 0.13 && < 0.15 + default-language: Haskell2010 + ghc-options: -O2 -threaded -rtsopts -with-rtsopts=-N diff --git a/src/DataFrame/IO/Parquet.hs b/src/DataFrame/IO/Parquet.hs --- a/src/DataFrame/IO/Parquet.hs +++ b/src/DataFrame/IO/Parquet.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MonoLocalBinds #-} @@ -7,16 +8,76 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -module DataFrame.IO.Parquet where +module DataFrame.IO.Parquet ( + -- * Reading + readParquet, + readParquetWithOpts, + readParquetFiles, + readParquetFilesWithOpts, + -- * Writing + writeParquet, + writeParquetWithOptions, + + -- * Options + ParquetReadOptions (..), + defaultParquetReadOptions, + ParquetWriteOptions (..), + WriterStrategy (..), + defaultParquetWriteOptions, + + -- * File metadata + parseFileMetadata, + readMetadataFromPath, + readMetadataFromHandle, + columnChunksForAll, + + -- * Schema description + ColumnDescription (..), + generateColumnDescriptions, + getColumnNames, + + -- * Decoding + parseParquetWithOpts, + parseColumnChunks, + getNonNullableColumn, + getNullableColumn, + getRepeatedColumn, + applyDescLogicalType, + applyLogicalType, + nativeTypeHints, + restoreNativeType, + + -- * Applying read options to a decoded frame + applyReadOptions, + applyPredicate, + applySelectedColumns, + applyRowRange, + applySafeRead, + + -- * Data sources + RandomAccess (..), + ReaderIO (runReaderIO), + FileBufferedOrSeekable, + ForceNonSeekable, + withFileBufferedOrSeekable, + _readParquetWithOpts, + readParquetWithMetadata, + FileMetadata, +) where + +import Control.Concurrent (getNumCapabilities) import Control.Exception (throw) import Control.Monad import Control.Monad.IO.Class (MonadIO (..)) +import Control.Monad.ST (stToIO) import Data.Bits (Bits (shiftL), (.|.)) 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 @@ -25,35 +86,54 @@ import qualified Data.List as L import qualified Data.Map as Map import qualified Data.Text as T -import Data.Time (UTCTime) -import Data.Time.Clock.POSIX (posixSecondsToUTCTime) +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, + appendNullableStringPageIO, + appendStringPageIO, boolDecoder, byteArrayDecoder, + decodeRequiredNumericInto, doubleDecoder, fixedLenByteArrayDecoder, floatDecoder, + foldColumnDataPagesM, + foldColumnPagesM, int32Decoder, int64Decoder, int96Decoder, - foldColumnPagesM, ) import DataFrame.IO.Parquet.Seeking ( - FileBufferedOrSeekable, + FileBufferedOrSeekable (..), ForceNonSeekable, withFileBufferedOrSeekable, ) import DataFrame.IO.Parquet.Thrift ( ColumnChunk (..), + ColumnMetaData (..), DecimalType (..), FileMetadata (..), + KeyValue (..), LogicalType (..), RowGroup (..), + Statistics (..), ThriftType (..), TimeUnit (..), TimestampType (..), @@ -70,62 +150,69 @@ generateColumnDescriptions, getColumnNames, ) +import DataFrame.IO.Parquet.Writer ( + ParquetWriteOptions (..), + WriterStrategy (..), + defaultParquetWriteOptions, + nativeTypeKeyPrefix, + writeParquet, + 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 ( + Schema (..), + SchemaType, + 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 @@ -133,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 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 @@ -198,7 +509,7 @@ matches <- glob pat - files <- filterM (fmap not . doesDirectoryExist) matches + files <- L.sort <$> filterM (fmap not . doesDirectoryExist) matches case files of [] -> @@ -209,99 +520,131 @@ 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 -> m DataFrame +{-# SPECIALIZE parseParquetWithOpts :: + ParquetReadOptions -> ReaderIO FileBufferedOrSeekable DataFrame + #-} +{-# 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] + } - let finalCols = zipWith applyDescLogicalType keptDescs rawCols +-- | 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 + (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 @@ -344,6 +687,19 @@ [ColumnChunk] -> ColumnDescription -> m Column +{-# SPECIALIZE parseColumnChunks :: + Int -> + [ColumnChunk] -> + ColumnDescription -> + ReaderIO FileBufferedOrSeekable Column + #-} +{-# SPECIALIZE parseColumnChunks :: + Int -> + [ColumnChunk] -> + ColumnDescription -> + ReaderIO ChunkBuffer Column + #-} +{-# INLINEABLE parseColumnChunks #-} parseColumnChunks totalRows chunks description | description.maxRepetitionLevel == 0 && description.maxDefinitionLevel == 0 = getNonNullableColumn totalRows description chunks @@ -353,6 +709,7 @@ getRepeatedColumn description chunks -- | Decode a required (non-nullable, non-repeated) column. +{-# INLINEABLE getNonNullableColumn #-} getNonNullableColumn :: forall m. (RandomAccess m, MonadIO m) => @@ -368,7 +725,7 @@ Just (INT96 _) -> go int96Decoder Just (FLOAT _) -> unboxedGo floatDecoder Just (DOUBLE _) -> unboxedGo doubleDecoder - Just (BYTE_ARRAY _) -> go byteArrayDecoder + Just (BYTE_ARRAY _) -> goPackedText Just (FIXED_LEN_BYTE_ARRAY _) -> case description.typeLength of Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type_length to be set" Just tl -> go (fixedLenByteArrayDecoder (fromIntegral tl)) @@ -382,6 +739,22 @@ go decoder = foldNonNullable totalRows (foldColumnPagesM description decoder chunks) + -- Decode non-nullable UTF-8 BYTE_ARRAY straight into one 'PackedText'. + goPackedText :: m Column + goPackedText = do + builder <- + liftIO $ stToIO (newTextBuilder totalRows (textByteHint totalRows chunks)) + _ <- + foldColumnDataPagesM + description + chunks + ( \() (dict, enc, nPresent, valBytes, _, _) -> + liftIO (appendStringPageIO builder dict enc nPresent valBytes) + ) + () + chunk <- liftIO $ stToIO (freezeTextChunk builder) + pure (mergeTextChunks [chunk]) + unboxedGo :: forall a. (Columnable a, VU.Unbox a) => @@ -390,7 +763,18 @@ 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 :: forall m. (RandomAccess m, MonadIO m) => @@ -406,7 +790,7 @@ Just (INT96 _) -> go int96Decoder Just (FLOAT _) -> unboxedGo floatDecoder Just (DOUBLE _) -> unboxedGo doubleDecoder - Just (BYTE_ARRAY _) -> go byteArrayDecoder + Just (BYTE_ARRAY _) -> goPackedTextNullable Just (FIXED_LEN_BYTE_ARRAY _) -> case description.typeLength of Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type_length to be set" Just tl -> go (fixedLenByteArrayDecoder (fromIntegral tl)) @@ -422,15 +806,36 @@ m Column go decoder = foldNullable maxDef totalRows (foldColumnPagesM description decoder chunks) + + -- Nullable UTF-8 BYTE_ARRAY: decode straight into a 'PackedText'. + goPackedTextNullable :: m Column + goPackedTextNullable = do + builder <- + liftIO $ stToIO (newTextBuilder totalRows (textByteHint totalRows chunks)) + _ <- + foldColumnDataPagesM + description + chunks + ( \() (dict, enc, nPresent, valBytes, defs, _) -> + liftIO + (appendNullableStringPageIO builder maxDef dict enc nPresent valBytes defs) + ) + () + chunk <- liftIO $ stToIO (freezeTextChunk builder) + pure (mergeTextChunks [chunk]) unboxedGo :: forall a. (Columnable a, VU.Unbox a) => UnboxedPageDecoder a -> m Column unboxedGo decoder = - foldNullableUnboxed maxDef totalRows (foldColumnPagesM description decoder chunks) + foldNullableUnboxed + maxDef + totalRows + (foldColumnPagesM description decoder chunks) -- | Decode a repeated (list/nested) column. +{-# INLINEABLE getRepeatedColumn #-} getRepeatedColumn :: forall m. (RandomAccess m, MonadIO m) => @@ -481,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) @@ -507,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) @@ -519,34 +918,86 @@ applyLogicalType (Just (LT_TIMESTAMP f)) col = let ts = unField f unit = unField ts.timestamp_unit - divisor = case unit of - MILLIS _ -> 1_000 - MICROS _ -> 1_000_000 - NANOS _ -> 1_000_000_000 - in fromRight col $ - DI.mapColumn - (microsecondsToUTCTime . (* (1_000_000 `div` divisor))) - col + -- 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 + NANOS _ -> epochToUTCTime 1_000_000_000 1_000 + in fromRight col $ DI.mapColumn conv col applyLogicalType (Just (LT_DECIMAL f)) col = let dt = unField f 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 -microsecondsToUTCTime :: Int64 -> UTCTime -microsecondsToUTCTime us = - posixSecondsToUTCTime (fromIntegral us / 1_000_000) +nativeTypeHints :: FileMetadata -> Schema +nativeTypeHints metadata = + makeSchema + [ (name, ty) + | kv <- concat (unField metadata.key_value_metadata) + , Just name <- [T.stripPrefix nativeTypeKeyPrefix (unField kv.kv_key)] + , Just value <- [unField kv.kv_value] + , Just ty <- [Map.lookup value stampedSchemaTypes] + ] + +stampedSchemaTypes :: Map.Map T.Text SchemaType +stampedSchemaTypes = + Map.fromList $ + concat + [ entry @Int "Int" + , entry @Int32 "Int32" + , entry @Int64 "Int64" + , entry @Integer "Integer" + , entry @Float "Float" + , entry @Double "Double" + , entry @Bool "Bool" + , entry @T.Text "Text" + , entry @UTCTime "UTCTime" + ] + where + entry :: + forall a. + (Columnable a, Read a, Columnable (Maybe a)) => + T.Text -> + [(T.Text, SchemaType)] + entry name = + [ (name, schemaType @a) + , ("Maybe " <> name, schemaType @(Maybe a)) + ] + +restoreNativeType :: Schema -> T.Text -> DI.Column -> DI.Column +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 + 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 with the given tick rate to 'UTCTime'. +epochToUTCTime :: Int64 -> Integer -> Int64 -> UTCTime +epochToUTCTime ticksPerSecond psPerTick v = + let (s, subTicks) = v `divMod` ticksPerSecond + (days, sInDay) = s `divMod` 86_400 + ps = fromIntegral sInDay * 1_000_000_000_000 + fromIntegral subTicks * psPerTick + in UTCTime + (ModifiedJulianDay (40_587 + fromIntegral days)) + (picosecondsToDiffTime ps) diff --git a/src/DataFrame/IO/Parquet/Binary.hs b/src/DataFrame/IO/Parquet/Binary.hs --- a/src/DataFrame/IO/Parquet/Binary.hs +++ b/src/DataFrame/IO/Parquet/Binary.hs @@ -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 diff --git a/src/DataFrame/IO/Parquet/Dictionary.hs b/src/DataFrame/IO/Parquet/Dictionary.hs --- a/src/DataFrame/IO/Parquet/Dictionary.hs +++ b/src/DataFrame/IO/Parquet/Dictionary.hs @@ -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 diff --git a/src/DataFrame/IO/Parquet/Encoding.hs b/src/DataFrame/IO/Parquet/Encoding.hs --- a/src/DataFrame/IO/Parquet/Encoding.hs +++ b/src/DataFrame/IO/Parquet/Encoding.hs @@ -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 #-} diff --git a/src/DataFrame/IO/Parquet/Levels.hs b/src/DataFrame/IO/Parquet/Levels.hs --- a/src/DataFrame/IO/Parquet/Levels.hs +++ b/src/DataFrame/IO/Parquet/Levels.hs @@ -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) = diff --git a/src/DataFrame/IO/Parquet/Page.hs b/src/DataFrame/IO/Parquet/Page.hs --- a/src/DataFrame/IO/Parquet/Page.hs +++ b/src/DataFrame/IO/Parquet/Page.hs @@ -17,11 +17,18 @@ fixedLenByteArrayDecoder, -- Page iteration foldColumnPagesM, + foldColumnDataPagesM, + RawPage, + decodeRequiredNumericInto, + appendStringPageIO, + appendNullableStringPageIO, ) where import Control.Monad.IO.Class (MonadIO (liftIO)) -import Data.Bits (shiftR, (.&.)) +import Control.Monad.ST (RealWorld, stToIO) +import Data.Bits (shiftL, shiftR, (.&.), (.|.)) import qualified Data.ByteString as BS +import qualified Data.ByteString.Unsafe as BSU import Data.Int (Int32, Int64) import Data.Maybe (fromJust, fromMaybe) import qualified Data.Text as T @@ -30,12 +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 (..), @@ -53,62 +69,57 @@ 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, + appendTextSliceFromPtr, ) +import Foreign.Ptr (Ptr, castPtr, plusPtr) +import Foreign.Storable (peekByteOff) import GHC.Float (castWord32ToFloat, castWord64ToDouble) 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 _ -> VU.fromList (readNBool nPresent bs) - RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getBool - PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getBool + -- 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 _ -> 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 @@ -117,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 @@ -147,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 @@ -157,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 -> @@ -175,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 :: @@ -208,37 +218,21 @@ rawBytes <- readBytes (Range offset compLen) return (codec, pType, rawBytes) -{- | Left-fold a monadic step over every DATA page (V1 or V2) of every column -chunk, in order, threading the running dictionary internally. - -Replaces the old Streamly @readPages@ 'Unfold' + @unfoldEach@/@fold@ pipeline -with a direct monadic fold: pages are decoded and handed to @step@ one at a -time, so only a single page's decoded values are live at once (the -@unsafeFreeze@-into-preallocated-buffer consumers rely on this — a materialised -page list would roughly double peak memory per column). - -Dictionary pages are consumed silently and update the running dictionary that -is threaded through the recursion. @INDEX_PAGE@s are skipped. - --- TODO: when a page index is available, use it here to compute which page --- byte ranges to request from the RandomAccess layer instead of reading the --- entire column chunk in one contiguous read. +-- | A decoded DATA page handed to a fold step. +type RawPage = + (Maybe DictVals, Encoding, Int, BS.ByteString, VU.Vector Int, VU.Vector Int) --- 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. --} -foldColumnPagesM :: - forall m v a acc. - (RandomAccess m, MonadIO m, VG.Vector v a) => +-- | Left-fold a monadic step over every DATA page of every column chunk. +{-# INLINEABLE foldColumnDataPagesM #-} +foldColumnDataPagesM :: + forall m acc. + (RandomAccess m, MonadIO m) => ColumnDescription -> - (Maybe DictVals -> Encoding -> Int -> BS.ByteString -> v a) -> [ColumnChunk] -> - (acc -> (v a, VU.Vector Int, VU.Vector Int) -> m acc) -> + (acc -> RawPage -> m acc) -> acc -> m acc -foldColumnPagesM description decoder chunks step = goChunks chunks +foldColumnDataPagesM description chunks step = goChunks chunks where maxDef = fromIntegral description.maxDefinitionLevel :: Int maxRep = fromIntegral description.maxRepetitionLevel :: Int @@ -252,7 +246,7 @@ goPages dict codec pType bs !acc | BS.null bs = return acc | otherwise = case parsePageHeader bs of - Left e -> error ("foldColumnPagesM: failed to parse page header: " ++ e) + Left e -> error ("foldColumnDataPagesM: 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 @@ -277,8 +271,7 @@ decompressed <- liftIO $ decompressData uncmpSz codec pageData let (defLvls, repLvls, nPresent, valBytes) = readLevelsV1 n maxDef maxRep decompressed - triple = (decoder dict enc nPresent valBytes, defLvls, repLvls) - acc' <- step acc triple + acc' <- step acc (dict, enc, nPresent, valBytes, defLvls, repLvls) goPages dict codec pType rest' acc' DATA_PAGE_V2 _ -> do let dph2 = @@ -289,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 @@ -298,49 +290,142 @@ if isCompressed then liftIO $ decompressData uncmpSz codec compValBytes else pure compValBytes - let triple = (decoder dict enc nPresent valBytes, defLvls, repLvls) - acc' <- step acc triple + acc' <- step acc (dict, enc, nPresent, valBytes, defLvls, repLvls) goPages dict codec pType rest' acc' INDEX_PAGE _ -> goPages dict codec pType rest' acc --- --------------------------------------------------------------------------- --- Page header parsing --- --------------------------------------------------------------------------- +-- | Left-fold over per-page value triples with a given decoder. +{-# INLINEABLE foldColumnPagesM #-} +foldColumnPagesM :: + forall m v a acc. + (RandomAccess m, MonadIO m, VG.Vector v a) => + ColumnDescription -> + (Maybe DictVals -> Encoding -> Int -> BS.ByteString -> v a) -> + [ColumnChunk] -> + (acc -> (v a, VU.Vector Int, VU.Vector Int) -> m acc) -> + acc -> + m acc +foldColumnPagesM description decoder chunks step = + foldColumnDataPagesM description chunks $ + \acc (dict, enc, nPresent, valBytes, defLvls, repLvls) -> + step acc (decoder dict enc nPresent valBytes, defLvls, repLvls) -parsePageHeader :: BS.ByteString -> Either String (BS.ByteString, PageHeader) -parsePageHeader = decodeWithLeftovers Pinch.compactProtocol +-- | Append a non-nullable BYTE_ARRAY page's strings into a 'TextBuilder'. +appendStringPageIO :: + TextBuilder RealWorld -> + Maybe DictVals -> + Encoding -> + Int -> + BS.ByteString -> + IO () +appendStringPageIO builder mDict enc nPresent bs = case enc of + PLAIN _ -> BSU.unsafeUseAsCStringLen bs $ \(cptr, _len) -> do + let p = castPtr cptr :: Ptr Word8 + go !_ !i | i >= nPresent = pure () + go !off !i = do + len <- readLenLE p off + stToIO (appendTextSliceFromPtr builder (p `plusPtr` (off + 4)) len) + go (off + 4 + len) (i + 1) + go 0 0 + RLE_DICTIONARY _ -> dictAppend + PLAIN_DICTIONARY _ -> dictAppend + _ -> error ("appendStringPageIO: unsupported encoding " ++ show enc) + where + dictAppend = case mDict of + 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" --- --------------------------------------------------------------------------- --- Batch value readers --- --------------------------------------------------------------------------- +-- | 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 + b1 <- peekByteOff p (off + 1) :: IO Word8 + b2 <- peekByteOff p (off + 2) :: IO Word8 + b3 <- peekByteOff p (off + 3) :: IO Word8 + pure $ + fromIntegral b0 + .|. (fromIntegral b1 `shiftL` 8) + .|. (fromIntegral b2 `shiftL` 16) + .|. (fromIntegral b3 `shiftL` 24) +{-# INLINE readLenLE #-} -readNBool :: Int -> BS.ByteString -> [Bool] -readNBool count bs = - let totalBytes = (count + 7) `div` 8 - bits = - concatMap - (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7]) - (BS.unpack (BS.take totalBytes bs)) - in take count bits +-- | Append a nullable BYTE_ARRAY page into a 'TextBuilder'. +appendNullableStringPageIO :: + TextBuilder RealWorld -> + Int -> + Maybe DictVals -> + Encoding -> + Int -> + BS.ByteString -> + VU.Vector Int -> + IO () +appendNullableStringPageIO builder maxDef mDict enc nPresent bs defs = case enc of + PLAIN _ -> BSU.unsafeUseAsCStringLen bs $ \(cptr, _len) -> do + let p = castPtr cptr :: Ptr Word8 + go !_ !i | i >= nDefs = pure () + go !off !i + | VU.unsafeIndex defs i == maxDef = do + len <- readLenLE p off + stToIO (appendTextSliceFromPtr builder (p `plusPtr` (off + 4)) len) + go (off + 4 + len) (i + 1) + | otherwise = stToIO (appendNull builder) >> go off (i + 1) + go 0 0 + RLE_DICTIONARY _ -> dictAppend + PLAIN_DICTIONARY _ -> dictAppend + _ -> error ("appendNullableStringPageIO: unsupported encoding " ++ show enc) + where + nDefs = VU.length defs + dictAppend = case mDict of + 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 = 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) + Nothing -> + error + "appendNullableStringPageIO: dictionary-encoded page but no dictionary seen" +parsePageHeader :: BS.ByteString -> Either String (BS.ByteString, PageHeader) +parsePageHeader = decodeWithLeftovers Pinch.compactProtocol + 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 _ = [] @@ -354,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 #-} diff --git a/src/DataFrame/IO/Parquet/Schema.hs b/src/DataFrame/IO/Parquet/Schema.hs --- a/src/DataFrame/IO/Parquet/Schema.hs +++ b/src/DataFrame/IO/Parquet/Schema.hs @@ -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 diff --git a/src/DataFrame/IO/Parquet/Seeking.hs b/src/DataFrame/IO/Parquet/Seeking.hs --- a/src/DataFrame/IO/Parquet/Seeking.hs +++ b/src/DataFrame/IO/Parquet/Seeking.hs @@ -1,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 diff --git a/src/DataFrame/IO/Parquet/Thrift.hs b/src/DataFrame/IO/Parquet/Thrift.hs --- a/src/DataFrame/IO/Parquet/Thrift.hs +++ b/src/DataFrame/IO/Parquet/Thrift.hs @@ -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 diff --git a/src/DataFrame/IO/Parquet/Time.hs b/src/DataFrame/IO/Parquet/Time.hs --- a/src/DataFrame/IO/Parquet/Time.hs +++ b/src/DataFrame/IO/Parquet/Time.hs @@ -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 diff --git a/src/DataFrame/IO/Parquet/Utils.hs b/src/DataFrame/IO/Parquet/Utils.hs --- a/src/DataFrame/IO/Parquet/Utils.hs +++ b/src/DataFrame/IO/Parquet/Utils.hs @@ -44,15 +44,11 @@ import DataFrame.Internal.Column ( Column (..), Columnable, - buildBitmapFromValid, fromList, ) +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,8 @@ 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. (MonadIO m, Columnable a) => @@ -186,6 +178,7 @@ v <- liftIO $ VB.unsafeFreeze mv return (BoxedColumn Nothing v) +{-# INLINEABLE foldNonNullableUnboxed #-} foldNonNullableUnboxed :: forall m a. (MonadIO m, Columnable a, VU.Unbox a) => @@ -198,29 +191,15 @@ 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. (MonadIO m, Columnable a) => @@ -229,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)) @@ -264,6 +237,7 @@ else return Nothing return (BoxedColumn maybeBm dat) +{-# INLINEABLE foldNullableUnboxed #-} foldNullableUnboxed :: forall m a. (MonadIO m, Columnable a, VU.Unbox a) => @@ -299,17 +273,8 @@ 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. ( MonadIO m @@ -333,6 +298,7 @@ fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef allReps allDefs allVals) _ -> fromList (stitchList maxDef allReps allDefs allVals) +{-# INLINEABLE foldRepeatedUnboxed #-} foldRepeatedUnboxed :: forall m a. ( MonadIO m @@ -357,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 -> diff --git a/src/DataFrame/IO/Parquet/Writer.hs b/src/DataFrame/IO/Parquet/Writer.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Writer.hs @@ -0,0 +1,438 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE OverloadedStrings #-} + +module DataFrame.IO.Parquet.Writer ( + writeParquet, + writeParquetWithOptions, + ParquetWriteOptions (..), + WriterStrategy (..), + defaultParquetWriteOptions, + nativeTypeKeyPrefix, + 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) +import DataFrame.IO.Parquet.Writer.DefLevels ( + DefLevels (..), + flushDef, + newDefLevels, + pushDef, + ) +import DataFrame.IO.Parquet.Writer.Encoder (Encoder (..), buildEncoder) +import DataFrame.IO.Parquet.Writer.Metadata ( + magic, + mkColumnChunk, + mkDataPageHeader, + mkRowGroup, + mkSchemaElem, + rootSchemaElement, + writeFooter, + ) +import DataFrame.IO.Parquet.Writer.Options ( + ParquetWriteOptions (..), + WriterStrategy (..), + defaultParquetWriteOptions, + ) +import DataFrame.IO.Utils.RandomAccess ( + MemoryBuffer (..), + WritableBinaryHandle, + atomicallyWriteFile, + bufferResidency, + ensureCapacity, + flushBufferToBuffer, + flushBufferToFile, + mallocBuffer, + resetPosition, + unsafeBufferView, + withWritableBinaryFile, + writeByteString, + writeByteStringToFile, + writeWord32LE, + ) +import DataFrame.Internal.Column (Column, columnTypeString, hasMissing) +import DataFrame.Internal.DataFrame ( + DataFrame, + columnNames, + dataframeDimensions, + getColumn, + ) +import qualified Pinch +import qualified Snappy +import System.Directory (createDirectoryIfMissing) +import System.FilePath (takeDirectory) +import Text.Printf (printf) + +data ParquetWriterState = ParquetWriterState + { outputFileHandle :: !WritableBinaryHandle + , columnChunks :: !(VB.Vector (ColumnChunkState RealWorld)) + , currentFileOffsetRef :: !(IORef Int64) + , scratchBuffer :: !(MemoryBuffer RealWorld) + , rowGroupMetadataRef :: !(IORef [RowGroup]) + , rowNumberRef :: !(IORef Int) + } + +data ColumnChunkState s = ColumnChunkState + { columnName :: !T.Text + , nullable :: !Bool + , schema :: !SchemaElement + , encoder :: !(Encoder s) + , buffer :: !(MemoryBuffer s) + , uncompressedBufferSize :: !(STRef s Int64) + , pageState :: !(PageState s) + } + +data PageState s = PageState + { pageBuffer :: !(MemoryBuffer s) + , definitionLevels :: !(DefLevels s) + , currentRowCount :: !(STRef s Int) + } + +writeParquet :: FilePath -> DataFrame -> IO () +writeParquet = writeParquetWithOptions defaultParquetWriteOptions + +writeParquetWithOptions :: ParquetWriteOptions -> FilePath -> DataFrame -> IO () +writeParquetWithOptions options path df = do + when (options.strategy == TwoPass) $ + error + "The Two Pass Strategy for the Parquet Writer has not yet been implemented" + case options.compressionCodec of + UNCOMPRESSED _ -> pure () + SNAPPY _ -> pure () + other -> error ("writeParquet: unsupported codec " <> show other) + let (totalRows, _) = dataframeDimensions df + case options.maxRowsPerFile of + Nothing -> do + when (isShardPattern path) $ + error + ( "writeParquet: path " + <> show path + <> " contains a '*' placeholder but maxRowsPerFile is not set" + ) + writeShard options path df 0 totalRows + Just rowsPerFile -> do + when (rowsPerFile <= 0) $ + error "writeParquet: maxRowsPerFile must be positive" + unless (isShardPattern path) $ + error + ( "writeParquet: maxRowsPerFile requires a path with a '*' placeholder, got " + <> show path + ) + let starts = case [0, rowsPerFile .. totalRows - 1] of + [] -> [0] -- empty frame still produces one (empty) shard + ss -> ss + forM_ (zip [0 ..] starts) $ \(shardIndex, start) -> do + let shardPath = shardPathFor path shardIndex + createDirectoryIfMissing True (takeDirectory shardPath) + writeShard options shardPath df start (min totalRows (start + rowsPerFile)) + +isShardPattern :: FilePath -> Bool +isShardPattern = elem '*' + +-- | Replace every @*@ in the pattern with a zero-padded shard index. +shardPathFor :: FilePath -> Int -> FilePath +shardPathFor pattern_ shardIndex = + concatMap (\c -> if c == '*' then printf "%05d" shardIndex else [c]) pattern_ + +-- | Write rows @[startRow, endRow)@ of the frame to a single Parquet file. +writeShard :: + ParquetWriteOptions -> + FilePath -> + DataFrame -> + Int -> + Int -> + IO () +writeShard options path_ df startRow endRow = do + let names = columnNames df + shardRows = max 0 (endRow - startRow) + columnChunks_ <- + 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 + rowGroupMetadataRef_ <- newIORef [] + rowNumberRef_ <- newIORef 0 + let writerState = + ParquetWriterState + output + columnChunks_ + currentFileOffsetRef_ + scratchBuffer_ + rowGroupMetadataRef_ + rowNumberRef_ + interval = max 1 options.batchRows + subBatch = max 1 options.subBatchRows + writeBatch :: Int -> Int -> IO () + writeBatch rowNum batchEnd + | rowNum >= batchEnd = pure () + | otherwise = do + let count = min subBatch (batchEnd - rowNum) + 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 () + loop rowNum + | rowNum >= endRow = pure () + | otherwise = do + let batchEnd = rowNum + min interval (endRow - rowNum) + writeBatch rowNum batchEnd + size <- stToIO (bufferedSize columnChunks_) + when (size >= options.rowGroupSize) $ + flushRowGroup options writerState + loop batchEnd + loop startRow + flushRowGroup options writerState + rowGroupMetadata <- reverse <$> readIORef rowGroupMetadataRef_ + let schemaElements = + rootSchemaElement (VB.length columnChunks_) + : VB.toList (VB.map schema columnChunks_) + writeFooter + output + schemaElements + shardRows + rowGroupMetadata + (nativeTypeKeyValues names df) + +nativeTypeKeyPrefix :: T.Text +nativeTypeKeyPrefix = "dataframe.type." + +-- | The type stamp for every column of @df@, as footer key-value pairs. +nativeTypeKeyValues :: [T.Text] -> DataFrame -> [(T.Text, T.Text)] +nativeTypeKeyValues names df = + [ (nativeTypeKeyPrefix <> name, T.pack (columnTypeString col)) + | name <- names + , Just col <- [getColumn name df] + ] + +-- | Encode @count@ rows from @firstRow@ into the column's page. +writeRows :: + 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 <- 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 = writeSTRef buf.positionRef pos + | pos + margin > size = do + -- Rare: buffer nearly full, grow it + writeSTRef buf.positionRef pos + arr' <- + ensureCapacity + buf + (pos + max margin ((end - row) * 64)) + size' <- getSizeofMutableByteArray arr' + go size' pos row + | otherwise = do + (pos', notNull) <- encode buf pos row + when ccs.nullable $ + pushDef dl (if notNull then 1 else 0) + go size pos' (row + 1) + + go size0 pos0 firstRow + + -- Batch bookkeeping: once per sub-batch instead of per value + modifySTRef' page.currentRowCount (+ count) + flushDef dl + pageRes <- bufferResidency buf + defRes <- bufferResidency dl.dlBuf + pure (pageRes + defRes >= options.pageSize) + +flushPage :: + ParquetWriteOptions -> + MemoryBuffer RealWorld -> + ColumnChunkState RealWorld -> + IO () +flushPage options scratch columnChunkState = do + let page = columnChunkState.pageState + numPageRows <- stToIO (readSTRef page.currentRowCount) + when (numPageRows > 0) $ do + 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 + stToIO $ do + resetPosition page.pageBuffer + resetPosition page.definitionLevels.dlBuf + resetPosition scratch + writeSTRef page.currentRowCount 0 + +assemblePageBody :: + MemoryBuffer s -> ColumnChunkState s -> ST s (MemoryBuffer s) +assemblePageBody scratch columnChunkState + | not columnChunkState.nullable = pure columnChunkState.pageState.pageBuffer + | otherwise = do + let page = columnChunkState.pageState + flushDef page.definitionLevels + resetPosition scratch + defLevelsSize <- bufferResidency page.definitionLevels.dlBuf + writeWord32LE scratch (fromIntegral defLevelsSize) + flushBufferToBuffer page.definitionLevels.dlBuf scratch + flushBufferToBuffer page.pageBuffer scratch + pure scratch + +writeDataPage :: + CompressionCodec -> + Int -> + MemoryBuffer RealWorld -> + ColumnChunkState RealWorld -> + IO () +writeDataPage codec numPageRows body columnChunkState = do + uncompressedPageSize <- stToIO (bufferResidency body) + compressedBody <- case codec of + UNCOMPRESSED _ -> pure Nothing + 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 = + Pinch.encode + Pinch.compactProtocol + (mkDataPageHeader numPageRows uncompressedPageSize compressedPageSize) + writeByteString columnChunkState.buffer headerBytes + case compressedBody of + Nothing -> stToIO (flushBufferToBuffer body columnChunkState.buffer) + Just bytes -> writeByteString columnChunkState.buffer bytes + stToIO $ + modifySTRef' + columnChunkState.uncompressedBufferSize + (+ fromIntegral (BS.length headerBytes + uncompressedPageSize)) + +flushRowGroup :: ParquetWriteOptions -> ParquetWriterState -> IO () +flushRowGroup options writerState = do + rowNumber <- readIORef writerState.rowNumberRef + when (rowNumber > 0) $ do + VB.forM_ + writerState.columnChunks + (flushPage options writerState.scratchBuffer) + (reversedColumnChunks, totalCompressed, totalUncompressed) <- + VB.foldM' + ( \(acc, totalCompressedSize, totalUncompressedSize) columnChunkState -> do + offset <- readIORef writerState.currentFileOffsetRef + (compressedSize, uncompressedSize) <- + stToIO $ + (,) + <$> bufferResidency columnChunkState.buffer + <*> readSTRef columnChunkState.uncompressedBufferSize + flushBufferToFile + writerState.outputFileHandle + columnChunkState.buffer + writeIORef + writerState.currentFileOffsetRef + (offset + fromIntegral compressedSize) + stToIO (writeSTRef columnChunkState.uncompressedBufferSize 0) + let columnChunk = + mkColumnChunk + options.compressionCodec + columnChunkState.encoder.encType + columnChunkState.columnName + offset + compressedSize + uncompressedSize + rowNumber + pure + ( columnChunk : acc + , totalCompressedSize + fromIntegral compressedSize + , totalUncompressedSize + uncompressedSize + ) + ) + ([], 0 :: Int64, 0 :: Int64) + writerState.columnChunks + modifyIORef' + writerState.rowGroupMetadataRef + ( mkRowGroup + (reverse reversedColumnChunks) + totalCompressed + totalUncompressed + rowNumber + : + ) + writeIORef writerState.rowNumberRef 0 + +bufferedSize :: VB.Vector (ColumnChunkState s) -> ST s Int +bufferedSize = + VB.foldM' + ( \total columnChunkState -> do + chunkSize <- bufferResidency columnChunkState.buffer + valuesSize <- + bufferResidency columnChunkState.pageState.pageBuffer + defLevelsSize <- + bufferResidency + columnChunkState.pageState.definitionLevels.dlBuf + pure (total + chunkSize + valuesSize + defLevelsSize) + ) + 0 + +initColumnChunkState :: + ParquetWriteOptions -> T.Text -> Column -> ST s (ColumnChunkState s) +initColumnChunkState options columnName_ column = do + encoder_ <- buildEncoder column + let nullable_ = hasMissing column + schema_ = + mkSchemaElem + columnName_ + encoder_.encType + nullable_ + encoder_.convertedType + encoder_.logicalType + bufferSize = max 1 options.pageSize + -- Buffers start at page size and grow, as chunk sizes vary by encoding. + buffer_ <- mallocBuffer bufferSize + uncompressedBufferSize_ <- newSTRef 0 + pageState_ <- initPageState bufferSize + pure + ColumnChunkState + { columnName = columnName_ + , nullable = nullable_ + , schema = schema_ + , encoder = encoder_ + , buffer = buffer_ + , uncompressedBufferSize = uncompressedBufferSize_ + , pageState = pageState_ + } + +initPageState :: Int -> ST s (PageState s) +initPageState bufferSize = do + pageBuffer_ <- mallocBuffer bufferSize + definitionLevels_ <- newDefLevels + currentRowCount_ <- newSTRef 0 + pure + PageState + { pageBuffer = pageBuffer_ + , definitionLevels = definitionLevels_ + , currentRowCount = currentRowCount_ + } diff --git a/src/DataFrame/IO/Parquet/Writer/DefLevels.hs b/src/DataFrame/IO/Parquet/Writer/DefLevels.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Writer/DefLevels.hs @@ -0,0 +1,62 @@ +{-# LANGUAGE OverloadedRecordDot #-} + +module DataFrame.IO.Parquet.Writer.DefLevels ( + DefLevels (..), + newDefLevels, + pushDef, + flushDef, +) where + +import Control.Monad (when) +import Control.Monad.ST (ST) +import Data.Bits (shiftL, shiftR, (.&.), (.|.)) +import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef) +import Data.Word (Word64) +import DataFrame.IO.Utils.RandomAccess (MemoryBuffer, mallocBuffer, writeWord8) + +data DefLevels s = DefLevels + { dlBuf :: !(MemoryBuffer s) + , dlValue :: !(STRef s Int) + , dlCount :: !(STRef s Int) + } + +newDefLevels :: ST s (DefLevels s) +newDefLevels = DefLevels <$> mallocBuffer 64 <*> newSTRef 0 <*> newSTRef 0 + +pushDef :: DefLevels s -> Int -> ST s () +pushDef dl value = do + count <- readSTRef dl.dlCount + if count == 0 + then writeSTRef dl.dlValue value >> writeSTRef dl.dlCount 1 + else do + current <- readSTRef dl.dlValue + if current == value + then writeSTRef dl.dlCount (count + 1) + else do + writeDefRun dl current count + writeSTRef dl.dlValue value + writeSTRef dl.dlCount 1 +{-# INLINE pushDef #-} + +flushDef :: DefLevels s -> ST s () +flushDef dl = do + count <- readSTRef dl.dlCount + when (count > 0) $ do + value <- readSTRef dl.dlValue + writeDefRun dl value count + writeSTRef dl.dlCount 0 +{-# INLINE flushDef #-} + +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 s -> Word64 -> ST s () +writeLeb128 buffer value + | value < 0x80 = writeWord8 buffer (fromIntegral value) + | otherwise = do + writeWord8 buffer (fromIntegral (value .&. 0x7f) .|. 0x80) + writeLeb128 buffer (value `shiftR` 7) +{-# INLINE writeLeb128 #-} diff --git a/src/DataFrame/IO/Parquet/Writer/Encoder.hs b/src/DataFrame/IO/Parquet/Writer/Encoder.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Writer/Encoder.hs @@ -0,0 +1,316 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module DataFrame.IO.Parquet.Writer.Encoder ( + Encoder (..), + buildEncoder, +) where + +import Control.Monad.ST (ST) +import Data.Bits (shiftL, (.|.)) +import Data.Int (Int32, Int64) +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)) +import Data.Time.Calendar (toModifiedJulianDay) +import Data.Time.Clock (UTCTime (UTCTime), diffTimeToPicoseconds) +import Data.Type.Equality (TestEquality (..), (:~:) (Refl)) +import qualified Data.Vector as VB +import qualified Data.Vector.Unboxed as VU +import Data.Word (Word8) +import DataFrame.IO.Parquet.Thrift +import DataFrame.IO.Utils.RandomAccess ( + MemoryBuffer (..), + ensureCapacity, + writeInteger64At, + writeWord32At, + writeWord64At, + ) +import DataFrame.Internal.Column ( + Column (..), + Columnable, + columnTypeString, + hasElemType, + ) +import DataFrame.Internal.Column.Bitmap ( + Bitmap, + bitmapTestBit, + ) +import DataFrame.Internal.Data.PackedText ( + PackedTextData (..), + offAt, + selAt, + ) +import Foreign (plusPtr) +import GHC.Float (castDoubleToWord64, castFloatToWord32) +import Pinch (enum, putField) +import Type.Reflection (typeRep) + +data Encoder s = Encoder + { encType :: !ThriftType + , convertedType :: !(Maybe ConvertedType) + , logicalType :: !(Maybe LogicalType) + , encodeValue :: !(MemoryBuffer s -> Int -> Int -> ST s (Int, Bool)) + , finishValues :: !(MemoryBuffer s -> Int -> ST s Int) + } + +buildEncoder :: Column -> ST s (Encoder s) +buildEncoder col + | hasElemType @Int32 col = + pure $ + scalarEncoder + (INT32 enum) + Nothing + Nothing + ( unboxedColumnWriter @Int32 + col + (\buffer pos v -> write32 buffer pos (fromIntegral v)) + ) + | hasElemType @Int64 col = + pure $ + scalarEncoder + (INT64 enum) + Nothing + Nothing + ( 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 + (INT64 enum) + Nothing + Nothing + ( unboxedColumnWriter @Int + col + (\buffer pos v -> write64 buffer pos (fromIntegral v)) + ) + | hasElemType @Integer col = + pure $ + scalarEncoder + (INT64 enum) + Nothing + Nothing + (columnWriter @Integer col writeInteger64At) + | hasElemType @Float col = + pure $ + scalarEncoder + (FLOAT enum) + Nothing + Nothing + ( unboxedColumnWriter @Float + col + (\buffer pos v -> write32 buffer pos (castFloatToWord32 v)) + ) + | hasElemType @Double col = + pure $ + scalarEncoder + (DOUBLE enum) + Nothing + Nothing + ( unboxedColumnWriter @Double + col + (\buffer pos v -> write64 buffer pos (castDoubleToWord64 v)) + ) + | 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 :: + ThriftType -> + Maybe ConvertedType -> + Maybe LogicalType -> + 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 -> + (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 -> + (MemoryBuffer s -> Int -> a -> ST s Int) -> + ColumnWriter s +unboxedColumnWriter = columnWriterWith (Just VU.unsafeIndex) +{-# INLINE unboxedColumnWriter #-} + +columnWriterWith :: + forall a s. + (Columnable a) => + Maybe (VU.Vector a -> Int -> a) -> + Column -> + (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 + (maybe (VU.unsafeIndex values) ($ values) unboxedIndex) + Nothing -> mismatch + _ -> mismatch + where + writeFrom bitmap at buffer pos row + | isPresent bitmap row = do + pos' <- writePrim buffer pos (at row) + pure (pos', True) + | otherwise = pure (pos, False) + mismatch = + error + ("writeParquet: incompatible column representation for " <> columnTypeString col) +{-# INLINE columnWriterWith #-} + +isPresent :: Maybe Bitmap -> Int -> Bool +isPresent Nothing _ = True +isPresent (Just bitmap) row = bitmapTestBit bitmap row +{-# INLINE isPresent #-} + +boolEncoder :: Column -> ST s (Encoder s) +boolEncoder col = do + bitsRef <- newSTRef (0 :: Word8) + countRef <- newSTRef (0 :: Int) + let addBit buffer pos value = do + 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 <- readSTRef buffer.arrayRef + writeByteArray arr pos bits' + writeSTRef bitsRef 0 + writeSTRef countRef 0 + pure (pos + 1) + else do + writeSTRef bitsRef bits' + writeSTRef countRef count' + pure pos + finish buffer pos = do + count <- readSTRef countRef + pos' <- + if count > 0 + then do + bits <- readSTRef bitsRef + arr <- readSTRef buffer.arrayRef + writeByteArray arr pos bits + pure (pos + 1) + else pure pos + writeSTRef bitsRef 0 + writeSTRef countRef 0 + pure pos' + pure + ( Encoder + (BOOLEAN enum) + Nothing + Nothing + (unboxedColumnWriter @Bool col addBit) + finish + ) + +textEncoder :: Column -> Encoder s +textEncoder col = + Encoder + (BYTE_ARRAY enum) + (Just (UTF8 enum)) + (Just (LT_STRING (putField StringType))) + writePresent + (\_ pos -> pure pos) + where + writePresent = case col of + BoxedColumn bitmap (values :: VB.Vector a) -> + case testEquality (typeRep @T.Text) (typeRep @a) of + Just Refl -> writeBoxed bitmap values + Nothing -> mismatch + PackedText bitmap packed -> writePacked bitmap packed + _ -> mismatch + writeBoxed bitmap values buffer pos row + | isPresent bitmap row = do + let Text bytes offset count = VB.unsafeIndex values row + pos' <- writeTextSlice buffer pos bytes offset count + pure (pos', True) + | otherwise = pure (pos, False) + writePacked bitmap packed buffer pos row + | isPresent bitmap row = do + let baseRow = maybe row (`selAt` row) packed.ptSel + start = offAt packed.ptOffsets baseRow + end = offAt packed.ptOffsets (baseRow + 1) + pos' <- writeTextSlice buffer pos packed.ptBytes start (end - start) + pure (pos', True) + | otherwise = pure (pos, False) + writeTextSlice buffer pos bytes offset count = do + writeSTRef buffer.positionRef pos + _ <- ensureCapacity buffer (pos + 4 + count) + writeWord32At buffer pos (fromIntegral count) + arr <- readSTRef buffer.arrayRef + withMutableByteArrayContents arr $ \ptr -> + 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 s +timestampEncoder col = + Encoder + (INT64 enum) + (Just (TIMESTAMP_MICROS enum)) + (Just timestampLogical) + (columnWriter @UTCTime col writeMicros) + (\_ pos -> pure pos) + where + writeMicros buffer pos t = do + writeWord64At buffer pos (fromIntegral (utcToMicros t)) + pure (pos + 8) + +timestampLogical :: LogicalType +timestampLogical = + LT_TIMESTAMP + ( putField + TimestampType + { timestamp_isAdjustedToUTC = putField True + , timestamp_unit = putField (MICROS (putField MicroSeconds)) + } + ) + +utcToMicros :: UTCTime -> Int64 +utcToMicros (UTCTime day dt) = + fromIntegral + ( (toModifiedJulianDay day - 40587) * 86400 * 1000000 + + diffTimeToPicoseconds dt `div` 1000000 + ) +{-# INLINE utcToMicros #-} diff --git a/src/DataFrame/IO/Parquet/Writer/Metadata.hs b/src/DataFrame/IO/Parquet/Writer/Metadata.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Writer/Metadata.hs @@ -0,0 +1,176 @@ +{-# LANGUAGE OverloadedStrings #-} + +module DataFrame.IO.Parquet.Writer.Metadata ( + mkSchemaElem, + rootSchemaElement, + mkDataPageHeader, + mkColumnChunk, + mkRowGroup, + writeFooter, + magic, +) where + +import Control.Monad.ST (stToIO) +import qualified Data.ByteString as BS +import Data.Int (Int64) +import qualified Data.Text as T +import DataFrame.IO.Parquet.Thrift +import DataFrame.IO.Utils.RandomAccess ( + WritableBinaryHandle, + flushBufferToFile, + mallocBuffer, + writeByteString, + writeWord32LE, + ) +import Pinch (enum, putField) +import qualified Pinch + +mkDataPageHeader :: Int -> Int -> Int -> PageHeader +mkDataPageHeader rows uncompressedSize compressedSize = + PageHeader + { ph_type = putField (DATA_PAGE enum) + , ph_uncompressed_page_size = putField (fromIntegral uncompressedSize) + , ph_compressed_page_size = putField (fromIntegral compressedSize) + , ph_crc = putField Nothing + , ph_data_page_header = putField (Just dph) + , ph_index_page_header = putField Nothing + , ph_dictionary_page_header = putField Nothing + , ph_data_page_header_v2 = putField Nothing + } + where + dph = + DataPageHeader + { dph_num_values = putField (fromIntegral rows) + , dph_encoding = putField (PLAIN enum) + , dph_definition_level_encoding = putField (RLE enum) + , dph_repetition_level_encoding = putField (RLE enum) + , dph_statistics = putField Nothing + } + +mkSchemaElem :: + T.Text -> + ThriftType -> + Bool -> + Maybe ConvertedType -> + Maybe LogicalType -> + SchemaElement +mkSchemaElem elementName elementType nullable converted logical = + SchemaElement + { schematype = putField (Just elementType) + , type_length = putField Nothing + , repetition_type = + putField (Just (if nullable then OPTIONAL enum else REQUIRED enum)) + , name = putField elementName + , num_children = putField Nothing + , converted_type = putField converted + , scale = putField Nothing + , precision = putField Nothing + , field_id = putField Nothing + , logicalType = putField logical + } + +rootSchemaElement :: Int -> SchemaElement +rootSchemaElement count = + SchemaElement + { schematype = putField Nothing + , type_length = putField Nothing + , repetition_type = putField Nothing + , name = putField "schema" + , num_children = putField (Just (fromIntegral count)) + , converted_type = putField Nothing + , scale = putField Nothing + , precision = putField Nothing + , field_id = putField Nothing + , logicalType = putField Nothing + } + +mkColumnChunk :: + CompressionCodec -> + ThriftType -> + T.Text -> + Int64 -> + Int -> + Int64 -> + Int -> + ColumnChunk +mkColumnChunk codec columnType columnName offset compressedSize uncompressedSize rgRows = + ColumnChunk + { cc_file_path = putField Nothing + , cc_file_offset = putField offset + , cc_meta_data = putField (Just metadata) + , cc_offset_index_offset = putField Nothing + , cc_offset_index_length = putField Nothing + , cc_column_index_offset = putField Nothing + , cc_column_index_length = putField Nothing + , cc_crypto_metadata = putField Nothing + , cc_encrypted_column_metadata = putField Nothing + } + where + metadata = + ColumnMetaData + { cmd_type = putField columnType + , cmd_encodings = putField [PLAIN enum, RLE enum] + , cmd_path_in_schema = putField [columnName] + , cmd_codec = putField codec + , cmd_num_values = putField (fromIntegral rgRows) + , cmd_total_uncompressed_size = putField uncompressedSize + , cmd_total_compressed_size = putField (fromIntegral compressedSize) + , cmd_key_value_metadata = putField Nothing + , cmd_data_page_offset = putField offset + , cmd_index_page_offset = putField Nothing + , cmd_dictionary_page_offset = putField Nothing + , cmd_statistics = putField Nothing + , cmd_encoding_stats = putField Nothing + , cmd_bloom_filter_offset = putField Nothing + , cmd_bloom_filter_length = putField Nothing + } + +mkRowGroup :: [ColumnChunk] -> Int64 -> Int64 -> Int -> RowGroup +mkRowGroup chunks totalCompressed totalUncompressed rgRows = + RowGroup + { rg_columns = putField chunks + , rg_total_byte_size = putField totalUncompressed + , rg_num_rows = putField (fromIntegral rgRows) + , rg_sorting_columns = putField Nothing + , rg_file_offset = putField Nothing + , rg_total_compressed_size = putField (Just totalCompressed) + , rg_ordinal = putField Nothing + } + +writeFooter :: + WritableBinaryHandle -> + [SchemaElement] -> + Int -> + [RowGroup] -> + [(T.Text, T.Text)] -> + IO () +writeFooter output schemaElements numRows rowGroupMetadata keyValues = do + let metadata = + FileMetadata + { version = putField 1 + , schema = putField schemaElements + , num_rows = putField (fromIntegral numRows) + , row_groups = putField rowGroupMetadata + , key_value_metadata = + putField $ + if null keyValues + then Nothing + else + Just + [ KeyValue (putField k) (putField (Just v)) + | (k, v) <- keyValues + ] + , created_by = putField (Just "dataframe-parquet") + , column_orders = putField Nothing + , encryption_algorithm = putField Nothing + , footer_signing_key_metadata = putField Nothing + } + footer = Pinch.encode Pinch.compactProtocol metadata + buffer <- stToIO (mallocBuffer (BS.length footer + 8)) + writeByteString buffer footer + stToIO (writeWord32LE buffer (fromIntegral (BS.length footer))) + writeByteString buffer magic + flushBufferToFile output buffer + +magic :: BS.ByteString +magic = "PAR1" diff --git a/src/DataFrame/IO/Parquet/Writer/Options.hs b/src/DataFrame/IO/Parquet/Writer/Options.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Writer/Options.hs @@ -0,0 +1,34 @@ +module DataFrame.IO.Parquet.Writer.Options ( + ParquetWriteOptions (..), + WriterStrategy (..), + defaultParquetWriteOptions, +) where + +import DataFrame.IO.Parquet.Thrift +import Pinch (enum) + +data WriterStrategy = InMemory | TwoPass + deriving (Eq, Show) + +data ParquetWriteOptions = ParquetWriteOptions + { pageSize :: !Int + , rowGroupSize :: !Int + , batchRows :: !Int + , subBatchRows :: !Int + , compressionCodec :: !CompressionCodec + , strategy :: !WriterStrategy + , maxRowsPerFile :: !(Maybe Int) + } + deriving (Eq, Show) + +defaultParquetWriteOptions :: ParquetWriteOptions +defaultParquetWriteOptions = + ParquetWriteOptions + { pageSize = 1048576 + , rowGroupSize = 134217728 + , batchRows = 8192 + , subBatchRows = 2048 + , compressionCodec = SNAPPY enum + , strategy = InMemory + , maxRowsPerFile = Nothing + } diff --git a/src/DataFrame/IO/Utils/RandomAccess.hs b/src/DataFrame/IO/Utils/RandomAccess.hs --- a/src/DataFrame/IO/Utils/RandomAccess.hs +++ b/src/DataFrame/IO/Utils/RandomAccess.hs @@ -1,21 +1,87 @@ {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE OverloadedRecordDot #-} -module DataFrame.IO.Utils.RandomAccess where +module DataFrame.IO.Utils.RandomAccess ( + uncurry3, + Range (..), + RandomAccess (..), + ReaderIO (runReaderIO), + LocalFile, + MMappedFile, + ChunkBuffer (..), + unsafeToByteString, + WritableBinaryHandle, + openWritableBinaryFile, + withWritableBinaryFile, + atomicallyWriteFile, + MemoryBuffer (..), + ensureCapacity, + mallocBuffer, + writeByteString, + appendTextArraySlice, + writeWord8, + writeWord32LE, + writeWord64LE, + writeInteger64, + writeWord32At, + writeWord64At, + writeInteger64At, + writeFloatLE, + writeDoubleLE, + bufferResidency, + bufferToByteString, + unsafeBufferView, + flushBufferToBuffer, + resetPosition, + flushBufferToFile, + writeByteStringToFile, +) where +import Control.Exception (bracket, bracketOnError, finally, throw) +import Control.Monad (when) import Control.Monad.IO.Class (MonadIO (..)) -import Data.ByteString (ByteString) -import Data.ByteString.Internal (ByteString (PS)) +import Control.Monad.ST (RealWorld, ST, stToIO) +import Data.Bits (shiftR) +import Data.ByteString.Internal (ByteString (BS, PS), create) +import qualified Data.ByteString.Unsafe as BU +import Data.Int (Int64) +import Data.Primitive.ByteArray ( + 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 (Word8) +import Data.Word (Word32, Word64, Word8) import DataFrame.IO.Parquet.Seeking ( FileBufferedOrSeekable, fGet, fSeek, readLastBytes, ) -import Foreign (castForeignPtr) +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 ( + BufferMode (NoBuffering), + Handle, + IOMode (WriteMode), SeekMode (AbsoluteSeek), + hClose, + hPutBuf, + hSetBinaryMode, + hSetBuffering, + openBinaryFile, + openBinaryTempFileWithDefaultPermissions, ) uncurry3 :: (a -> b -> c -> d) -> (a, b, c) -> d @@ -58,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 @@ -76,3 +150,242 @@ unsafeToByteString v = PS (castForeignPtr ptr) offset' len where (ptr, offset', len) = VS.unsafeToForeignPtr v + +-- Writer buffers are pinned MutableByteArrays so hPutBuf can take a Ptr. + +newtype WritableBinaryHandle = WritableBinaryHandle {unHandle :: Handle} + +openWritableBinaryFile :: FilePath -> IO WritableBinaryHandle +openWritableBinaryFile filepath = do + h <- openBinaryFile filepath WriteMode + hSetBinaryMode h True + hSetBuffering h NoBuffering + pure . WritableBinaryHandle $ h + +atomicallyWriteFile :: + FilePath -> + (FilePath -> IO a) -> + IO a +atomicallyWriteFile path action = + bracketOnError + openAction + removeFile + ( \tmpFile -> do + result <- action tmpFile + renameFile tmpFile path + pure result + ) + where + openAction = + bracketOnError + ( openBinaryTempFileWithDefaultPermissions + (takeDirectory path) + "dataframe-parquet.incomplete" + ) + (\(tmpFile, h) -> hClose h `finally` removeFile tmpFile) + ( \(tmpFile, h) -> do + hClose h + destinationExists <- doesFileExist path + when destinationExists (copyPermissions path tmpFile) + pure tmpFile + ) + +withWritableBinaryFile :: + FilePath -> + (WritableBinaryHandle -> IO a) -> + IO a +withWritableBinaryFile filepath = + bracket + (openWritableBinaryFile filepath) + (hClose . unHandle) + +data MemoryBuffer s = MemoryBuffer + { arrayRef :: !(STRef s (MutableByteArray s)) + , positionRef :: !(STRef s Int) + } + +mallocBuffer :: Int -> ST s (MemoryBuffer s) +mallocBuffer capacity + | capacity < 0 = throw (userError "mallocBuffer: negative capacity") + | otherwise = do + array <- newPinnedByteArray capacity + MemoryBuffer <$> newSTRef array <*> newSTRef 0 + +-- 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 <- readSTRef buffer.arrayRef + maxSize <- getSizeofMutableByteArray array + if needed <= maxSize + then pure array + else do + position <- readSTRef buffer.positionRef + grown <- newPinnedByteArray (needed + (needed `div` 2)) + copyMutableByteArray grown 0 array 0 position + writeSTRef buffer.arrayRef grown + pure grown +{-# INLINE ensureCapacity #-} + +writeWord8 :: MemoryBuffer s -> Word8 -> ST s () +writeWord8 buffer b = do + position <- readSTRef buffer.positionRef + array <- ensureCapacity buffer (position + 1) + writeByteArray array position b + writeSTRef buffer.positionRef (position + 1) +{-# INLINE writeWord8 #-} + +-- 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 <- stToIO (readSTRef buffer.positionRef) + array <- stToIO (ensureCapacity buffer (position + len)) + withMutableByteArrayContents array $ \dst -> + copyBytes (dst `plusPtr` position) (castPtr source) len + stToIO (writeSTRef buffer.positionRef (position + len)) +{-# INLINE writeByteString #-} + +writeWord32LE :: MemoryBuffer s -> Word32 -> ST s () +writeWord32LE buffer w = do + position <- readSTRef buffer.positionRef + writeWord32At buffer position w + writeSTRef buffer.positionRef (position + 4) +{-# INLINE writeWord32LE #-} + +writeWord64LE :: MemoryBuffer s -> Word64 -> ST s () +writeWord64LE buffer w = do + position <- readSTRef buffer.positionRef + writeWord64At buffer position w + writeSTRef buffer.positionRef (position + 8) +{-# INLINE writeWord64LE #-} + +writeWord32At :: MemoryBuffer s -> Int -> Word32 -> ST s () +writeWord32At buffer position w = do + array <- ensureCapacity buffer (position + 4) + writeByteArray array position (fromIntegral w :: Word8) + writeByteArray array (position + 1) (fromIntegral (w `shiftR` 8) :: Word8) + writeByteArray array (position + 2) (fromIntegral (w `shiftR` 16) :: Word8) + writeByteArray array (position + 3) (fromIntegral (w `shiftR` 24) :: Word8) +{-# INLINE writeWord32At #-} + +writeWord64At :: MemoryBuffer s -> Int -> Word64 -> ST s () +writeWord64At buffer position w = do + array <- ensureCapacity buffer (position + 8) + writeByteArray array position (fromIntegral w :: Word8) + writeByteArray array (position + 1) (fromIntegral (w `shiftR` 8) :: Word8) + writeByteArray array (position + 2) (fromIntegral (w `shiftR` 16) :: Word8) + writeByteArray array (position + 3) (fromIntegral (w `shiftR` 24) :: Word8) + writeByteArray array (position + 4) (fromIntegral (w `shiftR` 32) :: Word8) + writeByteArray array (position + 5) (fromIntegral (w `shiftR` 40) :: Word8) + writeByteArray array (position + 6) (fromIntegral (w `shiftR` 48) :: Word8) + writeByteArray array (position + 7) (fromIntegral (w `shiftR` 56) :: Word8) +{-# INLINE writeWord64At #-} + +writeInteger64 :: MemoryBuffer s -> Integer -> ST s () +writeInteger64 buffer value = do + position <- readSTRef buffer.positionRef + newPosition <- writeInteger64At buffer position value + writeSTRef buffer.positionRef newPosition +{-# INLINE writeInteger64 #-} + +writeInteger64At :: MemoryBuffer s -> Int -> Integer -> ST s Int +writeInteger64At buffer position value + | value < toInteger (minBound :: Int64) = outOfRange + | value > toInteger (maxBound :: Int64) = outOfRange + | otherwise = do + writeWord64At buffer position (fromIntegral value) + pure (position + 8) + where + outOfRange = + throw (userError "writeParquet: Integer value is outside the INT64 range") +{-# INLINE writeInteger64At #-} + +writeFloatLE :: MemoryBuffer s -> Float -> ST s () +writeFloatLE buffer = writeWord32LE buffer . castFloatToWord32 +{-# INLINE writeFloatLE #-} + +writeDoubleLE :: MemoryBuffer s -> Double -> ST s () +writeDoubleLE buffer = writeWord64LE buffer . castDoubleToWord64 +{-# INLINE writeDoubleLE #-} + +flushBufferToBuffer :: MemoryBuffer s -> MemoryBuffer s -> ST s () +flushBufferToBuffer source destination + | source.arrayRef == destination.arrayRef = pure () + | otherwise = do + sourceArray <- readSTRef source.arrayRef + sourcePosition <- readSTRef source.positionRef + destinationPosition <- readSTRef destination.positionRef + destinationArray <- + ensureCapacity destination (destinationPosition + sourcePosition) + copyMutableByteArray + destinationArray + destinationPosition + sourceArray + 0 + sourcePosition + writeSTRef destination.positionRef (destinationPosition + sourcePosition) + writeSTRef source.positionRef 0 +{-# INLINE flushBufferToBuffer #-} + +bufferToByteString :: MemoryBuffer RealWorld -> IO ByteString +bufferToByteString buffer = do + array <- stToIO (readSTRef buffer.arrayRef) + position <- stToIO (readSTRef buffer.positionRef) + create position $ \dst -> + withMutableByteArrayContents array $ \src -> + copyBytes dst (castPtr src) position + +-- | 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 s -> ST s () +resetPosition buffer = writeSTRef buffer.positionRef 0 +{-# INLINE resetPosition #-} + +-- 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 <- stToIO (readSTRef buffer.arrayRef) + position <- stToIO (readSTRef buffer.positionRef) + withMutableByteArrayContents array $ \ptr -> do + let chunkSize = 262144 + go offset + | offset >= position = pure () + | otherwise = do + let n = min chunkSize (position - offset) + hPutBuf h (ptr `plusPtr` offset) n + go (offset + n) + go 0 + stToIO (resetPosition buffer) + +writeByteStringToFile :: WritableBinaryHandle -> ByteString -> IO () +writeByteStringToFile (WritableBinaryHandle h) bs = + BU.unsafeUseAsCStringLen bs $ \(source, len) -> do + let chunkSize = 262144 + go offset + | offset >= len = pure () + | otherwise = do + let n = min chunkSize (len - offset) + hPutBuf h (source `plusPtr` offset) n + go (offset + n) + go 0 + +appendTextArraySlice :: MemoryBuffer s -> TA.Array -> Int -> Int -> ST s () +appendTextArraySlice buffer source offset count + | count < 0 = throw (userError "appendTextArraySlice: negative length") + | otherwise = do + position <- readSTRef buffer.positionRef + array <- ensureCapacity buffer (position + count) + withMutableByteArrayContents array $ \destination -> + TA.copyToPointer source offset (destination `plusPtr` position) count + writeSTRef buffer.positionRef (position + count) +{-# INLINE appendTextArraySlice #-} diff --git a/src/DataFrame/Typed/IO/Parquet.hs b/src/DataFrame/Typed/IO/Parquet.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/IO/Parquet.hs @@ -0,0 +1,65 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Typed Parquet reading, validated against a type-level schema. +module DataFrame.Typed.IO.Parquet ( + readParquet, + readParquetWithError, + readParquetWithOpts, + readParquetFiles, +) where + +import Control.Applicative ((<|>)) +import Control.Exception (SomeException, try) +import qualified Data.Text as T + +import DataFrame.IO.Parquet (ParquetReadOptions (..), defaultParquetReadOptions) +import qualified DataFrame.IO.Parquet as Parquet +import DataFrame.Typed.Freeze (freezeOrThrow, freezeWithError) +import DataFrame.Typed.Schema (KnownSchema, schemaColumnNames) +import DataFrame.Typed.Types (TypedDataFrame) + +-- | Read only the columns the schema names, unless the caller chose. +schemaOptions :: + forall cols. (KnownSchema cols) => ParquetReadOptions -> ParquetReadOptions +schemaOptions opts = + opts + { selectedColumns = selectedColumns opts <|> Just (schemaColumnNames @cols) + } + +-- | 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 an error on mismatch. +readParquetWithError :: + forall cols. + (KnownSchema cols) => + FilePath -> IO (Either T.Text (TypedDataFrame cols)) +readParquetWithError path = do + r <- + try + (Parquet.readParquetWithOpts (schemaOptions @cols defaultParquetReadOptions) path) + pure $ case r of + Left (e :: SomeException) -> Left (T.pack (show e)) + Right df -> freezeWithError df + +-- | Read a Parquet file with custom options, throwing on mismatch. +readParquetWithOpts :: + forall cols. + (KnownSchema cols) => + ParquetReadOptions -> FilePath -> IO (TypedDataFrame cols) +readParquetWithOpts opts path = + Parquet.readParquetWithOpts (schemaOptions @cols opts) path + >>= freezeOrThrow @cols + +-- | Read a directory\/glob of Parquet files into a typed DataFrame. +readParquetFiles :: + forall cols. (KnownSchema cols) => FilePath -> IO (TypedDataFrame cols) +readParquetFiles path = + Parquet.readParquetFilesWithOpts + (schemaOptions @cols defaultParquetReadOptions) + path + >>= freezeOrThrow @cols diff --git a/stress/DataFrame10GB.hs b/stress/DataFrame10GB.hs new file mode 100644 --- /dev/null +++ b/stress/DataFrame10GB.hs @@ -0,0 +1,137 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE NumericUnderscores #-} + +module DataFrame10GB ( + stressDataFrame, + stressRows, + stressColumns, + stressResidentBytesLowerBound, +) where + +import Control.Monad.ST (runST) +import Data.Int (Int32, Int64) +import qualified Data.Text as T +import qualified Data.Text.Array as A +import Data.Time (UTCTime (UTCTime), addDays, fromGregorian, secondsToDiffTime) +import qualified Data.Vector as VB +import qualified Data.Vector.Unboxed as VU +import Data.Word (Word8) +import DataFrame.Internal.Column (Column (..)) +import DataFrame.Internal.Column.Bitmap (Bitmap) +import DataFrame.Internal.Data.PackedText (mkPackedContiguous32) +import DataFrame.Internal.DataFrame (DataFrame, fromColumns) + +stressRows :: Int +stressRows = 1_000_000 + +stressGroups :: Int +stressGroups = 16 + +stressColumns :: Int +stressColumns = stressGroups * 14 + +textBytesPerRow :: Int +textBytesPerRow = 320 + +stressResidentBytesLowerBound :: Integer +stressResidentBytesLowerBound = + fromIntegral stressRows + * fromIntegral stressGroups + * fromIntegral (2 * textBytesPerRow + 2 * (4 + 8 + 4 + 8)) + +stressDataFrame :: DataFrame +stressDataFrame = fromColumns (concatMap columnGroup [0 .. stressGroups - 1]) + +columnGroup :: Int -> [(T.Text, Column)] +columnGroup group = + [ named "int32" (UnboxedColumn Nothing (int32Values group)) + , named "int64" (UnboxedColumn Nothing (int64Values group)) + , named "float" (UnboxedColumn Nothing (floatValues group)) + , named "double" (UnboxedColumn Nothing (doubleValues group)) + , named "bool" (UnboxedColumn Nothing (boolValues group)) + , named "timestamp" (BoxedColumn Nothing (timestampValues group)) + , named "text" (textColumn Nothing group) + , named + "nullable_int32" + (UnboxedColumn (Just nullableBitmap) (int32Values (group + stressGroups))) + , named + "nullable_int64" + (UnboxedColumn (Just nullableBitmap) (int64Values (group + stressGroups))) + , named + "nullable_float" + (UnboxedColumn (Just nullableBitmap) (floatValues (group + stressGroups))) + , named + "nullable_double" + (UnboxedColumn (Just nullableBitmap) (doubleValues (group + stressGroups))) + , named + "nullable_bool" + (UnboxedColumn (Just nullableBitmap) (boolValues (group + stressGroups))) + , named + "nullable_timestamp" + (BoxedColumn (Just nullableBitmap) (timestampValues (group + stressGroups))) + , named "nullable_text" (textColumn (Just nullableBitmap) (group + stressGroups)) + ] + where + named suffix column = (T.pack ("group_" <> show group <> "_" <> suffix), column) + +nullableBitmap :: Bitmap +nullableBitmap = VU.replicate (stressRows `div` 8) (0xFE :: Word8) + +int32Values :: Int -> VU.Vector Int32 +int32Values salt = + VU.generate stressRows $ \row -> + fromIntegral ((row + salt * 10_007) `mod` 2_000_001 - 1_000_000) + +int64Values :: Int -> VU.Vector Int64 +int64Values salt = + VU.generate stressRows $ \row -> + fromIntegral row * 1_000_003 - fromIntegral salt * 10_000_019 + +floatValues :: Int -> VU.Vector Float +floatValues salt = + VU.generate stressRows $ \row -> + fromIntegral ((row + salt * 101) `mod` 100_003) / 17 + +doubleValues :: Int -> VU.Vector Double +doubleValues salt = + VU.generate stressRows $ \row -> + fromIntegral row / 31.0 - fromIntegral salt * 1_000.25 + +boolValues :: Int -> VU.Vector Bool +boolValues salt = VU.generate stressRows (\row -> (row + salt) `mod` 3 == 0) + +timestampValues :: Int -> VB.Vector UTCTime +timestampValues salt = + VB.replicate + stressRows + ( UTCTime + (addDays (fromIntegral salt) (fromGregorian 2020 1 1)) + (secondsToDiffTime (fromIntegral (salt * 1_337 `mod` 86_400))) + ) + +textColumn :: Maybe Bitmap -> Int -> Column +textColumn bitmap salt = PackedText bitmap $ runST $ do + target <- A.new (stressRows * textBytesPerRow) + let template = textTemplate salt + fill !row + | row >= stressRows = pure () + | otherwise = do + A.copyI textBytesPerRow target (row * textBytesPerRow) template 0 + fill (row + 1) + fill 0 + bytes <- A.unsafeFreeze target + let offsets = + VU.generate + (stressRows + 1) + (\row -> fromIntegral (row * textBytesPerRow) :: Int32) + pure (mkPackedContiguous32 bytes offsets) + +textTemplate :: Int -> A.Array +textTemplate salt = A.run $ do + bytes <- A.new textBytesPerRow + let byte = fromIntegral (97 + salt `mod` 26) + fill !index + | index >= textBytesPerRow = pure () + | otherwise = A.unsafeWrite bytes index byte >> fill (index + 1) + fill 0 + pure bytes diff --git a/stress/StressMain.hs b/stress/StressMain.hs new file mode 100644 --- /dev/null +++ b/stress/StressMain.hs @@ -0,0 +1,38 @@ +module Main (main) where + +import Control.Exception (evaluate) +import Control.Monad (unless) +import DataFrame.IO.Parquet (readParquet) +import DataFrame.IO.Parquet.Writer (writeParquet) +import DataFrame.Internal.DataFrame (forceDataFrame) +import DataFrame10GB ( + stressColumns, + stressDataFrame, + stressResidentBytesLowerBound, + stressRows, + ) +import System.Exit (exitFailure) +import System.FilePath (()) +import System.IO (hPutStrLn, stderr) +import System.IO.Temp (withSystemTempDirectory) + +main :: IO () +main = withSystemTempDirectory "dataframe-parquet-10gb-stress" $ \directory -> do + expected <- evaluate (forceDataFrame stressDataFrame) + let output = directory "roundtrip.parquet" + putStrLn + ( "writing " + <> show stressRows + <> " rows x " + <> show stressColumns + <> " columns (at least " + <> show stressResidentBytesLowerBound + <> " resident payload bytes)" + ) + writeParquet output expected + putStrLn "reading the stress dataframe" + actual <- readParquet output + putStrLn "checking dataframe equivalence" + unless (expected == actual) $ do + hPutStrLn stderr "10 GiB Parquet roundtrip mismatch" + exitFailure diff --git a/tests/Main.hs b/tests/Main.hs new file mode 100644 --- /dev/null +++ b/tests/Main.hs @@ -0,0 +1,413 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Tests for the writer-buffer logic in "DataFrame.IO.Utils.RandomAccess". +module Main where + +import Control.Exception (SomeException, catch, evaluate) +import qualified Data.ByteString as BS +import Data.List (sortOn) +import qualified System.Exit as Exit +import System.FilePath (()) +import System.IO.Temp (withSystemTempDirectory) +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 ( + ParquetReadOptions (..), + defaultParquetReadOptions, + readParquet, + readParquetFiles, + readParquetWithOpts, + _readParquetWithOpts, + ) +import DataFrame.IO.Parquet.Writer ( + ParquetWriteOptions (..), + defaultParquetWriteOptions, + writeParquet, + writeParquetWithOptions, + ) +import DataFrame.IO.Utils.RandomAccess +import DataFrame.Internal.Column (columnTypeString, fromList) +import DataFrame.Internal.DataFrame ( + DataFrame, + columnNames, + 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 <- 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 <- stToIO (bufferResidency buffer) + bytes <- bufferToByteString buffer + assertEqual "direct write residency" 27 residency + assertEqual + "direct write bytes" + ( BS.pack + [ 0xaa + , 0x12 + , 0x34 + , 0x56 + , 0x78 + , 0x01 + , 0x02 + , 0x03 + , 0x04 + , 0x05 + , 0x06 + , 0x07 + , 0x08 + , 0x00 + , 0x00 + , 0x80 + , 0x3f + , 0x00 + , 0x00 + , 0x00 + , 0x00 + , 0x00 + , 0x00 + , 0xf0 + , 0x3f + , 0xfe + , 0xff + ] + ) + bytes + +directBufferFlush :: Test +directBufferFlush = TestCase $ do + source <- stToIO (mallocBuffer 0) + destination <- stToIO (mallocBuffer 0) + writeByteString destination (BS.pack [1, 2]) + writeByteString source (BS.pack [3, 4, 5]) + 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 + stToIO (flushBufferToBuffer destination destination) + selfFlushedBytes <- bufferToByteString destination + assertEqual "self flush is a no-op" destinationBytes selfFlushedBytes + stToIO (resetPosition destination) + destinationResidency <- stToIO (bufferResidency destination) + assertEqual "reset position" 0 destinationResidency + +directFileFlush :: Test +directFileFlush = TestCase $ + withSystemTempDirectory "dfpq-buffer" $ \dir -> do + let outPath = dir "out.bin" + payload = BS.pack (take 300000 (cycle [0 .. 255])) + buffer <- stToIO (mallocBuffer 1) + writeByteString buffer payload + withWritableBinaryFile outPath $ \output -> + flushBufferToFile output buffer + residency <- stToIO (bufferResidency buffer) + contents <- BS.readFile outPath + assertEqual "source cleared after file flush" 0 residency + assertEqual "large payload round-trips" payload contents + +writerRoundTrip :: String -> FilePath -> Test +writerRoundTrip label path = TestCase $ + withSystemTempDirectory "dfpq-writer" $ \dir -> do + df <- readParquet path + let out = dir "out.parquet" + writeParquet out df + df' <- readParquet out + assertEqual label df df' + +writerRoundTripTiny :: String -> FilePath -> Test +writerRoundTripTiny label path = TestCase $ + withSystemTempDirectory "dfpq-writer" $ \dir -> do + df <- readParquet path + let out = dir "out.parquet" + writeParquetWithOptions tinyWriteOpts out df + 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 = fromColumns [("text", fromList [payload, "short"])] + firstOut = dir "large-text-1.parquet" + secondOut = dir "large-text-2.parquet" + writeParquetWithOptions tinyWriteOpts firstOut df + firstRoundTrip <- readParquet firstOut + writeParquetWithOptions tinyWriteOpts secondOut firstRoundTrip + secondRoundTrip <- readParquet secondOut + assertEqual "large text first round-trip" df firstRoundTrip + assertEqual "large text second round-trip" df secondRoundTrip + +-- | 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 + df <- readParquet path + let pattern_ = dir "shards" "part-*.parquet" + writeParquetWithOptions + defaultParquetWriteOptions{maxRowsPerFile = Just rowsPerFile} + pattern_ + df + shards <- listDirectory (dir "shards") + assertEqual (label <> ": shard count") expectedShards (Prelude.length shards) + assertEqual + (label <> ": shard names") + ["part-" <> pad i <> ".parquet" | i <- [0 .. expectedShards - 1]] + (sortOn id shards) + df' <- readParquetFiles pattern_ + assertEqual (label <> ": shards round-trip") df df' + where + pad i = let s = show i in replicate (5 - Prelude.length s) '0' <> s + +-- | A path without a @*@ placeholder is rejected when sharding is requested. +shardedWriteRequiresPattern :: Test +shardedWriteRequiresPattern = TestCase $ + withSystemTempDirectory "dfpq-writer" $ \dir -> do + df <- readParquet "tests/data/mtcars.parquet" + threw <- + ( False + <$ writeParquetWithOptions + defaultParquetWriteOptions{maxRowsPerFile = Just 4} + (dir "out.parquet") + df + ) + `catch` (\e -> True <$ evaluate (Prelude.length (show (e :: SomeException)))) + unless threw (assertFailure "expected an error for a path without '*'") + +-- | Int, Int64 and Integer columns round-trip with their Haskell type. +writerRoundTripNativeIntTypes :: Test +writerRoundTripNativeIntTypes = TestCase $ + withSystemTempDirectory "dfpq-writer" $ \dir -> do + let df = + fromColumns + [ ("int", fromList [1 :: Int, 2, 3]) + , ("int64", fromList [1 :: Int64, 2, 3]) + , ("int32", fromList [1 :: Int32, 2, 3]) + , ("integer", fromList [1 :: Integer, 2, 3]) + , ("nullableInt", fromList [Just (1 :: Int), Nothing, Just 3]) + , ("nullableInt64", fromList [Just (1 :: Int64), Nothing, Just 3]) + ] + out = dir "int-types.parquet" + writeParquet out df + df' <- readParquet out + assertEqual + "native int types: column types" + (columnTypes df) + (columnTypes df') + assertEqual "native int types: frame" df df' + +columnTypes :: DataFrame -> [(String, String)] +columnTypes df = + [ (T.unpack name, columnTypeString (fromJust (getColumn name df))) + | name <- columnNames df + ] + +tinyWriteOpts :: ParquetWriteOptions +tinyWriteOpts = + defaultParquetWriteOptions + { pageSize = 64 + , rowGroupSize = 512 + , batchRows = 4 + , 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 "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 + "writer roundtrip: alltypes_plain" + (writerRoundTrip "alltypes_plain" "tests/data/alltypes_plain.parquet") + , TestLabel + "writer roundtrip: alltypes_plain.snappy" + ( writerRoundTrip + "alltypes_plain.snappy" + "tests/data/alltypes_plain.snappy.parquet" + ) + , TestLabel + "writer roundtrip: alltypes_dictionary" + (writerRoundTrip "alltypes_dictionary" "tests/data/alltypes_dictionary.parquet") + , TestLabel + "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 + "writer roundtrip: mtcars" + (writerRoundTrip "mtcars" "tests/data/mtcars.parquet") + , TestLabel + "writer roundtrip: int32_decimal" + (writerRoundTrip "int32_decimal" "tests/data/int32_decimal.parquet") + , TestLabel + "writer roundtrip: int64_decimal" + (writerRoundTrip "int64_decimal" "tests/data/int64_decimal.parquet") + , TestLabel + "writer roundtrip: sharded mtcars" + (writerRoundTripSharded "sharded mtcars" "tests/data/mtcars.parquet" 10 4) + , TestLabel + "writer roundtrip: sharded exact multiple" + (writerRoundTripSharded "sharded exact" "tests/data/mtcars.parquet" 32 1) + , TestLabel + "sharded write requires a '*' pattern" + shardedWriteRequiresPattern + , TestLabel + "writer roundtrip: Int/Integer keep their Haskell type" + writerRoundTripNativeIntTypes + , TestLabel + "writer roundtrip: alltypes_plain multi-page" + ( writerRoundTripTiny + "alltypes_plain multi-page" + "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 + result <- runTestTT tests + if failures result > 0 || errors result > 0 + then Exit.exitFailure + else Exit.exitSuccess diff --git a/tests/data/alltypes_dictionary.parquet b/tests/data/alltypes_dictionary.parquet new file mode 100644 Binary files /dev/null and b/tests/data/alltypes_dictionary.parquet differ diff --git a/tests/data/alltypes_plain.parquet b/tests/data/alltypes_plain.parquet new file mode 100644 Binary files /dev/null and b/tests/data/alltypes_plain.parquet differ diff --git a/tests/data/alltypes_plain.snappy.parquet b/tests/data/alltypes_plain.snappy.parquet new file mode 100644 Binary files /dev/null and b/tests/data/alltypes_plain.snappy.parquet differ diff --git a/tests/data/alltypes_tiny_pages.parquet b/tests/data/alltypes_tiny_pages.parquet new file mode 100644 Binary files /dev/null and b/tests/data/alltypes_tiny_pages.parquet differ diff --git a/tests/data/int32_decimal.parquet b/tests/data/int32_decimal.parquet new file mode 100644 Binary files /dev/null and b/tests/data/int32_decimal.parquet differ diff --git a/tests/data/int64_decimal.parquet b/tests/data/int64_decimal.parquet new file mode 100644 Binary files /dev/null and b/tests/data/int64_decimal.parquet differ diff --git a/tests/data/mtcars.parquet b/tests/data/mtcars.parquet new file mode 100644 Binary files /dev/null and b/tests/data/mtcars.parquet differ diff --git a/tests/data/transactions.parquet b/tests/data/transactions.parquet new file mode 100644 Binary files /dev/null and b/tests/data/transactions.parquet differ