dataframe-parquet (empty) → 1.0.0.0
raw patch · 15 files changed
+3217/−0 lines, 15 filesdep +Globdep +aesondep +attoparsec
Dependencies added: Glob, aeson, attoparsec, base, bytestring, containers, dataframe-core, dataframe-operations, dataframe-parsing, directory, filepath, hashable, http-conduit, pinch, snappy-hs, streamly-bytestring, streamly-core, text, time, unordered-containers, vector, zlib, zstd
Files
- LICENSE +20/−0
- dataframe-parquet.cabal +70/−0
- src/DataFrame/IO/Parquet.hs +740/−0
- src/DataFrame/IO/Parquet/Binary.hs +141/−0
- src/DataFrame/IO/Parquet/Decompress.hs +32/−0
- src/DataFrame/IO/Parquet/Dictionary.hs +164/−0
- src/DataFrame/IO/Parquet/Encoding.hs +132/−0
- src/DataFrame/IO/Parquet/Levels.hs +210/−0
- src/DataFrame/IO/Parquet/Page.hs +352/−0
- src/DataFrame/IO/Parquet/Schema.hs +85/−0
- src/DataFrame/IO/Parquet/Seeking.hs +159/−0
- src/DataFrame/IO/Parquet/Thrift.hs +584/−0
- src/DataFrame/IO/Parquet/Time.hs +67/−0
- src/DataFrame/IO/Parquet/Utils.hs +383/−0
- src/DataFrame/IO/Utils/RandomAccess.hs +78/−0
+ LICENSE view
@@ -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.
+ dataframe-parquet.cabal view
@@ -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
+ src/DataFrame/IO/Parquet.hs view
@@ -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]
+ src/DataFrame/IO/Parquet/Binary.hs view
@@ -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
+ src/DataFrame/IO/Parquet/Decompress.hs view
@@ -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)
+ src/DataFrame/IO/Parquet/Dictionary.hs view
@@ -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)
+ src/DataFrame/IO/Parquet/Encoding.hs view
@@ -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 #-}
+ src/DataFrame/IO/Parquet/Levels.hs view
@@ -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
+ src/DataFrame/IO/Parquet/Page.hs view
@@ -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)
+ src/DataFrame/IO/Parquet/Schema.hs view
@@ -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
+ src/DataFrame/IO/Parquet/Seeking.hs view
@@ -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))
+ src/DataFrame/IO/Parquet/Thrift.hs view
@@ -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
+ src/DataFrame/IO/Parquet/Time.hs view
@@ -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
+ src/DataFrame/IO/Parquet/Utils.hs view
@@ -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)
+ src/DataFrame/IO/Utils/RandomAccess.hs view
@@ -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