diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,20 @@ +Copyright (c) 2026 Michael Chavinda + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be included +in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/dataframe-parquet.cabal b/dataframe-parquet.cabal new file mode 100644 --- /dev/null +++ b/dataframe-parquet.cabal @@ -0,0 +1,70 @@ +cabal-version: 2.4 +name: dataframe-parquet +version: 1.0.0.0 + +synopsis: Parquet reader and writer for the dataframe ecosystem. +description: + @DataFrame.IO.Parquet@ — pure-Haskell Parquet 2.0 reader and writer + (with snappy, zstd, gzip codecs, dictionary decoding, nested + list/repeated columns, predicate pushdown, and HTTP range reads). + Heavy package — pulls in @pinch@, @streamly@, compression codecs, + and @http-conduit@. Most users want @dataframe-csv@ instead unless + they specifically need Parquet. + +bug-reports: https://github.com/mchav/dataframe/issues +license: MIT +license-file: LICENSE +author: Michael Chavinda +maintainer: mschavinda@gmail.com +copyright: (c) 2024-2025 Michael Chavinda +category: Data +tested-with: GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2 + +common warnings + ghc-options: + -Wincomplete-patterns + -Wincomplete-uni-patterns + -Wunused-imports + -Wunused-local-binds + +library + import: warnings + exposed-modules: + DataFrame.IO.Parquet + DataFrame.IO.Parquet.Binary + DataFrame.IO.Parquet.Decompress + DataFrame.IO.Parquet.Dictionary + DataFrame.IO.Parquet.Encoding + DataFrame.IO.Parquet.Levels + DataFrame.IO.Parquet.Page + DataFrame.IO.Parquet.Schema + DataFrame.IO.Parquet.Seeking + DataFrame.IO.Parquet.Thrift + DataFrame.IO.Parquet.Time + DataFrame.IO.Parquet.Utils + DataFrame.IO.Utils.RandomAccess + build-depends: base >= 4 && < 5, + aeson >= 0.11.0.0 && < 3, + attoparsec >= 0.12 && < 0.15, + bytestring >= 0.11 && < 0.13, + containers >= 0.6.7 && < 0.9, + dataframe-core ^>= 1.0, + dataframe-operations ^>= 1.0, + dataframe-parsing ^>= 1.0, + directory >= 1.3.0.0 && < 2, + filepath >= 1.4 && < 2, + Glob >= 0.10 && < 1, + hashable >= 1.2 && < 2, + http-conduit >= 2.3 && < 3, + pinch >= 0.5 && < 1, + snappy-hs ^>= 0.1, + streamly-bytestring >= 0.2.0 && < 0.4, + streamly-core >= 0.2.3 && < 0.4, + text >= 2.0 && < 3, + time >= 1.12 && < 2, + unordered-containers >= 0.1 && < 1, + vector ^>= 0.13, + zlib >= 0.5 && < 1, + zstd >= 0.1.2.0 && < 0.2 + hs-source-dirs: src + default-language: Haskell2010 diff --git a/src/DataFrame/IO/Parquet.hs b/src/DataFrame/IO/Parquet.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet.hs @@ -0,0 +1,740 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE NumericUnderscores #-} +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module DataFrame.IO.Parquet where + +import Control.Exception (throw, try) +import Control.Monad +import Control.Monad.IO.Class (MonadIO (..)) +import Data.Aeson (FromJSON (..), eitherDecodeStrict, withObject, (.:)) +import Data.Bits (Bits (shiftL), (.|.)) +import qualified Data.ByteString as BS +import Data.Either (fromRight) +import Data.Functor ((<&>)) +import Data.Int (Int32, Int64) +#if !MIN_VERSION_base(4,20,0) +import Data.List (foldl', transpose) +#else +import Data.List (transpose) +#endif +import qualified Data.List as L +import qualified Data.Map as Map +import qualified Data.Text as T +import Data.Text.Encoding (encodeUtf8) +import Data.Time (UTCTime) +import Data.Time.Clock.POSIX (posixSecondsToUTCTime) +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, + boolDecoder, + byteArrayDecoder, + doubleDecoder, + fixedLenByteArrayDecoder, + floatDecoder, + int32Decoder, + int64Decoder, + int96Decoder, + readPages, + ) +import DataFrame.IO.Parquet.Seeking ( + FileBufferedOrSeekable, + ForceNonSeekable, + withFileBufferedOrSeekable, + ) +import DataFrame.IO.Parquet.Thrift ( + ColumnChunk (..), + DecimalType (..), + FileMetadata (..), + LogicalType (..), + RowGroup (..), + ThriftType (..), + TimeUnit (..), + TimestampType (..), + unField, + ) +import DataFrame.IO.Parquet.Utils ( + ColumnDescription (..), + foldNonNullable, + foldNonNullableUnboxed, + foldNullable, + foldNullableUnboxed, + foldRepeated, + foldRepeatedUnboxed, + generateColumnDescriptions, + getColumnNames, + ) +import DataFrame.IO.Utils.RandomAccess ( + RandomAccess (..), + ReaderIO (runReaderIO), + ) +import DataFrame.Internal.Column (Column, Columnable) +import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.DataFrame (DataFrame (..)) +import DataFrame.Internal.Expression (Expr, getColumns) +import DataFrame.Operations.Merge () +import qualified DataFrame.Operations.Subset as DS +import Network.HTTP.Simple ( + getResponseBody, + getResponseStatusCode, + httpBS, + parseRequest, + setRequestHeader, + ) +import qualified Pinch +import qualified Streamly.Data.Stream as Stream +import System.Directory ( + doesDirectoryExist, + getHomeDirectory, + getTemporaryDirectory, + ) +import System.Environment (lookupEnv) +import System.FilePath (()) +import System.FilePath.Glob (compile, glob, match) +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 + +Column selection for @selectedColumns@ uses leaf column names only. +-} +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. + -} + , 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. + , safeColumns :: Bool + -- ^ When True, every column is promoted to OptionalColumn after read, regardless of nullability in the schema. + } + deriving (Show) + +{- | Default Parquet read options. + +Equivalent to: + +@ +ParquetReadOptions + { selectedColumns = Nothing + , predicate = Nothing + , rowRange = Nothing + , safeColumns = False + } +@ +-} +defaultParquetReadOptions :: ParquetReadOptions +defaultParquetReadOptions = + ParquetReadOptions + { selectedColumns = Nothing + , predicate = Nothing + , rowRange = Nothing + , safeColumns = False + } + +-- Public API -------------------------------------------------------------- + +{- | Read a parquet file from path and load it into a dataframe. + +==== __Example__ +@ +ghci> D.readParquet ".\/data\/mtcars.parquet" +@ +-} +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. +-} +readParquetWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame +readParquetWithOpts opts path + | isHFUri path = do + paths <- fetchHFParquetFiles path + let optsNoRange = opts{rowRange = Nothing} + dfs <- mapM (_readParquetWithOpts Nothing optsNoRange) paths + pure (applyRowRange opts (mconcat dfs)) + | otherwise = _readParquetWithOpts Nothing opts path + +-- | 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 + +{- | Read Parquet files from a directory or glob path. + +This is equivalent to calling 'readParquetFilesWithOpts' with 'defaultParquetReadOptions'. +-} +readParquetFiles :: FilePath -> IO DataFrame +readParquetFiles = readParquetFilesWithOpts defaultParquetReadOptions + +{- | Read multiple Parquet files (directory or glob) using explicit options. + +If @path@ is a directory, all non-directory entries are read. +If @path@ is a glob, matching files are read. + +For multi-file reads, @rowRange@ is applied once after concatenation (global range semantics). + +==== __Example__ +@ +ghci> D.readParquetFilesWithOpts +ghci| (D.defaultParquetReadOptions{D.selectedColumns = Just ["id"], D.rowRange = Just (0, 5)}) +ghci| "./tests/data/alltypes_plain*.parquet" +@ +-} +readParquetFilesWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame +readParquetFilesWithOpts opts path + | isHFUri path = do + files <- fetchHFParquetFiles path + let optsWithoutRowRange = opts{rowRange = Nothing} + dfs <- mapM (_readParquetWithOpts Nothing optsWithoutRowRange) files + pure (applyRowRange opts (mconcat dfs)) + | otherwise = do + isDir <- doesDirectoryExist path + + let pat = if isDir then path "*.parquet" else path + + matches <- glob pat + + files <- filterM (fmap not . doesDirectoryExist) matches + + case files of + [] -> + error $ + "readParquetFiles: no parquet files found for " ++ path + _ -> do + let optsWithoutRowRange = opts{rowRange = Nothing} + 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'. +-} +parseParquetWithOpts :: + (RandomAccess m, MonadIO m) => + ParquetReadOptions -> + m DataFrame +parseParquetWithOpts opts = do + metadata <- parseFileMetadata + + let schemaElems = unField metadata.schema + allNames = getColumnNames (drop 1 schemaElems) + leafNames = L.nub (map (last . T.splitOn ".") allNames) + predicateColumns = maybe [] (L.nub . getColumns) (predicate opts) + selectedColumnsForRead = case selectedColumns opts of + Nothing -> Nothing + Just selected -> Just (L.nub (selected ++ predicateColumns)) + + -- TODO: When selectedColumnsForRead is Just, pass the set of required + -- column indices into the chunk parsers so that RandomAccess reads are + -- skipped for columns not in the selection, rather than decoding all + -- columns and projecting afterward. + + -- TODO: When rowRange is set, compute cumulative row offsets from + -- rg_num_rows in each RowGroup and skip any group whose row interval does + -- not overlap the requested range, avoiding all decoding for those groups. + + -- 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. + + -- 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 + ) + + let 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 + + rawCols <- zipWithM (parseColumnChunks vectorLength) chunks descriptions + + let finalCols = zipWith applyDescLogicalType descriptions rawCols + indices = Map.fromList $ zip allNames [0 ..] + dimensions = (vectorLength, length finalCols) + + let df = + DataFrame + (Vector.fromListN (length finalCols) finalCols) + indices + dimensions + Map.empty + + return (applyReadOptions opts df) + +{- | Parse the file-level Thrift metadata from the Parquet file footer. +Validates the trailing 4-byte magic marker (\"PAR1\") before decoding. +-} +parseFileMetadata :: (RandomAccess m) => m FileMetadata +parseFileMetadata = do + footerBytes <- readSuffix 8 + let magic = BS.drop 4 footerBytes + when (magic /= "PAR1") $ + error + ( "Not a valid Parquet file: expected magic bytes \"PAR1\", got " + ++ show magic + ) + let size = getMetadataSize footerBytes + rawMetadata <- readSuffix (size + 8) <&> BS.take size + case Pinch.decode Pinch.compactProtocol rawMetadata of + Left e -> error $ "Failed to parse Parquet metadata: " ++ show e + Right metadata -> return metadata + where + getMetadataSize footer = + let sizes :: [Int] + sizes = map (fromIntegral . BS.index footer) [0 .. 3] + in foldl' (.|.) 0 $ zipWith shiftL sizes [0, 8 .. 24] + +-- | Read the file metadata from a Parquet file at the given path. +readMetadataFromPath :: FilePath -> IO FileMetadata +readMetadataFromPath path = + withFileBufferedOrSeekable Nothing path ReadMode $ + runReaderIO parseFileMetadata + +-- | Read only the file metadata from an open 'FileBufferedOrSeekable' handle. +readMetadataFromHandle :: FileBufferedOrSeekable -> IO FileMetadata +readMetadataFromHandle = runReaderIO parseFileMetadata + +-- | Collect column chunks per column (transposed across all row groups). +columnChunksForAll :: FileMetadata -> [[ColumnChunk]] +columnChunksForAll = + transpose . map (unField . rg_columns) . unField . row_groups + +-- | Dispatch a column's chunks to the correct decoder path. +parseColumnChunks :: + (RandomAccess m, MonadIO m) => + Int -> + [ColumnChunk] -> + ColumnDescription -> + m Column +parseColumnChunks totalRows chunks description + | description.maxRepetitionLevel == 0 && description.maxDefinitionLevel == 0 = + getNonNullableColumn totalRows description chunks + | description.maxRepetitionLevel == 0 = + getNullableColumn totalRows description chunks + | otherwise = + getRepeatedColumn description chunks + +-- | Decode a required (non-nullable, non-repeated) column. +getNonNullableColumn :: + forall m. + (RandomAccess m, MonadIO m) => + Int -> + ColumnDescription -> + [ColumnChunk] -> + m Column +getNonNullableColumn totalRows description chunks = + case description.colElementType of + Just (BOOLEAN _) -> unboxedGo boolDecoder + Just (INT32 _) -> unboxedGo int32Decoder + Just (INT64 _) -> unboxedGo int64Decoder + Just (INT96 _) -> go int96Decoder + Just (FLOAT _) -> unboxedGo floatDecoder + Just (DOUBLE _) -> unboxedGo doubleDecoder + Just (BYTE_ARRAY _) -> go byteArrayDecoder + 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)) + Nothing -> error "Column has no Parquet type" + where + go :: + forall a. + (Columnable a) => + PageDecoder a -> + m Column + go decoder = + foldNonNullable totalRows $ + (\(vs, _, _) -> vs) + <$> Stream.unfoldEach (readPages description decoder) (Stream.fromList chunks) + + unboxedGo :: + forall a. + (Columnable a, VU.Unbox a) => + UnboxedPageDecoder a -> + m Column + unboxedGo decoder = + foldNonNullableUnboxed totalRows $ + (\(vs, _, _) -> vs) + <$> Stream.unfoldEach + (readPages description decoder) + (Stream.fromList chunks) + +-- | Decode an optional (nullable) column. +getNullableColumn :: + forall m. + (RandomAccess m, MonadIO m) => + Int -> + ColumnDescription -> + [ColumnChunk] -> + m Column +getNullableColumn totalRows description chunks = + case description.colElementType of + Just (BOOLEAN _) -> unboxedGo boolDecoder + Just (INT32 _) -> unboxedGo int32Decoder + Just (INT64 _) -> unboxedGo int64Decoder + Just (INT96 _) -> go int96Decoder + Just (FLOAT _) -> unboxedGo floatDecoder + Just (DOUBLE _) -> unboxedGo doubleDecoder + Just (BYTE_ARRAY _) -> go byteArrayDecoder + 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)) + Nothing -> error "Column has no Parquet type" + where + maxDef :: Int + maxDef = fromIntegral description.maxDefinitionLevel + + go :: + forall a. + (Columnable a) => + PageDecoder a -> + m Column + go decoder = + foldNullable maxDef totalRows $ + (\(vs, ds, _) -> (vs, ds)) + <$> Stream.unfoldEach (readPages description decoder) (Stream.fromList chunks) + unboxedGo :: + forall a. + (Columnable a, VU.Unbox a) => + UnboxedPageDecoder a -> + m Column + unboxedGo decoder = + foldNullableUnboxed maxDef totalRows $ + (\(vs, ds, _) -> (vs, ds)) + <$> Stream.unfoldEach + (readPages description decoder) + (Stream.fromList chunks) + +-- | Decode a repeated (list/nested) column. +getRepeatedColumn :: + forall m. + (RandomAccess m, MonadIO m) => + ColumnDescription -> + [ColumnChunk] -> + m Column +getRepeatedColumn description chunks = + case description.colElementType of + Just (BOOLEAN _) -> unboxedGo boolDecoder + Just (INT32 _) -> unboxedGo int32Decoder + Just (INT64 _) -> unboxedGo int64Decoder + Just (INT96 _) -> go int96Decoder + Just (FLOAT _) -> unboxedGo floatDecoder + Just (DOUBLE _) -> unboxedGo doubleDecoder + Just (BYTE_ARRAY _) -> go byteArrayDecoder + 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)) + Nothing -> error "Column has no Parquet type" + where + maxRep :: Int + maxRep = fromIntegral description.maxRepetitionLevel + maxDef :: Int + maxDef = fromIntegral description.maxDefinitionLevel + + go :: + forall a. + ( Columnable a + , Columnable (Maybe [Maybe a]) + , Columnable (Maybe [Maybe [Maybe a]]) + , Columnable (Maybe [Maybe [Maybe [Maybe a]]]) + ) => + PageDecoder a -> + m Column + go decoder = + foldRepeated maxRep maxDef $ + Stream.unfoldEach (readPages description decoder) (Stream.fromList chunks) + + unboxedGo :: + forall a. + ( VU.Unbox a + , Columnable a + , Columnable (Maybe [Maybe a]) + , Columnable (Maybe [Maybe [Maybe a]]) + , Columnable (Maybe [Maybe [Maybe [Maybe a]]]) + ) => + UnboxedPageDecoder a -> + m Column + unboxedGo decoder = + foldRepeatedUnboxed maxRep maxDef $ + Stream.unfoldEach + (readPages description decoder) + (Stream.fromList chunks) + +-- Options application ----------------------------------------------------- + +applyRowRange :: ParquetReadOptions -> DataFrame -> DataFrame +applyRowRange opts df = + maybe df (`DS.range` df) (rowRange opts) + +applySelectedColumns :: ParquetReadOptions -> DataFrame -> DataFrame +applySelectedColumns opts df = + maybe df (`DS.select` df) (selectedColumns opts) + +applyPredicate :: ParquetReadOptions -> DataFrame -> DataFrame +applyPredicate opts df = + maybe df (`DS.filterWhere` df) (predicate opts) + +applySafeRead :: ParquetReadOptions -> DataFrame -> DataFrame +applySafeRead opts df + | safeColumns opts = df{columns = Vector.map DI.ensureOptional (columns df)} + | otherwise = df + +applyReadOptions :: ParquetReadOptions -> DataFrame -> DataFrame +applyReadOptions opts = + applySafeRead opts + . applyRowRange opts + . 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'). +-} +applyDescLogicalType :: ColumnDescription -> DI.Column -> DI.Column +applyDescLogicalType desc = applyLogicalType (colLogicalType desc) + +applyLogicalType :: Maybe LogicalType -> DI.Column -> DI.Column +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 +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 + 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 + else col +applyLogicalType _ col = col + +microsecondsToUTCTime :: Int64 -> UTCTime +microsecondsToUTCTime us = + posixSecondsToUTCTime (fromIntegral us / 1_000_000) + +-- HuggingFace support ----------------------------------------------------- + +data HFRef = HFRef + { hfOwner :: T.Text + , hfDataset :: T.Text + , hfGlob :: T.Text + } + +data HFParquetFile = HFParquetFile + { hfpUrl :: T.Text + , hfpConfig :: T.Text + , hfpSplit :: T.Text + , hfpFilename :: T.Text + } + deriving (Show) + +instance FromJSON HFParquetFile where + parseJSON = withObject "HFParquetFile" $ \o -> + HFParquetFile + <$> o .: "url" + <*> o .: "config" + <*> o .: "split" + <*> o .: "filename" + +newtype HFParquetResponse = HFParquetResponse {hfParquetFiles :: [HFParquetFile]} + +instance FromJSON HFParquetResponse where + parseJSON = withObject "HFParquetResponse" $ \o -> + HFParquetResponse <$> o .: "parquet_files" + +isHFUri :: FilePath -> Bool +isHFUri = L.isPrefixOf "hf://" + +parseHFUri :: FilePath -> Either String HFRef +parseHFUri path = + let stripped = drop (length ("hf://datasets/" :: String)) path + in case T.splitOn "/" (T.pack stripped) of + (owner : dataset : rest) + | not (null rest) -> + Right $ HFRef owner dataset (T.intercalate "/" rest) + _ -> + Left $ "Invalid hf:// URI (expected hf://datasets/owner/dataset/glob): " ++ path + +getHFToken :: IO (Maybe BS.ByteString) +getHFToken = do + envToken <- lookupEnv "HF_TOKEN" + case envToken of + Just t -> pure (Just (encodeUtf8 (T.pack t))) + Nothing -> do + home <- getHomeDirectory + let tokenPath = home ".cache" "huggingface" "token" + result <- try (BS.readFile tokenPath) :: IO (Either IOError BS.ByteString) + case result of + Right bs -> pure (Just (BS.takeWhile (/= 10) bs)) + Left _ -> pure Nothing + +{- | Extract the repo-relative path from a HuggingFace download URL. +URL format: https://huggingface.co/datasets/{owner}/{dataset}/resolve/{ref}/{path} +Returns the {path} portion (e.g. "data/train-00000-of-00001.parquet"). +-} +hfUrlRepoPath :: HFParquetFile -> String +hfUrlRepoPath f = + case T.breakOn "/resolve/" (hfpUrl f) of + (_, rest) + | not (T.null rest) -> + -- Drop "/resolve/", then drop the ref component (up to and including "/") + T.unpack $ T.drop 1 $ T.dropWhile (/= '/') $ T.drop (T.length "/resolve/") rest + _ -> + T.unpack (hfpConfig f) T.unpack (hfpSplit f) T.unpack (hfpFilename f) + +matchesGlob :: T.Text -> HFParquetFile -> Bool +matchesGlob g f = match (compile (T.unpack g)) (hfUrlRepoPath f) + +resolveHFUrls :: Maybe BS.ByteString -> HFRef -> IO [HFParquetFile] +resolveHFUrls mToken ref = do + let dataset = hfOwner ref <> "/" <> hfDataset ref + let apiUrl = "https://datasets-server.huggingface.co/parquet?dataset=" ++ T.unpack dataset + req0 <- parseRequest apiUrl + let req = case mToken of + Nothing -> req0 + Just tok -> setRequestHeader "Authorization" ["Bearer " <> tok] req0 + resp <- httpBS req + let status = getResponseStatusCode resp + when (status /= 200) $ + ioError $ + userError $ + "HuggingFace API returned status " + ++ show status + ++ " for dataset " + ++ T.unpack dataset + case eitherDecodeStrict (getResponseBody resp) of + Left err -> ioError $ userError $ "Failed to parse HF API response: " ++ err + Right hfResp -> pure $ filter (matchesGlob (hfGlob ref)) (hfParquetFiles hfResp) + +downloadHFFiles :: Maybe BS.ByteString -> [HFParquetFile] -> IO [FilePath] +downloadHFFiles mToken files = do + tmpDir <- getTemporaryDirectory + forM files $ \f -> do + -- Derive a collision-resistant temp name from the URL path components + let fname = case (hfpConfig f, hfpSplit f) of + (c, s) | T.null c && T.null s -> T.unpack (hfpFilename f) + (c, s) -> T.unpack c <> "_" <> T.unpack s <> "_" <> T.unpack (hfpFilename f) + let destPath = tmpDir fname + req0 <- parseRequest (T.unpack (hfpUrl f)) + let req = case mToken of + Nothing -> req0 + Just tok -> setRequestHeader "Authorization" ["Bearer " <> tok] req0 + resp <- httpBS req + let status = getResponseStatusCode resp + when (status /= 200) $ + ioError $ + userError $ + "Failed to download " ++ T.unpack (hfpUrl f) ++ " (HTTP " ++ show status ++ ")" + BS.writeFile destPath (getResponseBody resp) + pure destPath + +-- | True when the path contains glob wildcard characters. +hasGlob :: T.Text -> Bool +hasGlob = T.any (\c -> c == '*' || c == '?' || c == '[') + +{- | Build the direct HF repo download URL for a path with no wildcards. +Format: https://huggingface.co/datasets/{owner}/{dataset}/resolve/main/{path} +-} +directHFUrl :: HFRef -> T.Text +directHFUrl ref = + "https://huggingface.co/datasets/" + <> hfOwner ref + <> "/" + <> hfDataset ref + <> "/resolve/main/" + <> hfGlob ref + +fetchHFParquetFiles :: FilePath -> IO [FilePath] +fetchHFParquetFiles uri = do + ref <- case parseHFUri uri of + Left err -> ioError (userError err) + Right r -> pure r + mToken <- getHFToken + if hasGlob (hfGlob ref) + then do + hfFiles <- resolveHFUrls mToken ref + when (null hfFiles) $ + ioError $ + userError $ + "No parquet files found for " ++ uri + downloadHFFiles mToken hfFiles + else do + -- Direct repo file download — no datasets-server needed + let url = directHFUrl ref + let filename = last $ T.splitOn "/" (hfGlob ref) + downloadHFFiles mToken [HFParquetFile url "" "" filename] diff --git a/src/DataFrame/IO/Parquet/Binary.hs b/src/DataFrame/IO/Parquet/Binary.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Binary.hs @@ -0,0 +1,141 @@ +{-# LANGUAGE TypeApplications #-} + +module DataFrame.IO.Parquet.Binary where + +import Control.Exception (bracketOnError) +import Control.Monad +import Data.Bits +import qualified Data.ByteString as BS +import qualified Data.ByteString.Unsafe as BSU +import Data.Char +import Data.IORef +import Data.Int +import Data.Word +import qualified Foreign.Marshal.Alloc as Foreign +import qualified Foreign.Ptr as Foreign +import qualified Foreign.Storable as Foreign + +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 + -} + 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 + Nothing -> error "readUVarInt: not enough input bytes" + Just (b, bs') -> + if b < 0x80 + then (result .|. (fromIntegral b `shiftL` shiftAmt), bs') + else + let payloadBits = fromIntegral (b .&. 0x7f) :: Word64 + in loop bs' (result .|. (payloadBits `shiftL` shiftAmt)) (shiftAmt + 7) (i + 1) + +readVarIntFromBytes :: (Integral a) => BS.ByteString -> (a, BS.ByteString) +readVarIntFromBytes bs = (fromIntegral n, remainder) + where + (n, remainder) = loop 0 0 bs + loop shiftAmt result bs' = case BS.uncons bs' of + Nothing -> (result, BS.empty) + Just (x, xs) -> + let res = result .|. (fromIntegral (x .&. 0x7f) :: Integer) `shiftL` shiftAmt + in if x .&. 0x80 /= 0x80 then (res, xs) else loop (shiftAmt + 7) res xs + +readIntFromBytes :: (Integral a) => BS.ByteString -> (a, BS.ByteString) +readIntFromBytes bs = + let (n, remainder) = readVarIntFromBytes bs + u = fromIntegral n :: Word32 + in ( fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)) + , remainder + ) + +readInt32FromBytes :: BS.ByteString -> (Int32, BS.ByteString) +readInt32FromBytes bs = + let (n', remainder) = readVarIntFromBytes @Int64 bs + n = fromIntegral n' :: Int32 + u = fromIntegral n :: Word32 + in ((fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)), remainder) + +readAndAdvance :: IORef Int -> BS.ByteString -> IO Word8 +readAndAdvance bufferPos buffer = do + pos <- readIORef bufferPos + let b = BS.index buffer pos + modifyIORef' bufferPos (+ 1) + return b + +readVarIntFromBuffer :: (Integral a) => BS.ByteString -> IORef Int -> IO a +readVarIntFromBuffer buf bufferPos = do + start <- readIORef bufferPos + let loop i shiftAmt result = do + b <- readAndAdvance bufferPos buf + let res = result .|. (fromIntegral (b .&. 0x7f) :: Integer) `shiftL` shiftAmt + if b .&. 0x80 /= 0x80 + then return res + else loop (i + 1) (shiftAmt + 7) res + fromIntegral <$> loop start 0 0 + +readIntFromBuffer :: (Integral a) => BS.ByteString -> IORef Int -> IO a +readIntFromBuffer buf bufferPos = do + n <- readVarIntFromBuffer buf bufferPos + let u = fromIntegral n :: Word32 + return $ fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)) + +readInt32FromBuffer :: BS.ByteString -> IORef Int -> IO Int32 +readInt32FromBuffer buf bufferPos = do + n <- (fromIntegral <$> readVarIntFromBuffer @Int64 buf bufferPos) :: IO Int32 + let u = fromIntegral n :: Word32 + return $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)) + +readString :: BS.ByteString -> IORef Int -> IO String +readString buf pos = do + nameSize <- readVarIntFromBuffer @Int buf pos + replicateM nameSize (chr . fromIntegral <$> readAndAdvance pos buf) + +readByteStringFromBytes :: BS.ByteString -> (BS.ByteString, BS.ByteString) +readByteStringFromBytes xs = + let + (size, remainder) = readVarIntFromBytes @Int xs + in + BS.splitAt size remainder + +readByteString :: BS.ByteString -> IORef Int -> IO BS.ByteString +readByteString buf pos = do + size <- readVarIntFromBuffer @Int buf pos + fillByteStringByWord8 size (\_ -> readAndAdvance pos buf) + +readByteString' :: BS.ByteString -> Int64 -> IO BS.ByteString +readByteString' buf size = + fillByteStringByWord8 + (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. +-} +fillByteStringByWord8 :: Int -> (Int -> IO Word8) -> IO BS.ByteString +fillByteStringByWord8 size getByte = do + bracketOnError + (Foreign.mallocBytes size :: IO (Foreign.Ptr Word8)) + Foreign.free + -- \^ ensures p is freed if (IO Word8) throws. + ( \p -> do + fill 0 p + BSU.unsafePackCStringFinalizer p size (Foreign.free p) + ) + where + fill i p + | i >= size = pure () + | otherwise = getByte i >>= Foreign.pokeByteOff p i >> fill (i + 1) p +{-# INLINE fillByteStringByWord8 #-} + +readSingleByte :: Int64 -> BS.ByteString -> IO Word8 +readSingleByte pos buffer = return $ BS.index buffer (fromIntegral pos) + +readNoAdvance :: IORef Int -> BS.ByteString -> IO Word8 +readNoAdvance bufferPos buffer = do + pos <- readIORef bufferPos + return $ BS.index buffer pos diff --git a/src/DataFrame/IO/Parquet/Decompress.hs b/src/DataFrame/IO/Parquet/Decompress.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Decompress.hs @@ -0,0 +1,32 @@ +module DataFrame.IO.Parquet.Decompress where + +import qualified Codec.Compression.GZip as GZip +import qualified Codec.Compression.Zstd.Base as Zstd +import qualified Data.ByteString as BS +import qualified Data.ByteString as LB +import Data.ByteString.Internal (createAndTrim, toForeignPtr) +import DataFrame.IO.Parquet.Thrift (CompressionCodec (..)) +import Foreign.ForeignPtr (withForeignPtr) +import Foreign.Ptr (plusPtr) +import qualified Snappy + +decompressData :: Int -> CompressionCodec -> BS.ByteString -> IO BS.ByteString +decompressData uncompressedSize codec compressed = case codec of + (ZSTD _) -> createAndTrim uncompressedSize $ \dstPtr -> + let (srcFP, offset, compressedSize) = toForeignPtr compressed + in withForeignPtr srcFP $ \srcPtr -> do + result <- + Zstd.decompress + dstPtr + uncompressedSize + (srcPtr `plusPtr` offset) + compressedSize + case result of + Left e -> error $ "ZSTD error: " <> e + Right actualSize -> return actualSize + (SNAPPY _) -> case Snappy.decompress compressed of + Left e -> error (show e) + Right res -> pure res + (UNCOMPRESSED _) -> pure compressed + (GZIP _) -> pure (LB.toStrict (GZip.decompress (BS.fromStrict compressed))) + other -> error ("Unsupported compression type: " <> show other) diff --git a/src/DataFrame/IO/Parquet/Dictionary.hs b/src/DataFrame/IO/Parquet/Dictionary.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Dictionary.hs @@ -0,0 +1,164 @@ +{-# LANGUAGE BangPatterns #-} + +module DataFrame.IO.Parquet.Dictionary (DictVals (..), readDictVals, decodeRLEBitPackedHybrid) where + +import Data.Bits +import qualified Data.ByteString as BS +import qualified Data.ByteString.Unsafe as BSU +import Data.Int (Int32, Int64) +import qualified Data.Text as T +import Data.Text.Encoding +import Data.Time (UTCTime) +import qualified Data.Vector as V +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 GHC.Float + +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) + 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. +-} +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 (FIXED_LEN_BYTE_ARRAY _) bs _ (Just len) = + DText (V.fromList (readPageFixedBytes bs (fromIntegral len))) +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) + +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) + +readPageBool :: BS.ByteString -> [Bool] +readPageBool bs = + concatMap (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7]) (BS.unpack bs) + +readPageInt64 :: BS.ByteString -> [Int64] +readPageInt64 xs + | BS.null xs = [] + | otherwise = + fromIntegral (littleEndianWord64 (BS.take 8 xs)) : readPageInt64 (BS.drop 8 xs) + +readPageFloat :: BS.ByteString -> [Float] +readPageFloat xs + | BS.null xs = [] + | otherwise = + castWord32ToFloat (littleEndianWord32 (BS.take 4 xs)) + : readPageFloat (BS.drop 4 xs) + +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) + +readPageInt96Times :: BS.ByteString -> [UTCTime] +readPageInt96Times bs + | BS.null bs = [] + | otherwise = + let (times, _) = readNInt96Times (BS.length bs `div` 12) bs + in times + +readPageFixedBytes :: BS.ByteString -> Int -> [T.Text] +readPageFixedBytes xs len + | BS.null xs = [] + | otherwise = + decodeUtf8Lenient (BS.take len xs) : readPageFixedBytes (BS.drop len xs) len + +unpackBitPacked :: Int -> Int -> BS.ByteString -> ([Word32], BS.ByteString) +unpackBitPacked bw count bs + | count <= 0 = ([], bs) + | BS.null bs = ([], bs) + | otherwise = + let totalBytes = (bw * count + 7) `div` 8 + chunk = BS.take totalBytes bs + rest = BS.drop totalBytes bs + in (extractBits bw count chunk, rest) + +-- | LSB-first bit accumulator: reads each byte once with no intermediate ByteString allocation. +extractBits :: Int -> Int -> BS.ByteString -> [Word32] +extractBits bw count bs = go 0 (0 :: Word64) 0 count + where + !mask = if bw == 32 then maxBound else (1 `shiftL` bw) - 1 :: Word64 + !len = BS.length bs + go !byteIdx !acc !accBits !remaining + | remaining <= 0 = [] + | accBits >= bw = + fromIntegral (acc .&. mask) + : go byteIdx (acc `shiftR` bw) (accBits - bw) (remaining - 1) + | byteIdx >= len = [] + | otherwise = + let b = fromIntegral (BSU.unsafeIndex bs byteIdx) :: Word64 + in go (byteIdx + 1) (acc .|. (b `shiftL` accBits)) (accBits + 8) remaining + +decodeRLEBitPackedHybrid :: Int -> BS.ByteString -> ([Word32], BS.ByteString) +decodeRLEBitPackedHybrid bitWidth bs + | 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 #-}. + let (hdr64, afterHdr) = readUVarInt bs + isPacked = (hdr64 .&. 1) == 1 + in if isPacked + then + let groups = fromIntegral (hdr64 `shiftR` 1) :: Int + totalVals = groups * 8 + in unpackBitPacked bitWidth totalVals afterHdr + else + let mask = if bitWidth == 32 then maxBound else (1 `shiftL` bitWidth) - 1 + runLen = fromIntegral (hdr64 `shiftR` 1) :: Int + nBytes = (bitWidth + 7) `div` 8 :: Int + word32 = littleEndianWord32 (BS.take 4 afterHdr) + value = word32 .&. mask + in (replicate runLen value, BS.drop nBytes afterHdr) diff --git a/src/DataFrame/IO/Parquet/Encoding.hs b/src/DataFrame/IO/Parquet/Encoding.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Encoding.hs @@ -0,0 +1,132 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} + +module DataFrame.IO.Parquet.Encoding ( + -- Kept from the original Encoding module (used by Levels) + ceilLog2, + bitWidthForMaxLevel, + -- Vector-based RLE/bit-packed decoder (from new parser) + decodeRLEBitPackedHybrid, + extractBitsInto, + fillRun, + decodeDictIndices, +) where + +import Control.Monad.ST (ST, runST) +import Data.Bits +import qualified Data.ByteString as BS +import qualified Data.ByteString.Unsafe as BSU +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM +import Data.Word +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 + | otherwise = 1 + ceilLog2 ((x + 1) `div` 2) + +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 -> + -- | Exact number of values to decode + Int -> + BS.ByteString -> + (VU.Vector Word32, BS.ByteString) +decodeRLEBitPackedHybrid bw need bs + | bw == 0 = (VU.replicate need 0, bs) + | otherwise = runST $ do + mv <- VUM.new need + rest <- go mv 0 bs + dat <- VU.unsafeFreeze mv + return (dat, rest) + where + !mask = if bw == 32 then maxBound else (1 `shiftL` bw) - 1 :: Word32 + go :: VUM.STVector s Word32 -> Int -> BS.ByteString -> ST s BS.ByteString + go mv !filled !buf + | filled >= need = return buf + | BS.null buf = return buf + | 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 + -- Consume all the bytes for this group even if we + -- only need a subset of the values. + bytesN = (bw * totalVals + 7) `div` 8 + (chunk, rest) = BS.splitAt bytesN afterHdr + extractBitsInto bw takeN chunk mv filled + go mv (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 + -- Fill the run directly — no list, no reverse. + fillRun mv filled (filled + takeN) val + go mv (filled + takeN) (BS.drop nbytes afterHdr) +{-# INLINE decodeRLEBitPackedHybrid #-} + +-- | Fill @mv[start..end-1]@ with @val@ using a bulk @memset@-style write. +fillRun :: VUM.STVector s Word32 -> Int -> Int -> Word32 -> ST s () +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. +-} +extractBitsInto :: + -- | Bit width + Int -> + -- | Number of values to extract + Int -> + BS.ByteString -> + VUM.STVector s Word32 -> + -- | Write offset into @mv@ + Int -> + ST s () +extractBitsInto bw count bs mv off = 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 + VUM.unsafeWrite mv (off + 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 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. +-} +decodeDictIndices :: Int -> BS.ByteString -> (VU.Vector Int, BS.ByteString) +decodeDictIndices need bs = case BS.uncons bs of + Nothing -> error "decodeDictIndices: empty stream" + Just (w0, rest0) -> + let bw = fromIntegral w0 :: Int + (raw, rest1) = decodeRLEBitPackedHybrid bw need rest0 + in (VU.map fromIntegral raw, rest1) +{-# INLINE decodeDictIndices #-} diff --git a/src/DataFrame/IO/Parquet/Levels.hs b/src/DataFrame/IO/Parquet/Levels.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Levels.hs @@ -0,0 +1,210 @@ +{-# LANGUAGE BangPatterns #-} + +module DataFrame.IO.Parquet.Levels ( + -- Level readers + readLevelsV1, + readLevelsV2, + -- Stitch functions + stitchList, + stitchList2, + stitchList3, +) where + +import Control.Monad.ST (runST) +import qualified Data.ByteString as BS +import Data.Int (Int32) +import qualified Data.Vector as VB +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM +import Data.Word (Word32) +import DataFrame.IO.Parquet.Encoding ( + bitWidthForMaxLevel, + decodeRLEBitPackedHybrid, + ) +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@. +-} +convertAndCount :: Int -> VU.Vector Word32 -> (VU.Vector Int, Int) +convertAndCount maxDef raw = runST $ do + let !n = VU.length raw + mv <- VUM.unsafeNew n + let !maxDefW = fromIntegral maxDef :: Word32 + go !i !nPresent + | i >= n = pure nPresent + | otherwise = do + let !w = VU.unsafeIndex raw i + !d = fromIntegral w :: Int + VUM.unsafeWrite mv i d + if w == maxDefW + then go (i + 1) (nPresent + 1) + else go (i + 1) nPresent + !nPresent <- go 0 0 + !out <- VU.unsafeFreeze mv + pure (out, nPresent) + +readLevelsV1 :: + -- | Total number of values in the page + Int -> + -- | maxDefinitionLevel + Int -> + -- | maxRepetitionLevel + Int -> + BS.ByteString -> + (VU.Vector Int, VU.Vector Int, Int, BS.ByteString) +readLevelsV1 n maxDef maxRep bs = + let bwRep = bitWidthForMaxLevel maxRep + bwDef = bitWidthForMaxLevel maxDef + (repVec, _, afterRep) = decodeLevelBlock bwRep n bs + (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) + decodeLevelBlock bw n' buf = + let blockLen = fromIntegral (littleEndianWord32 (BS.take 4 buf)) :: Int + blockData = BS.take blockLen (BS.drop 4 buf) + after = BS.drop (4 + blockLen) buf + (raw, _) = decodeRLEBitPackedHybrid bw n' blockData + (out, np) = convertAndCount maxDef raw + in (out, np, after) + +readLevelsV2 :: + -- | Total number of values + Int -> + -- | maxDefinitionLevel + Int -> + -- | maxRepetitionLevel + Int -> + -- | Repetition-level byte length (from page header) + Int32 -> + -- | Definition-level byte length (from page header) + Int32 -> + BS.ByteString -> + (VU.Vector Int, VU.Vector Int, Int, BS.ByteString) +readLevelsV2 n maxDef maxRep repLen defLen bs = + let (repBytes, afterRepBytes) = BS.splitAt (fromIntegral repLen) bs + (defBytes, afterDefBytes) = BS.splitAt (fromIntegral defLen) afterRepBytes + bwRep = bitWidthForMaxLevel maxRep + bwDef = bitWidthForMaxLevel maxDef + repVec + | bwRep == 0 = VU.replicate n 0 + | 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)) + | 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. +-} +stitchList :: + Int -> + VU.Vector Int -> + VU.Vector Int -> + VB.Vector a -> + [Maybe [Maybe a]] +stitchList maxDef repVec defVec values = + map toRow (splitAtRepBound 0 (pairWithValsV maxDef repVec defVec values)) + where + toRow [] = Nothing + 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. +-} +stitchList2 :: + Int -> + Int -> + VU.Vector Int -> + VU.Vector Int -> + VB.Vector a -> + [Maybe [Maybe [Maybe a]]] +stitchList2 defT1 maxDef repVec defVec values = + map toRow (splitAtRepBound 0 triplets) + where + triplets = pairWithValsV maxDef repVec defVec values + toRow [] = Nothing + toRow ((_, d, _) : _) | d == 0 = Nothing + toRow row = Just (map toOuter (splitAtRepBound 1 row)) + toOuter [] = Nothing + toOuter ((_, d, _) : _) | d < defT1 = Nothing + toOuter outer = Just (map toLeaf (splitAtRepBound 2 outer)) + 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. +-} +stitchList3 :: + Int -> + Int -> + Int -> + VU.Vector Int -> + VU.Vector Int -> + VB.Vector a -> + [Maybe [Maybe [Maybe [Maybe a]]]] +stitchList3 defT1 defT2 maxDef repVec defVec values = + map toRow (splitAtRepBound 0 triplets) + where + triplets = pairWithValsV maxDef repVec defVec values + toRow [] = Nothing + toRow ((_, d, _) : _) | d == 0 = Nothing + toRow row = Just (map toOuter (splitAtRepBound 1 row)) + toOuter [] = Nothing + toOuter ((_, d, _) : _) | d < defT1 = Nothing + toOuter outer = Just (map toMiddle (splitAtRepBound 2 outer)) + toMiddle [] = Nothing + toMiddle ((_, d, _) : _) | d < defT2 = Nothing + toMiddle middle = Just (map toLeaf (splitAtRepBound 3 middle)) + 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. +-} +pairWithValsV :: + Int -> + VU.Vector Int -> + VU.Vector Int -> + VB.Vector a -> + [(Int, Int, Maybe a)] +pairWithValsV maxDef repVec defVec values = go 0 0 + where + n = VU.length defVec + go i j + | i >= n = [] + | otherwise = + let r = VU.unsafeIndex repVec i + d = VU.unsafeIndex defVec i + in if d == maxDef + 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@. +-} +splitAtRepBound :: Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]] +splitAtRepBound _ [] = [] +splitAtRepBound bound (t : ts) = + let (rest, remaining) = span (\(r, _, _) -> r > bound) ts + in (t : rest) : splitAtRepBound bound remaining diff --git a/src/DataFrame/IO/Parquet/Page.hs b/src/DataFrame/IO/Parquet/Page.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Page.hs @@ -0,0 +1,352 @@ +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module DataFrame.IO.Parquet.Page ( + -- Types + PageDecoder, + UnboxedPageDecoder, + -- Per-type decoders + boolDecoder, + int32Decoder, + int64Decoder, + int96Decoder, + floatDecoder, + doubleDecoder, + byteArrayDecoder, + fixedLenByteArrayDecoder, + -- Page iteration + readPages, +) where + +import Control.Monad.IO.Class (MonadIO (liftIO)) +import Data.Bits (shiftR, (.&.)) +import qualified Data.ByteString as BS +import Data.Int (Int32, Int64) +import Data.Maybe (fromJust, fromMaybe) +import qualified Data.Text as T +import Data.Text.Encoding (decodeUtf8Lenient) +import Data.Time (UTCTime) +import qualified Data.Vector as VB +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Unboxed as VU +import DataFrame.IO.Parquet.Decompress (decompressData) +import DataFrame.IO.Parquet.Dictionary ( + DictVals (..), + readDictVals, + ) +import DataFrame.IO.Parquet.Encoding (decodeDictIndices) +import DataFrame.IO.Parquet.Levels (readLevelsV1, readLevelsV2) +import DataFrame.IO.Parquet.Thrift ( + ColumnChunk (..), + ColumnMetaData (..), + CompressionCodec, + DataPageHeader (..), + DataPageHeaderV2 (..), + DictionaryPageHeader (..), + Encoding (..), + PageHeader (..), + PageType (..), + ThriftType (..), + unField, + ) +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 GHC.Float (castWord32ToFloat, castWord64ToDouble) +import Pinch (decodeWithLeftovers) +import qualified Pinch +import Streamly.Internal.Data.Unfold (Step (..), Unfold, mkUnfoldM) + +-- --------------------------------------------------------------------------- +-- 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. +-} +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 + _ -> error ("boolDecoder: unsupported encoding " ++ show enc) + where + getBool (DBool ds) i = ds VB.! i + getBool d _ = error ("boolDecoder: wrong dict type, got " ++ show d) + +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 + _ -> error ("int32Decoder: unsupported encoding " ++ show enc) + where + getInt32 (DInt32 ds) i = ds VB.! i + getInt32 d _ = error ("int32Decoder: wrong dict type, got " ++ show d) + +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 + _ -> error ("int64Decoder: unsupported encoding " ++ show enc) + where + getInt64 (DInt64 ds) i = ds VB.! i + getInt64 d _ = error ("int64Decoder: wrong dict type, got " ++ show d) + +int96Decoder :: PageDecoder UTCTime +int96Decoder mDict enc nPresent bs = case enc of + PLAIN _ -> VB.fromList (readNInt96 nPresent bs) + RLE_DICTIONARY _ -> lookupDict mDict nPresent bs getInt96 + PLAIN_DICTIONARY _ -> lookupDict mDict nPresent bs getInt96 + _ -> error ("int96Decoder: unsupported encoding " ++ show enc) + where + getInt96 (DInt96 ds) i = ds VB.! 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 + _ -> error ("floatDecoder: unsupported encoding " ++ show enc) + where + getFloat (DFloat ds) i = ds VB.! i + getFloat d _ = error ("floatDecoder: wrong dict type, got " ++ show d) + +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 + _ -> error ("doubleDecoder: unsupported encoding " ++ show enc) + where + getDouble (DDouble ds) i = ds VB.! i + getDouble d _ = error ("doubleDecoder: wrong dict type, got " ++ show d) + +byteArrayDecoder :: PageDecoder T.Text +byteArrayDecoder mDict enc nPresent bs = case enc of + PLAIN _ -> VB.fromList (readNTexts nPresent bs) + RLE_DICTIONARY _ -> lookupDict mDict nPresent bs getText + 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) + +fixedLenByteArrayDecoder :: Int -> PageDecoder T.Text +fixedLenByteArrayDecoder len mDict enc nPresent bs = case enc of + PLAIN _ -> VB.fromList (readNFixedTexts len nPresent bs) + RLE_DICTIONARY _ -> lookupDict mDict nPresent bs getText + 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) + +{- | Shared dictionary-path helper: decode @nPresent@ RLE/bit-packed indices +and look each one up in the dictionary. +-} +lookupDict :: + Maybe DictVals -> + Int -> + BS.ByteString -> + (DictVals -> Int -> a) -> + VB.Vector a +lookupDict 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 VB.generate nPresent (f dict . VU.unsafeIndex idxs) + +unboxedLookupDict :: + (VU.Unbox 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 +-- --------------------------------------------------------------------------- + +-- | Read the raw (compressed) byte range for a column chunk. +readChunkBytes :: + (RandomAccess m) => + ColumnChunk -> + m (CompressionCodec, ThriftType, BS.ByteString) +readChunkBytes columnChunk = do + let meta = fromJust . unField $ columnChunk.cc_meta_data + codec = unField meta.cmd_codec + pType = unField meta.cmd_type + dataOffset = fromIntegral . unField $ meta.cmd_data_page_offset + dictOffset = fromIntegral <$> unField meta.cmd_dictionary_page_offset + offset = fromMaybe dataOffset dictOffset + compLen = fromIntegral . unField $ meta.cmd_total_compressed_size + rawBytes <- readBytes (Range offset compLen) + return (codec, pType, rawBytes) + +{- | An 'Unfold' from a 'ColumnChunk' to per-page value triples. + +The seed is a 'ColumnChunk'. The inject step reads the chunk's compressed +bytes and discovers the codec and physical type from the column metadata. +Codec and type are then threaded through the unfold state along with the +running dictionary and remaining bytes, so no intermediate list or +concatenation step is needed. Use with 'Stream.unfoldEach' to produce a +flat stream of per-page results directly from a stream of column chunks. + +Dictionary pages are consumed silently and update the running dictionary +that is threaded through the unfold state. + +The internal state is +@(Maybe DictVals, BS.ByteString, CompressionCodec, ThriftType)@. + +-- TODO: when a page index is available, use it here to compute which page +-- byte ranges to request from the RandomAccess layer instead of reading the +-- entire column chunk in one contiguous read. + +-- TODO: accept an optional row-range and use the column/offset page index +-- (when present in file metadata) to Skip pages whose row range does not +-- overlap the requested range, avoiding decompression of irrelevant pages +-- entirely. +-} +readPages :: + (RandomAccess m, MonadIO m, VG.Vector v a) => + ColumnDescription -> + (Maybe DictVals -> Encoding -> Int -> BS.ByteString -> v a) -> + Unfold m ColumnChunk (v a, VU.Vector Int, VU.Vector Int) +readPages description decoder = mkUnfoldM step inject + where + maxDef = fromIntegral description.maxDefinitionLevel :: Int + maxRep = fromIntegral description.maxRepetitionLevel :: Int + + -- Inject: read chunk bytes; put codec and pType into state. + inject cc = do + (codec, pType, rawBytes) <- readChunkBytes cc + return (Nothing, rawBytes, codec, pType) + + step (dict, bs, codec, pType) + | BS.null bs = return Stop + | otherwise = case parsePageHeader bs of + Left e -> error ("readPages: 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 + return $ Skip (Just d, rest', codec, pType) + 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 + let (defLvls, repLvls, nPresent, valBytes) = + readLevelsV1 n maxDef maxRep decompressed + triple = (decoder dict enc nPresent valBytes, defLvls, repLvls) + return $ Yield triple (dict, rest', codec, pType) + 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 + 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. + isCompressed = fromMaybe True (unField dph2.dph2_is_compressed) + (defLvls, repLvls, nPresent, compValBytes) = + readLevelsV2 n maxDef maxRep repLen defLen pageData + valBytes <- + if isCompressed + then liftIO $ decompressData uncmpSz codec compValBytes + else pure compValBytes + let triple = (decoder dict enc nPresent valBytes, defLvls, repLvls) + return $ Yield triple (dict, rest', codec, pType) + INDEX_PAGE _ -> return $ Skip (dict, rest', codec, pType) + +-- --------------------------------------------------------------------------- +-- Page header parsing +-- --------------------------------------------------------------------------- + +parsePageHeader :: BS.ByteString -> Either String (BS.ByteString, PageHeader) +parsePageHeader = decodeWithLeftovers Pinch.compactProtocol + +-- --------------------------------------------------------------------------- +-- 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) + +readNInt64 :: Int -> BS.ByteString -> VU.Vector Int64 +readNInt64 n bs = VU.generate n $ \i -> + fromIntegral (littleEndianWord64 (BS.drop (8 * i) bs)) + +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)) + +readNDouble :: Int -> BS.ByteString -> VU.Vector Double +readNDouble n bs = VU.generate n $ \i -> + castWord64ToDouble (littleEndianWord64 (BS.drop (8 * i) bs)) + +readNTexts :: Int -> BS.ByteString -> [T.Text] +readNTexts 0 _ = [] +readNTexts n bs = + let len = fromIntegral . littleEndianInt32 . BS.take 4 $ bs + text = decodeUtf8Lenient . BS.take len . BS.drop 4 $ bs + in text : readNTexts (n - 1) (BS.drop (4 + len) bs) + +readNFixedTexts :: Int -> Int -> BS.ByteString -> [T.Text] +readNFixedTexts _ 0 _ = [] +readNFixedTexts len n bs = + decodeUtf8Lenient (BS.take len bs) + : readNFixedTexts len (n - 1) (BS.drop len bs) diff --git a/src/DataFrame/IO/Parquet/Schema.hs b/src/DataFrame/IO/Parquet/Schema.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Schema.hs @@ -0,0 +1,85 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedRecordDot #-} +{-# 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'. +-} +module DataFrame.IO.Parquet.Schema ( + schemaToEmptyDataFrame, + schemaElemToColumn, + emptyColumnForType, + emptyNullableColumnForType, +) where + +import Data.Int (Int32, Int64) +import qualified Data.Maybe as Maybe +import qualified Data.Set as S +import qualified Data.Text as T + +import DataFrame.IO.Parquet.Thrift ( + SchemaElement, + ThriftType (..), + name, + num_children, + schematype, + unField, + ) +import DataFrame.Internal.Column (Column, fromList) +import DataFrame.Internal.DataFrame (DataFrame, fromNamedColumns) + +{- | 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@. +-} +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) + +{- | Convert a single parquet 'SchemaElement' into a named empty 'Column', +picking a nullable or non-nullable representation based on @nullableCols@. +-} +schemaElemToColumn :: S.Set T.Text -> SchemaElement -> (T.Text, Column) +schemaElemToColumn nullableCols element = + let colName = unField element.name + isNull = colName `S.member` nullableCols + column = + if isNull + then emptyNullableColumnForType (unField element.schematype) + else emptyColumnForType (unField element.schematype) + in (colName, column) + +-- | An empty 'Column' of the given parquet physical type. +emptyColumnForType :: Maybe ThriftType -> Column +emptyColumnForType = \case + Just (BOOLEAN _) -> fromList @Bool [] + Just (INT32 _) -> fromList @Int32 [] + Just (INT64 _) -> fromList @Int64 [] + Just (INT96 _) -> fromList @Int64 [] + Just (FLOAT _) -> fromList @Float [] + Just (DOUBLE _) -> fromList @Double [] + Just (BYTE_ARRAY _) -> fromList @T.Text [] + Just (FIXED_LEN_BYTE_ARRAY _) -> fromList @T.Text [] + other -> error $ "Unsupported parquet type for column: " <> show other + +-- | Like 'emptyColumnForType' but produces a nullable @Maybe a@ column. +emptyNullableColumnForType :: Maybe ThriftType -> Column +emptyNullableColumnForType = \case + Just (BOOLEAN _) -> fromList @(Maybe Bool) [] + Just (INT32 _) -> fromList @(Maybe Int32) [] + Just (INT64 _) -> fromList @(Maybe Int64) [] + Just (INT96 _) -> fromList @(Maybe Int64) [] + Just (FLOAT _) -> fromList @(Maybe Float) [] + Just (DOUBLE _) -> fromList @(Maybe Double) [] + Just (BYTE_ARRAY _) -> fromList @(Maybe T.Text) [] + Just (FIXED_LEN_BYTE_ARRAY _) -> fromList @(Maybe T.Text) [] + other -> error $ "Unsupported parquet type for column: " <> show other diff --git a/src/DataFrame/IO/Parquet/Seeking.hs b/src/DataFrame/IO/Parquet/Seeking.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Seeking.hs @@ -0,0 +1,159 @@ +{- | This module contains low-level utilities around file seeking + +potentially also contains all Streamly related low-level utilities. + +later this module can be renamed / moved to an internal module. +-} +module DataFrame.IO.Parquet.Seeking ( + SeekableHandle (getSeekableHandle), + SeekMode (..), + FileBufferedOrSeekable (..), + ForceNonSeekable, + advanceBytes, + mkFileBufferedOrSeekable, + mkSeekableHandle, + readLastBytes, + seekAndReadBytes, + seekAndStreamBytes, + withFileBufferedOrSeekable, + fSeek, + fGet, +) where + +import Control.Monad +import Control.Monad.IO.Class +import qualified Data.ByteString as BS +import Data.ByteString.Unsafe (unsafeDrop, unsafeTake) +import Data.IORef +import Data.Int +import Data.Word +import Streamly.Data.Stream (Stream) +import qualified Streamly.Data.Stream as S +import qualified Streamly.External.ByteString as SBS +import qualified Streamly.FileSystem.Handle as SHandle +import System.IO + +{- | This handle carries a proof that it must be seekable. +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. +-} +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. +-} +data FileBufferedOrSeekable + = FileBuffered !(IORef Int64) !BS.ByteString + | FileSeekable !SeekableHandle + +-- | Smart constructor for SeekableHandle +mkSeekableHandle :: Handle -> IO (Maybe SeekableHandle) +mkSeekableHandle h = do + seekable <- hIsSeekable h + pure $ if seekable then Just (SeekableHandle h) else Nothing + +-- | For testing only +type ForceNonSeekable = Maybe Bool + +{- | Smart constructor for FileBufferedOrSeekable, tries to keep in the seekable case +if possible. +-} +mkFileBufferedOrSeekable :: + ForceNonSeekable -> Handle -> IO FileBufferedOrSeekable +mkFileBufferedOrSeekable forceNonSeek h = do + seekable <- hIsSeekable h + if not seekable || forceNonSeek == Just True + 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. +-} +withFileBufferedOrSeekable :: + ForceNonSeekable -> + FilePath -> + IOMode -> + (FileBufferedOrSeekable -> IO a) -> + IO a +withFileBufferedOrSeekable forceNonSeek path ioMode action = withFile path ioMode $ \h -> do + fbos <- mkFileBufferedOrSeekable forceNonSeek h + action fbos + +-- | Read from the end, useful for reading metadata without loading entire file +readLastBytes :: Integer -> FileBufferedOrSeekable -> IO BS.ByteString +readLastBytes n (FileSeekable sh) = do + let h = getSeekableHandle sh + hSeek h SeekFromEnd (negate n) + S.fold SBS.write (SHandle.read h) +readLastBytes n (FileBuffered i bs) = do + writeIORef i (fromIntegral $ BS.length bs) + when (n > fromIntegral (BS.length bs)) $ error "lastBytes: n > length bs" + pure $ BS.drop (BS.length bs - fromIntegral n) bs + +-- | Note: this does not guarantee n bytes (if it ends early) +advanceBytes :: Int -> FileBufferedOrSeekable -> IO BS.ByteString +advanceBytes = seekAndReadBytes Nothing + +-- | Note: this does not guarantee n bytes (if it ends early) +seekAndReadBytes :: + Maybe (SeekMode, Integer) -> Int -> FileBufferedOrSeekable -> IO BS.ByteString +seekAndReadBytes mSeek len f = seekAndStreamBytes mSeek len f >>= S.fold SBS.write + +{- | Warning: the stream produced from this function accesses to the mutable handler. +if multiple streams are pulled from the same handler at the same time, chaos happen. +Make sure there is only one stream running at one time for each SeekableHandle, +and streams are not read again when they are not used anymore. +-} +seekAndStreamBytes :: + (MonadIO m) => + Maybe (SeekMode, Integer) -> Int -> FileBufferedOrSeekable -> m (Stream m Word8) +seekAndStreamBytes mSeek len f = do + liftIO $ + case mSeek of + Nothing -> pure () + Just (seekMode, seekTo) -> fSeek f seekMode seekTo + pure $ S.take len $ fRead f + +fSeek :: FileBufferedOrSeekable -> SeekMode -> Integer -> IO () +fSeek (FileSeekable (SeekableHandle h)) seekMode seekTo = hSeek h seekMode seekTo +fSeek (FileBuffered i _bs) AbsoluteSeek seekTo = writeIORef i (fromIntegral seekTo) +fSeek (FileBuffered i _bs) RelativeSeek seekTo = modifyIORef' i (+ fromIntegral seekTo) +fSeek (FileBuffered i bs) SeekFromEnd seekTo = writeIORef i (fromIntegral $ BS.length bs + fromIntegral seekTo) + +fGet :: FileBufferedOrSeekable -> Int -> IO BS.ByteString +fGet (FileSeekable (SeekableHandle h)) n = BS.hGet h n +fGet (FileBuffered iRef bs) n + | n == 0 = pure BS.empty + | n > 0 = do + i <- fromIntegral <$> readIORef iRef + if (BS.length bs - i) < n + then if i <= BS.length bs then pure $ unsafeDrop i bs else pure BS.empty + else pure . unsafeTake n . unsafeDrop i $ bs + | otherwise = error "Can't read a negative number of bytes" + +fRead :: (MonadIO m) => FileBufferedOrSeekable -> Stream m Word8 +fRead (FileSeekable (SeekableHandle h)) = SHandle.read h +fRead (FileBuffered i bs) = S.concatEffect $ do + pos <- liftIO $ readIORef i + pure $ + S.mapM + ( \x -> do + liftIO (modifyIORef' i (+ 1)) + pure x + ) + (S.unfold SBS.reader (BS.drop (fromIntegral pos) bs)) diff --git a/src/DataFrame/IO/Parquet/Thrift.hs b/src/DataFrame/IO/Parquet/Thrift.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Thrift.hs @@ -0,0 +1,584 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE TypeFamilies #-} + +module DataFrame.IO.Parquet.Thrift where + +import Data.ByteString (ByteString) +import Data.Int (Int16, Int32, Int64, Int8) +import Data.Text (Text) +import GHC.Generics (Generic) +import GHC.TypeLits (KnownNat) +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 +data ThriftType + = BOOLEAN (Enumeration 0) + | INT32 (Enumeration 1) + | INT64 (Enumeration 2) + | INT96 (Enumeration 3) + | FLOAT (Enumeration 4) + | DOUBLE (Enumeration 5) + | BYTE_ARRAY (Enumeration 6) + | FIXED_LEN_BYTE_ARRAY (Enumeration 7) + deriving (Eq, Show, Generic) + +instance Pinchable ThriftType + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L183 +data FieldRepetitionType + = REQUIRED (Enumeration 0) + | OPTIONAL (Enumeration 1) + | REPEATED (Enumeration 2) + deriving (Eq, Show, Generic) + +instance Pinchable FieldRepetitionType + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/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 + PLAIN_DICTIONARY (Enumeration 2) + | RLE (Enumeration 3) + | BIT_PACKED (Enumeration 4) + | DELTA_BINARY_PACKED (Enumeration 5) + | DELTA_LENGTH_BYTE_ARRAY (Enumeration 6) + | DELTA_BYTE_ARRAY (Enumeration 7) + | RLE_DICTIONARY (Enumeration 8) + | BYTE_STREAM_SPLIT (Enumeration 9) + deriving (Eq, Show, Generic) + +instance Pinchable Encoding + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L244 +data CompressionCodec + = UNCOMPRESSED (Enumeration 0) + | SNAPPY (Enumeration 1) + | GZIP (Enumeration 2) + | LZO (Enumeration 3) + | BROTLI (Enumeration 4) + | LZ4 (Enumeration 5) + | ZSTD (Enumeration 6) + | LZ4_RAW (Enumeration 7) + deriving (Eq, Show, Generic) + +instance Pinchable CompressionCodec + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L261 +data PageType + = DATA_PAGE (Enumeration 0) + | INDEX_PAGE (Enumeration 1) + | DICTIONARY_PAGE (Enumeration 2) + | DATA_PAGE_V2 (Enumeration 3) + deriving (Eq, Show, Generic) + +instance Pinchable PageType + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L271 +data BoundaryOrder + = UNORDERED (Enumeration 0) + | ASCENDING (Enumeration 1) + | DESCENDING (Enumeration 2) + deriving (Eq, Show, Generic) + +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. + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L283 +-- struct StringType {} +data StringType = StringType deriving (Eq, Show) +instance Pinchable StringType where + type Tag StringType = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure StringType + +data UUIDType = UUIDType deriving (Eq, Show) +instance Pinchable UUIDType where + type Tag UUIDType = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure UUIDType + +data MapType = MapType deriving (Eq, Show) +instance Pinchable MapType where + type Tag MapType = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure MapType + +data ListType = ListType deriving (Eq, Show) +instance Pinchable ListType where + type Tag ListType = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure ListType + +data EnumType = EnumType deriving (Eq, Show) +instance Pinchable EnumType where + type Tag EnumType = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure EnumType + +data DateType = DateType deriving (Eq, Show) +instance Pinchable DateType where + type Tag DateType = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure DateType + +data Float16Type = Float16Type deriving (Eq, Show) +instance Pinchable Float16Type where + type Tag Float16Type = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure Float16Type + +data NullType = NullType deriving (Eq, Show) +instance Pinchable NullType where + type Tag NullType = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure NullType + +data JsonType = JsonType deriving (Eq, Show) +instance Pinchable JsonType where + type Tag JsonType = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure JsonType + +data BsonType = BsonType deriving (Eq, Show) +instance Pinchable BsonType where + type Tag BsonType = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure BsonType + +data VariantType = VariantType deriving (Eq, Show) +instance Pinchable VariantType where + type Tag VariantType = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure VariantType + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L290 +data TimeUnit + = MILLIS (Field 1 MilliSeconds) + | MICROS (Field 2 MicroSeconds) + | NANOS (Field 3 NanoSeconds) + deriving (Eq, Show, Generic) + +instance Pinchable TimeUnit + +data MilliSeconds = MilliSeconds deriving (Eq, Show) +instance Pinchable MilliSeconds where + type Tag MilliSeconds = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure MilliSeconds + +data MicroSeconds = MicroSeconds deriving (Eq, Show) +instance Pinchable MicroSeconds where + type Tag MicroSeconds = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure MicroSeconds + +data NanoSeconds = NanoSeconds deriving (Eq, Show) +instance Pinchable NanoSeconds where + type Tag NanoSeconds = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure NanoSeconds + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L317 +data DecimalType + = DecimalType + { decimal_scale :: Field 1 Int32 + , decimal_precision :: Field 2 Int32 + } + deriving (Eq, Show, Generic) + +instance Pinchable DecimalType + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L328 +data IntType + = IntType + { int_bitWidth :: Field 1 Int8 + , int_isSigned :: Field 2 Bool + } + deriving (Eq, Show, Generic) + +instance Pinchable IntType + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L338 +data TimeType + = TimeType + { time_isAdjustedToUTC :: Field 1 Bool + , time_unit :: Field 2 TimeUnit + } + deriving (Eq, Show, Generic) + +instance Pinchable TimeType + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L349 +data TimestampType + = TimestampType + { timestamp_isAdjustedToUTC :: Field 1 Bool + , timestamp_unit :: Field 2 TimeUnit + } + deriving (Eq, Show, Generic) + +instance Pinchable TimestampType + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L360 +-- union LogicalType +data LogicalType + = LT_STRING (Field 1 StringType) + | LT_MAP (Field 2 MapType) + | LT_LIST (Field 3 ListType) + | LT_ENUM (Field 4 EnumType) + | LT_DECIMAL (Field 5 DecimalType) + | LT_DATE (Field 6 DateType) + | LT_TIME (Field 7 TimeType) + | LT_TIMESTAMP (Field 8 TimestampType) + | LT_INTEGER (Field 10 IntType) + | LT_NULL (Field 11 NullType) + | LT_JSON (Field 12 JsonType) + | LT_BSON (Field 13 BsonType) + | LT_UUID (Field 14 UUIDType) + | LT_FLOAT16 (Field 15 Float16Type) + | LT_VARIANT (Field 16 VariantType) + deriving (Eq, Show, Generic) + +instance Pinchable LogicalType + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L270 +data ConvertedType + = UTF8 (Enumeration 0) + | MAP (Enumeration 1) + | MAP_KEY_VALUE (Enumeration 2) + | LIST (Enumeration 3) + | ENUM (Enumeration 4) + | DECIMAL (Enumeration 5) + | DATE (Enumeration 6) + | TIME_MILLIS (Enumeration 7) + | TIME_MICROS (Enumeration 8) + | TIMESTAMP_MILLIS (Enumeration 9) + | TIMESTAMP_MICROS (Enumeration 10) + | UINT_8 (Enumeration 11) + | UINT_16 (Enumeration 12) + | UINT_32 (Enumeration 13) + | UINT_64 (Enumeration 14) + | INT_8 (Enumeration 15) + | INT_16 (Enumeration 16) + | INT_32 (Enumeration 17) + | INT_64 (Enumeration 18) + | JSON (Enumeration 19) + | BSON (Enumeration 20) + | INTERVAL (Enumeration 21) + deriving (Eq, Show, Generic) + +instance Pinchable ConvertedType + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L505 +data SchemaElement + = SchemaElement + { schematype :: Field 1 (Maybe ThriftType) -- called just type in parquet.thrift + , type_length :: Field 2 (Maybe Int32) + , repetition_type :: Field 3 (Maybe FieldRepetitionType) + , name :: Field 4 Text + , num_children :: Field 5 (Maybe Int32) + , converted_type :: Field 6 (Maybe ConvertedType) + , scale :: Field 7 (Maybe Int32) + , precision :: Field 8 (Maybe Int32) + , field_id :: Field 9 (Maybe Int32) + , logicalType :: Field 10 (Maybe LogicalType) + } + deriving (Eq, Show, Generic) + +instance Pinchable SchemaElement + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L560 +data Statistics + = Statistics + { stats_max :: Field 1 (Maybe ByteString) + , stats_min :: Field 2 (Maybe ByteString) + , stats_null_count :: Field 3 (Maybe Int64) + , stats_distinct_count :: Field 4 (Maybe Int64) + , stats_max_value :: Field 5 (Maybe ByteString) + , stats_min_value :: Field 6 (Maybe ByteString) + , stats_is_max_value_exact :: Field 7 (Maybe Bool) + , stats_is_min_value_exact :: Field 8 (Maybe Bool) + } + deriving (Eq, Show, Generic) + +instance Pinchable Statistics + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L600 +data PageEncodingStats + = PageEncodingStats + { pes_page_type :: Field 1 PageType + , pes_encoding :: Field 2 Encoding + , pes_count :: Field 3 Int32 + } + deriving (Eq, Show, Generic) + +instance Pinchable PageEncodingStats + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L614 +data ColumnMetaData + = ColumnMetaData + { cmd_type :: Field 1 ThriftType + , cmd_encodings :: Field 2 [Encoding] + , cmd_path_in_schema :: Field 3 [Text] + , cmd_codec :: Field 4 CompressionCodec + , cmd_num_values :: Field 5 Int64 + , cmd_total_uncompressed_size :: Field 6 Int64 + , cmd_total_compressed_size :: Field 7 Int64 + , cmd_key_value_metadata :: Field 8 (Maybe [KeyValue]) + , cmd_data_page_offset :: Field 9 Int64 + , cmd_index_page_offset :: Field 10 (Maybe Int64) + , cmd_dictionary_page_offset :: Field 11 (Maybe Int64) + , cmd_statistics :: Field 12 (Maybe Statistics) + , cmd_encoding_stats :: Field 13 (Maybe [PageEncodingStats]) + , cmd_bloom_filter_offset :: Field 14 (Maybe Int64) + , cmd_bloom_filter_length :: Field 15 (Maybe Int32) + } + deriving (Eq, Show, Generic) + +instance Pinchable ColumnMetaData + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/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 +data EncryptionWithColumnKey + = EncryptionWithColumnKey + { ewck_path_in_schema :: Field 1 [Text] + , ewck_key_metadata :: Field 2 (Maybe ByteString) + } + deriving (Eq, Show, Generic) + +instance Pinchable EncryptionWithColumnKey + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L893 +-- union ColumnCryptoMetaData +data ColumnCryptoMetaData + = CCM_ENCRYPTION_WITH_FOOTER_KEY (Field 1 EncryptionWithFooterKey) + | CCM_ENCRYPTION_WITH_COLUMN_KEY (Field 2 EncryptionWithColumnKey) + deriving (Eq, Show, Generic) + +instance Pinchable ColumnCryptoMetaData + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L899 +data ColumnChunk + = ColumnChunk + { cc_file_path :: Field 1 (Maybe Text) + , cc_file_offset :: Field 2 Int64 + , cc_meta_data :: Field 3 (Maybe ColumnMetaData) + , cc_offset_index_offset :: Field 4 (Maybe Int64) + , cc_offset_index_length :: Field 5 (Maybe Int32) + , cc_column_index_offset :: Field 6 (Maybe Int64) + , cc_column_index_length :: Field 7 (Maybe Int32) + , cc_crypto_metadata :: Field 8 (Maybe ColumnCryptoMetaData) + , cc_encrypted_column_metadata :: Field 9 (Maybe ByteString) + } + deriving (Eq, Show, Generic) + +instance Pinchable ColumnChunk + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L940 +data SortingColumn + = SortingColumn + { sc_column_idx :: Field 1 Int32 + , sc_descending :: Field 2 Bool + , sc_nulls_first :: Field 3 Bool + } + deriving (Eq, Show, Generic) + +instance Pinchable SortingColumn + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L958 +data RowGroup + = RowGroup + { rg_columns :: Field 1 [ColumnChunk] + , rg_total_byte_size :: Field 2 Int64 + , rg_num_rows :: Field 3 Int64 + , rg_sorting_columns :: Field 4 (Maybe [SortingColumn]) + , rg_file_offset :: Field 5 (Maybe Int64) + , rg_total_compressed_size :: Field 6 (Maybe Int64) + , rg_ordinal :: Field 7 (Maybe Int16) + } + deriving (Eq, Show, Generic) + +instance Pinchable RowGroup + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L980 +data KeyValue + = KeyValue + { kv_key :: Field 1 Text + , kv_value :: Field 2 (Maybe Text) + } + deriving (Eq, Show, Generic) + +instance Pinchable KeyValue + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L990 +-- union ColumnOrder +newtype ColumnOrder + = TYPE_ORDER (Field 1 TypeDefinedOrder) + deriving (Eq, Show, Generic) + +instance Pinchable ColumnOrder + +-- Empty struct for TYPE_ORDER +data TypeDefinedOrder = TypeDefinedOrder deriving (Eq, Show) +instance Pinchable TypeDefinedOrder where + type Tag TypeDefinedOrder = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure TypeDefinedOrder + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1094 +data AesGcmV1 + = AesGcmV1 + { aes_gcm_v1_aad_prefix :: Field 1 (Maybe ByteString) + , aes_gcm_v1_aad_file_unique :: Field 2 (Maybe ByteString) + , aes_gcm_v1_supply_aad_prefix :: Field 3 (Maybe Bool) + } + deriving (Eq, Show, Generic) + +instance Pinchable AesGcmV1 + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1107 +data AesGcmCtrV1 + = AesGcmCtrV1 + { aes_gcm_ctr_v1_aad_prefix :: Field 1 (Maybe ByteString) + , aes_gcm_ctr_v1_aad_file_unique :: Field 2 (Maybe ByteString) + , aes_gcm_ctr_v1_supply_aad_prefix :: Field 3 (Maybe Bool) + } + deriving (Eq, Show, Generic) + +instance Pinchable AesGcmCtrV1 + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1118 +-- union EncryptionAlgorithm +data EncryptionAlgorithm + = AES_GCM_V1 (Field 1 AesGcmV1) + | AES_GCM_CTR_V1 (Field 2 AesGcmCtrV1) + deriving (Eq, Show, Generic) + +instance Pinchable EncryptionAlgorithm + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1001 +data PageLocation + = PageLocation + { pl_offset :: Field 1 Int64 + , pl_compressed_page_size :: Field 2 Int32 + , pl_first_row_index :: Field 3 Int64 + } + deriving (Eq, Show, Generic) + +instance Pinchable PageLocation + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1017 +data OffsetIndex + = OffsetIndex + { oi_page_locations :: Field 1 [PageLocation] + , oi_unencoded_byte_array_data_bytes :: Field 2 (Maybe [Int64]) + } + deriving (Eq, Show, Generic) + +instance Pinchable OffsetIndex + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1033 +data ColumnIndex + = ColumnIndex + { ci_null_pages :: Field 1 [Bool] + , ci_min_values :: Field 2 [ByteString] + , ci_max_values :: Field 3 [ByteString] + , ci_boundary_order :: Field 4 BoundaryOrder + , ci_null_counts :: Field 5 (Maybe [Int64]) + , ci_repetition_level_histograms :: Field 6 (Maybe [Int64]) + , ci_definition_level_histograms :: Field 7 (Maybe [Int64]) + } + deriving (Eq, Show, Generic) + +instance Pinchable ColumnIndex + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1248 +data DataPageHeader + = DataPageHeader + { dph_num_values :: Field 1 Int32 + , dph_encoding :: Field 2 Encoding + , dph_definition_level_encoding :: Field 3 Encoding + , dph_repetition_level_encoding :: Field 4 Encoding + , dph_statistics :: Field 5 (Maybe Statistics) + } + deriving (Eq, Show, Generic) + +instance Pinchable DataPageHeader + +data IndexPageHeader = IndexPageHeader deriving (Eq, Show) +instance Pinchable IndexPageHeader where + type Tag IndexPageHeader = Pinch.TStruct + pinch _ = Pinch.struct [] + unpinch _ = pure IndexPageHeader + +data DictionaryPageHeader + = DictionaryPageHeader + { diph_num_values :: Field 1 Int32 + , diph_encoding :: Field 2 Encoding + , diph_is_sorted :: Field 3 (Maybe Bool) + } + deriving (Eq, Show, Generic) + +instance Pinchable DictionaryPageHeader + +data DataPageHeaderV2 + = DataPageHeaderV2 + { dph2_num_values :: Field 1 Int32 + , dph2_num_nulls :: Field 2 Int32 + , dph2_num_rows :: Field 3 Int32 + , dph2_encoding :: Field 4 Encoding + , dph2_definition_levels_byte_length :: Field 5 Int32 + , dph2_repetition_levels_byte_length :: Field 6 Int32 + , dph2_is_compressed :: Field 7 (Maybe Bool) + , dph2_statistics :: Field 8 (Maybe Statistics) + } + deriving (Eq, Show, Generic) + +instance Pinchable DataPageHeaderV2 + +data PageHeader + = PageHeader + { ph_type :: Field 1 PageType + , ph_uncompressed_page_size :: Field 2 Int32 + , ph_compressed_page_size :: Field 3 Int32 + , ph_crc :: Field 4 (Maybe Int32) + , ph_data_page_header :: Field 5 (Maybe DataPageHeader) + , ph_index_page_header :: Field 6 (Maybe IndexPageHeader) + , ph_dictionary_page_header :: Field 7 (Maybe DictionaryPageHeader) + , ph_data_page_header_v2 :: Field 8 (Maybe DataPageHeaderV2) + } + deriving (Eq, Show, Generic) + +instance Pinchable PageHeader + +-- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1277 +data FileMetadata + = FileMetadata + { version :: Field 1 Int32 + , schema :: Field 2 [SchemaElement] + , num_rows :: Field 3 Int64 + , row_groups :: Field 4 [RowGroup] + , key_value_metadata :: Field 5 (Maybe [KeyValue]) + , created_by :: Field 6 (Maybe Text) + , column_orders :: Field 7 (Maybe [ColumnOrder]) + , encryption_algorithm :: Field 8 (Maybe EncryptionAlgorithm) + , footer_signing_key_metadata :: Field 9 (Maybe ByteString) + } + deriving (Eq, Show, Generic) + +instance Pinchable FileMetadata + +unField :: (KnownNat n) => Field n a -> a +unField (Pinch.Field a) = a diff --git a/src/DataFrame/IO/Parquet/Time.hs b/src/DataFrame/IO/Parquet/Time.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Time.hs @@ -0,0 +1,67 @@ +{-# LANGUAGE NumericUnderscores #-} + +module DataFrame.IO.Parquet.Time where + +import qualified Data.ByteString as BS +import Data.Time +import Data.Word + +import DataFrame.Internal.Binary ( + littleEndianWord32, + littleEndianWord64, + word32ToLittleEndian, + word64ToLittleEndian, + ) + +int96ToUTCTime :: BS.ByteString -> UTCTime +int96ToUTCTime bytes + | BS.length bytes /= 12 = error "INT96 must be exactly 12 bytes" + | otherwise = + let (nanosBytes, julianBytes) = BS.splitAt 8 bytes + nanosSinceMidnight = littleEndianWord64 nanosBytes + julianDay = littleEndianWord32 julianBytes + in julianDayAndNanosToUTCTime (fromIntegral julianDay) nanosSinceMidnight + +julianDayAndNanosToUTCTime :: Integer -> Word64 -> UTCTime +julianDayAndNanosToUTCTime julianDay nanosSinceMidnight = + let day = julianDayToDay julianDay + secondsSinceMidnight = fromIntegral nanosSinceMidnight / (1_000_000_000 :: Double) + diffTime = secondsToDiffTime (floor secondsSinceMidnight) + in UTCTime day diffTime + +julianDayToDay :: Integer -> Day +julianDayToDay julianDay = + let a = julianDay + 32_044 + b = (4 * a + 3) `div` 146_097 + c = a - (146_097 * b) `div` 4 + d = (4 * c + 3) `div` 1461 + e = c - (1461 * d) `div` 4 + m = (5 * e + 2) `div` 153 + day = e - (153 * m + 2) `div` 5 + 1 + month = m + 3 - 12 * (m `div` 10) + 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. +utcTimeToInt96 :: UTCTime -> BS.ByteString +utcTimeToInt96 (UTCTime day diffTime) = + let julianDay = dayToJulianDay day + nanosSinceMidnight = floor (realToFrac diffTime * (1_000_000_000 :: Double)) + nanosBytes = word64ToLittleEndian nanosSinceMidnight + julianBytes = word32ToLittleEndian (fromIntegral julianDay) + in nanosBytes `BS.append` julianBytes + +dayToJulianDay :: Day -> Integer +dayToJulianDay day = + let (year, month, dayOfMonth) = toGregorian day + a = (fromIntegral $ (14 - fromIntegral month) `div` (12 :: Integer)) :: Integer + y = fromIntegral $ year + 4800 - a + m = fromIntegral $ month + 12 * fromIntegral a - 3 + in fromIntegral dayOfMonth + + (153 * m + 2) `div` 5 + + 365 * y + + y `div` 4 + - y `div` 100 + + y `div` 400 + - 32_045 diff --git a/src/DataFrame/IO/Parquet/Utils.hs b/src/DataFrame/IO/Parquet/Utils.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Utils.hs @@ -0,0 +1,383 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module DataFrame.IO.Parquet.Utils ( + ColumnDescription (..), + generateColumnDescriptions, + getColumnNames, + foldNonNullable, + foldNonNullableUnboxed, + foldNullable, + foldNullableUnboxed, + foldRepeated, + foldRepeatedUnboxed, +) where + +import Control.Monad.IO.Class (MonadIO (..)) +import Data.Int (Int32) +import Data.Maybe (fromMaybe) +import Data.Text (Text) +import qualified Data.Text as T +import qualified Data.Vector as VB +import qualified Data.Vector.Mutable as VBM +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM +import Data.Word (Word8) +import DataFrame.IO.Parquet.Levels ( + stitchList, + stitchList2, + stitchList3, + ) +import DataFrame.IO.Parquet.Thrift ( + ConvertedType (..), + FieldRepetitionType (..), + LogicalType (..), + SchemaElement (..), + ThriftType, + unField, + ) +import DataFrame.IO.Utils.RandomAccess (RandomAccess) +import DataFrame.Internal.Column ( + Column (..), + Columnable, + buildBitmapFromValid, + fromList, + ) +import qualified Streamly.Data.Fold as Fold +import Streamly.Data.Stream (Stream) +import qualified Streamly.Data.Stream as Stream + +data ColumnDescription = ColumnDescription + { colElementType :: !(Maybe ThriftType) + , maxDefinitionLevel :: !Int32 + , maxRepetitionLevel :: !Int32 + , colLogicalType :: !(Maybe LogicalType) + , colConvertedType :: !(Maybe ConvertedType) + , typeLength :: !(Maybe Int32) + } + deriving (Show, Eq) + +levelContribution :: Maybe FieldRepetitionType -> (Int, Int) +levelContribution = \case + Just (REPEATED _) -> (1, 1) + Just (OPTIONAL _) -> (1, 0) + _ -> (0, 0) -- REQUIRED or absent + +data SchemaTree = SchemaTree SchemaElement [SchemaTree] + +buildTree :: [SchemaElement] -> (SchemaTree, [SchemaElement]) +buildTree [] = error "buildTree: schema ended unexpectedly" +buildTree (se : rest) = + let n = fromIntegral $ fromMaybe 0 (unField (num_children se)) :: Int + (children, rest') = buildChildren n rest + in (SchemaTree se children, rest') + +-- | Build a forest of sibling trees from a flat depth-first element list. +buildForest :: [SchemaElement] -> ([SchemaTree], [SchemaElement]) +buildForest [] = ([], []) +buildForest xs = + let (tree, rest') = buildTree xs + (siblings, rest'') = buildForest rest' + in (tree : siblings, rest'') + +-- | Build exactly @n@ child trees, each consuming only its own subtree. +buildChildren :: Int -> [SchemaElement] -> ([SchemaTree], [SchemaElement]) +buildChildren 0 xs = ([], xs) +buildChildren n xs = + let (child, rest') = buildTree xs + (siblings, rest'') = buildChildren (n - 1) rest' + in (child : siblings, rest'') + +collectLeaves :: Int -> Int -> SchemaTree -> [ColumnDescription] +collectLeaves defAcc repAcc (SchemaTree se children) = + let (dInc, rInc) = levelContribution (unField (repetition_type se)) + defLevel = defAcc + dInc + repLevel = repAcc + rInc + in case children of + [] -> + -- leaf: emit a description + let pType = unField (schematype se) + in [ ColumnDescription + pType + (fromIntegral defLevel) + (fromIntegral repLevel) + (unField (logicalType se)) + (unField (converted_type se)) + (unField (type_length se)) + ] + _ -> + -- internal node: recurse into children + concatMap (collectLeaves defLevel repLevel) children + +generateColumnDescriptions :: [SchemaElement] -> [ColumnDescription] +generateColumnDescriptions [] = [] +generateColumnDescriptions (_ : rest) = + -- drop schema root + let (forest, _) = buildForest rest + in concatMap (collectLeaves 0 0) forest + +getColumnNames :: [SchemaElement] -> [Text] +getColumnNames [] = [] +getColumnNames schemaElements = + let (forest, _) = buildForest schemaElements + in go forest [] False + where + isRepeated se = case unField (repetition_type se) of + Just (REPEATED _) -> True + _ -> False + + go [] _ _ = [] + go (SchemaTree se children : rest) path skipThis = + case children of + -- Leaf node + [] -> + let newPath = if skipThis then path else path ++ [unField (name se)] + fullName = T.intercalate "." newPath + in fullName : go rest path skipThis + -- REPEATED intermediate: skip this name; skip single child too + _ + | isRepeated se -> + let skipChildren = length children == 1 + childLeaves = go children path skipChildren + in childLeaves ++ go rest path skipThis + -- Name-skipped intermediate: recurse with skip cleared + _ + | skipThis -> + let childLeaves = go children path False + in childLeaves ++ go rest path skipThis + -- Normal intermediate: add name to path, recurse + _ -> + let subPath = path ++ [unField (name se)] + childLeaves = go children subPath False + in childLeaves ++ go rest path skipThis + +{- | Fold a stream of value chunks into a non-nullable 'Column'. + +Pre-allocates a mutable vector of @totalRows@ and fills it chunk-by-chunk +using a single 'Fold.foldlM\'' pass, avoiding any intermediate list or +concatenation allocation. + +For unboxable element types the chunks (which are always boxed) are +unboxed element-by-element directly into the pre-allocated unboxed +buffer, eliminating the boxing round-trip that a 'fromVector' call on a +boxed concat would otherwise require. +-} +foldNonNullable :: + forall m a. + (RandomAccess m, MonadIO m, Columnable a) => + Int -> + Stream m (VB.Vector a) -> + m Column +foldNonNullable totalRows stream = do + mv <- liftIO $ VBM.unsafeNew totalRows + _ <- + Stream.fold + ( Fold.foldlM' + ( \off chunk -> liftIO $ do + let n = VB.length chunk + VB.copy (VBM.unsafeSlice off n mv) chunk + return (off + n) + ) + (return 0) + ) + stream + v <- liftIO $ VB.unsafeFreeze mv + return (BoxedColumn Nothing v) + +foldNonNullableUnboxed :: + forall m a. + (RandomAccess m, MonadIO m, Columnable a, VU.Unbox a) => + Int -> + Stream m (VU.Vector a) -> + m Column +foldNonNullableUnboxed totalRows stream = do + mv <- liftIO $ VUM.unsafeNew totalRows + _ <- + Stream.fold + ( Fold.foldlM' + ( \off chunk -> liftIO $ do + let n = VU.length chunk + go i + | i >= n = return () + | otherwise = do + VUM.unsafeWrite + mv + (off + i) + (VU.unsafeIndex chunk i) + go (i + 1) + go 0 + return (off + n) + ) + (return 0) + ) + stream + dat <- liftIO $ VU.unsafeFreeze mv + return (UnboxedColumn Nothing dat) + +{- | Fold a stream of (values, def-levels) pairs into a nullable 'Column'. + +Pre-allocates the output buffer and a valid-mask vector of @totalRows@, +then scatters values inline during a single 'Fold.foldlM\'' pass. +This eliminates the @allVals@ intermediate vector that the old +'Stream.toList' + concat approach required. + +A 'hasNull' flag is accumulated during the scatter so the +'buildBitmapFromValid' call (and the second 'VU.all' scan) is skipped +entirely when all values are present. +-} +foldNullable :: + forall m a. + (RandomAccess m, MonadIO m, Columnable a) => + Int -> + Int -> + Stream m (VB.Vector a, VU.Vector Int) -> + m Column +foldNullable maxDef totalRows stream = 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. + mvDat <- + liftIO $ VBM.replicate totalRows (error "parquet: null slot accessed") + mvValid <- liftIO (VUM.new totalRows :: IO (VUM.IOVector Word8)) + (_, hasNull) <- + Stream.fold + ( Fold.foldlM' + ( \(rowOff, anyNull) (vals, defs) -> liftIO $ do + let nDefs = VU.length defs + go i j acc + | i >= nDefs = return acc + | VU.unsafeIndex defs i == maxDef = do + let !v = VB.unsafeIndex vals j + VBM.unsafeWrite mvDat (rowOff + i) v + VUM.unsafeWrite mvValid (rowOff + i) 1 + go (i + 1) (j + 1) acc + | otherwise = go (i + 1) j True + newNull <- go 0 0 False + return (rowOff + nDefs, anyNull || newNull) + ) + (return (0, False)) + ) + stream + dat <- liftIO $ VB.unsafeFreeze mvDat + maybeBm <- + if hasNull + then do + validV <- liftIO $ VU.unsafeFreeze mvValid + return (Just (buildBitmapFromValid validV)) + else return Nothing + return (BoxedColumn maybeBm dat) + +foldNullableUnboxed :: + forall m a. + (RandomAccess m, MonadIO m, Columnable a, VU.Unbox a) => + Int -> + Int -> + Stream m (VU.Vector a, VU.Vector Int) -> + m Column +foldNullableUnboxed maxDef totalRows stream = do + -- zero-init means null slots silently hold 0, guarded by bitmap. + mvDat <- liftIO $ VUM.new totalRows + mvValid <- liftIO (VUM.new totalRows :: IO (VUM.IOVector Word8)) + -- Drain the stream into a list once, then run a tight IO loop. This + -- avoids per-page Streamly polymorphic-monad dispatch in the inner + -- scatter loop. + chunks <- Stream.toList stream + hasNull <- liftIO $ scatterChunks mvDat mvValid maxDef chunks + dat <- liftIO $ VU.unsafeFreeze mvDat + maybeBm <- + if hasNull + then do + validV <- liftIO $ VU.unsafeFreeze mvValid + return (Just (buildBitmapFromValid validV)) + else return Nothing + return (UnboxedColumn maybeBm dat) + where + scatterChunks :: + VUM.IOVector a -> + VUM.IOVector Word8 -> + Int -> + [(VU.Vector a, VU.Vector Int)] -> + IO Bool + scatterChunks mvDat mvValid !md = goChunks 0 False + where + goChunks !_ !anyNull [] = pure anyNull + goChunks !rowOff !anyNull ((vals, defs) : rest) = do + let !nDefs = VU.length defs + go !i !j !acc + | i >= nDefs = pure acc + | VU.unsafeIndex defs i == md = do + VUM.unsafeWrite mvDat (rowOff + i) (VU.unsafeIndex vals j) + VUM.unsafeWrite mvValid (rowOff + i) 1 + go (i + 1) (j + 1) acc + | otherwise = go (i + 1) j True + !newNull <- go 0 0 False + goChunks (rowOff + nDefs) (anyNull || newNull) rest +{-# INLINE foldNullableUnboxed #-} + +{- | Fold a stream of (values, def-levels, rep-levels) triples into a +repeated (list) 'Column' using Dremel-style level stitching. + +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)@. +-} +foldRepeated :: + forall m a. + ( RandomAccess m + , MonadIO m + , Columnable a + , Columnable (Maybe [Maybe a]) + , Columnable (Maybe [Maybe [Maybe a]]) + , Columnable (Maybe [Maybe [Maybe [Maybe a]]]) + ) => + Int -> + Int -> + Stream m (VB.Vector a, VU.Vector Int, VU.Vector Int) -> + m Column +foldRepeated maxRep maxDef stream = do + chunks <- Stream.toList stream + let allVals = VB.concat [vs | (vs, _, _) <- chunks] + allDefs = VU.concat [ds | (_, ds, _) <- chunks] + allReps = VU.concat [rs | (_, _, rs) <- chunks] + return $ case maxRep of + 2 -> fromList (stitchList2 (maxDef - 2) maxDef allReps allDefs allVals) + 3 -> + fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef allReps allDefs allVals) + _ -> fromList (stitchList maxDef allReps allDefs allVals) + +foldRepeatedUnboxed :: + forall m a. + ( RandomAccess m + , MonadIO m + , Columnable a + , VU.Unbox a + , Columnable (Maybe [Maybe a]) + , Columnable (Maybe [Maybe [Maybe a]]) + , Columnable (Maybe [Maybe [Maybe [Maybe a]]]) + ) => + Int -> + Int -> + Stream m (VU.Vector a, VU.Vector Int, VU.Vector Int) -> + m Column +foldRepeatedUnboxed maxRep maxDef stream = do + chunks <- Stream.toList stream + let allVals = VB.convert $ VU.concat [vs | (vs, _, _) <- chunks] + allDefs = VU.concat [ds | (_, ds, _) <- chunks] + allReps = VU.concat [rs | (_, _, rs) <- chunks] + return $ case maxRep of + 2 -> fromList (stitchList2 (maxDef - 2) maxDef allReps allDefs allVals) + 3 -> + fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef allReps allDefs allVals) + _ -> fromList (stitchList maxDef allReps allDefs allVals) diff --git a/src/DataFrame/IO/Utils/RandomAccess.hs b/src/DataFrame/IO/Utils/RandomAccess.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Utils/RandomAccess.hs @@ -0,0 +1,78 @@ +{-# LANGUAGE FlexibleInstances #-} + +module DataFrame.IO.Utils.RandomAccess where + +import Control.Monad.IO.Class (MonadIO (..)) +import Data.ByteString (ByteString) +import Data.ByteString.Internal (ByteString (PS)) +import qualified Data.Vector.Storable as VS +import Data.Word (Word8) +import DataFrame.IO.Parquet.Seeking ( + FileBufferedOrSeekable, + fGet, + fSeek, + readLastBytes, + ) +import Foreign (castForeignPtr) +import System.IO ( + SeekMode (AbsoluteSeek), + ) + +uncurry3 :: (a -> b -> c -> d) -> (a, b, c) -> d +uncurry3 f (a, b, c) = f a b c + +data Range = Range {offset :: !Integer, length :: !Int} deriving (Eq, Show) + +class (Monad m) => RandomAccess m where + readBytes :: Range -> m ByteString + readRanges :: [Range] -> m [ByteString] + readRanges = mapM readBytes + readSuffix :: Int -> m ByteString + +newtype ReaderIO r a = ReaderIO {runReaderIO :: r -> IO a} + +instance Functor (ReaderIO r) where + fmap f (ReaderIO run) = ReaderIO $ fmap f . run + +instance Applicative (ReaderIO r) where + pure a = ReaderIO $ \_ -> pure a + (ReaderIO fg) <*> (ReaderIO fa) = ReaderIO $ \r -> do + a <- fa r + g <- fg r + pure (g a) + +instance Monad (ReaderIO r) where + return = pure + (ReaderIO ma) >>= f = ReaderIO $ \r -> do + a <- ma r + runReaderIO (f a) r + +instance MonadIO (ReaderIO r) where + liftIO io = ReaderIO $ const io + +type LocalFile = ReaderIO FileBufferedOrSeekable + +instance RandomAccess LocalFile where + readBytes (Range offset' length') = ReaderIO $ \handle -> do + fSeek handle AbsoluteSeek offset' + fGet handle length' + readSuffix n = ReaderIO (readLastBytes $ fromIntegral n) + +type MMappedFile = ReaderIO (VS.Vector Word8) + +-- The instance exists but we don't have the means to mmap the file currently +instance RandomAccess MMappedFile where + readBytes (Range offset' length') = + ReaderIO $ + pure . unsafeToByteString . VS.slice (fromInteger offset') length' + readSuffix n = + ReaderIO $ \v -> + let len = VS.length v + n' = min n len + start = len - n' + in pure . unsafeToByteString $ VS.slice start n' v + +unsafeToByteString :: VS.Vector Word8 -> ByteString +unsafeToByteString v = PS (castForeignPtr ptr) offset' len + where + (ptr, offset', len) = VS.unsafeToForeignPtr v