diff --git a/dataframe-parquet.cabal b/dataframe-parquet.cabal --- a/dataframe-parquet.cabal +++ b/dataframe-parquet.cabal @@ -1,6 +1,6 @@ -cabal-version: 2.4 +cabal-version: 3.4 name: dataframe-parquet -version: 1.1.0.1 +version: 1.2.0.0 synopsis: Parquet reader and writer for the dataframe ecosystem. description: @@ -46,21 +46,23 @@ DataFrame.IO.Parquet.Time DataFrame.IO.Parquet.Utils 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.0 && < 2.1, + dataframe-core:internal >= 2.0 && < 2.1, + dataframe-operations >= 2.0 && < 2.1, + dataframe-parsing:internal >= 2.0 && < 2.1, 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 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 @@ -12,6 +12,7 @@ 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) @@ -25,23 +26,26 @@ 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 qualified Data.Vector as Vector import qualified Data.Vector.Unboxed as VU import DataFrame.Errors (DataFrameException (ColumnsNotFoundException)) import DataFrame.IO.Parquet.Page ( PageDecoder, UnboxedPageDecoder, + appendNullableStringPageIO, + appendStringPageIO, boolDecoder, byteArrayDecoder, doubleDecoder, fixedLenByteArrayDecoder, floatDecoder, + foldColumnDataPagesM, + foldColumnPagesM, int32Decoder, int64Decoder, int96Decoder, - foldColumnPagesM, ) import DataFrame.IO.Parquet.Seeking ( FileBufferedOrSeekable, @@ -76,6 +80,11 @@ ) import DataFrame.Internal.Column (Column, Columnable) import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.ColumnBuilder ( + freezeTextChunk, + mergeTextChunks, + newTextBuilder, + ) import DataFrame.Internal.DataFrame (DataFrame (..)) import DataFrame.Internal.Expression (Expr, getColumns) import DataFrame.Operations.Merge () @@ -219,6 +228,10 @@ (RandomAccess m, MonadIO m) => ParquetReadOptions -> m DataFrame +{-# SPECIALIZE parseParquetWithOpts :: + ParquetReadOptions -> ReaderIO FileBufferedOrSeekable DataFrame + #-} +{-# INLINEABLE parseParquetWithOpts #-} parseParquetWithOpts opts = do metadata <- parseFileMetadata @@ -344,6 +357,13 @@ [ColumnChunk] -> ColumnDescription -> m Column +{-# SPECIALIZE parseColumnChunks :: + Int -> + [ColumnChunk] -> + ColumnDescription -> + ReaderIO FileBufferedOrSeekable Column + #-} +{-# INLINEABLE parseColumnChunks #-} parseColumnChunks totalRows chunks description | description.maxRepetitionLevel == 0 && description.maxDefinitionLevel == 0 = getNonNullableColumn totalRows description chunks @@ -353,6 +373,7 @@ getRepeatedColumn description chunks -- | Decode a required (non-nullable, non-repeated) column. +{-# INLINEABLE getNonNullableColumn #-} getNonNullableColumn :: forall m. (RandomAccess m, MonadIO m) => @@ -368,7 +389,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 +403,27 @@ go decoder = foldNonNullable totalRows (foldColumnPagesM description decoder chunks) + -- Decode a non-nullable BYTE_ARRAY (UTF-8) column straight into a single + -- shared byte buffer + offsets ('PackedText'), instead of a boxed vector + -- of per-row 'Text'. Each page's decoded 'Text' values (which share the + -- chunk dictionary for dictionary-encoded pages) are appended by memcpy + -- into one builder across all pages/chunks, then frozen once. This is the + -- same representation the fast CSV reader uses and matches Arrow's string + -- layout: no retained per-row 'Text' headers, no eager UTF-8 validation. + goPackedText :: m Column + goPackedText = do + builder <- liftIO $ stToIO (newTextBuilder totalRows (totalRows * 8)) + _ <- + 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) => @@ -391,6 +433,7 @@ foldNonNullableUnboxed totalRows (foldColumnPagesM description decoder chunks) -- | Decode an optional (nullable) column. +{-# INLINEABLE getNullableColumn #-} getNullableColumn :: forall m. (RandomAccess m, MonadIO m) => @@ -406,7 +449,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 +465,38 @@ m Column go decoder = foldNullable maxDef totalRows (foldColumnPagesM description decoder chunks) + + -- Nullable BYTE_ARRAY (UTF-8): decode straight into a 'PackedText' (shared + -- byte buffer + offsets + validity bitmap) via the text builder, walking + -- def-levels to interleave nulls. Avoids the boxed @Vector Text@ the + -- generic 'foldNullable' path would build. + goPackedTextNullable :: m Column + goPackedTextNullable = do + builder <- liftIO $ stToIO (newTextBuilder totalRows (totalRows * 8)) + _ <- + 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) => @@ -519,14 +585,15 @@ 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 + -- (ticks per second, picoseconds per tick) for each unit. Convert at + -- native precision: a millisecond grid keeps ms, a nanosecond grid + -- keeps ns. (The old code multiplied by @1_000_000 `div` divisor@, + -- which truncated NANOS to 0 and collapsed every value to the epoch.) + 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 @@ -547,6 +614,20 @@ else col applyLogicalType _ col = col -microsecondsToUTCTime :: Int64 -> UTCTime -microsecondsToUTCTime us = - posixSecondsToUTCTime (fromIntegral us / 1_000_000) +{- | Convert an epoch timestamp expressed as @ticksPerSecond@ ticks/second +(each tick = @psPerTick@ picoseconds) to 'UTCTime', at full precision. + +Splits into days + within-day picoseconds with integer 'divMod' (which floors, +so the split is correct for pre-epoch negative values too), and never forms +picoseconds-since-epoch — that would overflow 'Int64' for modern dates — only +the bounded within-day picosecond count. 40587 is the Modified Julian Day of +the Unix epoch (1970-01-01). +-} +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/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,17 @@ fixedLenByteArrayDecoder, -- Page iteration foldColumnPagesM, + foldColumnDataPagesM, + RawPage, + 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,6 +36,7 @@ import qualified Data.Vector as VB import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU +import Data.Word (Word8) import DataFrame.IO.Parquet.Decompress (decompressData) import DataFrame.IO.Parquet.Dictionary ( DictVals (..), @@ -58,6 +65,14 @@ littleEndianWord32, littleEndianWord64, ) +import DataFrame.Internal.ColumnBuilder ( + TextBuilder, + appendNull, + appendText, + appendTextSliceFromPtr, + ) +import Foreign.Ptr (Ptr, castPtr, plusPtr) +import Foreign.Storable (peekByteOff) import GHC.Float (castWord32ToFloat, castWord64ToDouble) import Pinch (decodeWithLeftovers) import qualified Pinch @@ -82,7 +97,12 @@ boolDecoder :: UnboxedPageDecoder Bool boolDecoder mDict enc nPresent bs = case enc of - PLAIN _ -> VU.fromList (readNBool nPresent bs) + -- PLAIN bools are bit-packed (1 bit/value, LSB-first). Generate the + -- unboxed vector directly by indexing the bit for each row, avoiding an + -- intermediate @[Bool]@ list. + PLAIN _ -> + VU.generate nPresent $ \i -> + (BSU.unsafeIndex bs (i `shiftR` 3) `shiftR` (i .&. 7)) .&. 1 == 1 RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getBool PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getBool _ -> error ("boolDecoder: unsupported encoding " ++ show enc) @@ -208,17 +228,20 @@ 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. +{- | A decoded DATA page handed to a fold step: the running dictionary, the +page encoding, the present-value count, the (decompressed) value bytes, and the +definition/repetition level vectors. +-} +type RawPage = + (Maybe DictVals, Encoding, Int, BS.ByteString, VU.Vector Int, VU.Vector Int) -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). +{- | Left-fold a monadic step over every DATA page (V1 or V2) of every column +chunk, in order, threading the running dictionary internally and handing each +page to @step@ as a 'RawPage' (no value decoding — the step decides how to +materialize, e.g. into a typed vector or straight into a text buffer). -Dictionary pages are consumed silently and update the running dictionary that -is threaded through the recursion. @INDEX_PAGE@s are skipped. +Dictionary pages update the running dictionary; @INDEX_PAGE@s are skipped. Only +one page's bytes are live at a time. -- TODO: when a page index is available, use it here to compute which page -- byte ranges to request from the RandomAccess layer instead of reading the @@ -229,16 +252,16 @@ -- overlap the requested range, avoiding decompression of irrelevant pages -- entirely. -} -foldColumnPagesM :: - forall m v a acc. - (RandomAccess m, MonadIO m, VG.Vector v a) => +{-# 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 +275,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 +300,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 = @@ -298,11 +320,127 @@ 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 +{- | Left-fold over per-page value triples, decoding each page with @decoder@. +A thin wrapper over 'foldColumnDataPagesM' for the typed (numeric/boxed) column +paths. +-} +{-# 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) + +{- | Append a non-nullable BYTE_ARRAY page's strings straight into a +'TextBuilder', no intermediate boxed 'Text' vector. + +PLAIN pages copy each value's UTF-8 bytes directly from the page buffer pointer +('appendTextSliceFromPtr'), skipping the per-value 'decodeUtf8Lenient' + +'ByteString' slicing + 'Text' allocation that @readNTexts@ would do. +Dictionary pages append the shared dictionary 'Text's by reference. +-} +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 ds) -> + let (idxs, _) = decodeDictIndices nPresent bs + in stToIO (VU.mapM_ (\i -> appendText builder (ds VB.! i)) idxs) + Just d -> error ("appendStringPageIO: wrong dict type, got " ++ show d) + Nothing -> error "appendStringPageIO: dictionary-encoded page but no dictionary seen" + +{- | Read a little-endian 4-byte length prefix at byte @off@ from a raw page +pointer. +-} +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 #-} + +{- | Append a NULLABLE BYTE_ARRAY page into a 'TextBuilder', interleaving nulls +by walking the definition levels: a present row (@def == maxDef@) consumes the +next stored value (dictionary entry by reference, or the next length-prefixed +PLAIN slice by memcpy); a null row ('appendNull') writes an offset and clears a +validity bit. Builds 'PackedText' + validity bitmap with no per-row 'Text' or +@[Maybe a]@ allocation. Only present values are stored in the page payload. +-} +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 ds) -> do + let (idxs, _) = decodeDictIndices nPresent bs + go !_ !i | i >= nDefs = pure () + go !j !i + | VU.unsafeIndex defs i == maxDef = + stToIO (appendText builder (ds VB.! VU.unsafeIndex idxs j)) + >> go (j + 1) (i + 1) + | 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" + -- --------------------------------------------------------------------------- -- Page header parsing -- --------------------------------------------------------------------------- @@ -313,15 +451,6 @@ -- --------------------------------------------------------------------------- -- Batch value readers -- --------------------------------------------------------------------------- - -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 readNInt32 :: Int -> BS.ByteString -> VU.Vector Int32 readNInt32 n bs = VU.generate n $ \i -> littleEndianInt32 (BS.drop (4 * i) bs) diff --git a/src/DataFrame/IO/Parquet/Utils.hs b/src/DataFrame/IO/Parquet/Utils.hs --- a/src/DataFrame/IO/Parquet/Utils.hs +++ b/src/DataFrame/IO/Parquet/Utils.hs @@ -167,6 +167,7 @@ single streaming left fold ('PageFold'), avoiding any intermediate list or concatenation allocation. Only one page's values are live at a time. -} +{-# INLINEABLE foldNonNullable #-} foldNonNullable :: forall m a. (MonadIO m, Columnable a) => @@ -186,6 +187,7 @@ v <- liftIO $ VB.unsafeFreeze mv return (BoxedColumn Nothing v) +{-# INLINEABLE foldNonNullableUnboxed #-} foldNonNullableUnboxed :: forall m a. (MonadIO m, Columnable a, VU.Unbox a) => @@ -221,6 +223,7 @@ A 'hasNull' flag is accumulated during the scatter so the 'buildBitmapFromValid' call is skipped entirely when all values are present. -} +{-# INLINEABLE foldNullable #-} foldNullable :: forall m a. (MonadIO m, Columnable a) => @@ -264,6 +267,7 @@ else return Nothing return (BoxedColumn maybeBm dat) +{-# INLINEABLE foldNullableUnboxed #-} foldNullableUnboxed :: forall m a. (MonadIO m, Columnable a, VU.Unbox a) => @@ -310,6 +314,7 @@ Threshold formula: @defT_r = maxDef - 2 * (maxRep - r)@. -} +{-# INLINEABLE foldRepeated #-} foldRepeated :: forall m a. ( MonadIO m @@ -333,6 +338,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 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,57 @@ +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Typed Parquet reading. + +The reader validates the file against a type-level schema as it loads, supplied +by type application: + +@ +type Trips = '[Column \"id\" Int, Column \"fare\" Double] + +trips <- readParquet \@Trips \"trips.parquet\" -- IO (TypedDataFrame Trips) +@ + +'readParquet' (and 'readParquetWithOpts'\/'readParquetFiles') throw a +'DataFrameException' on schema mismatch; 'readParquetWithError' returns the +mismatch as an 'Either' instead. +-} +module DataFrame.Typed.IO.Parquet ( + readParquet, + readParquetWithError, + readParquetWithOpts, + readParquetFiles, +) where + +import qualified Data.Text as T + +import DataFrame.IO.Parquet (ParquetReadOptions) +import qualified DataFrame.IO.Parquet as Parquet +import DataFrame.Typed.Freeze (freezeOrThrow, freezeWithError) +import DataFrame.Typed.Schema (KnownSchema) +import DataFrame.Typed.Types (TypedDataFrame) + +-- | Read a Parquet file into a typed DataFrame, throwing on schema mismatch. +readParquet :: + forall cols. (KnownSchema cols) => FilePath -> IO (TypedDataFrame cols) +readParquet path = Parquet.readParquet path >>= freezeOrThrow @cols + +-- | Read a Parquet file, returning a descriptive error on schema mismatch. +readParquetWithError :: + forall cols. + (KnownSchema cols) => + FilePath -> IO (Either T.Text (TypedDataFrame cols)) +readParquetWithError path = freezeWithError <$> Parquet.readParquet path + +-- | Read a Parquet file with custom options, throwing on schema mismatch. +readParquetWithOpts :: + forall cols. + (KnownSchema cols) => + ParquetReadOptions -> FilePath -> IO (TypedDataFrame cols) +readParquetWithOpts opts path = + Parquet.readParquetWithOpts 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.readParquetFiles path >>= freezeOrThrow @cols