dataframe 1.2.0.0 → 1.3.0.0
raw patch · 38 files changed
+4347/−3251 lines, 38 filesdep +pinchdep ~directoryPVP ok
version bump matches the API change (PVP)
Dependencies added: pinch
Dependency ranges changed: directory
API changes (from Hackage documentation)
Files
- CHANGELOG.md +9/−0
- README.md +104/−2
- dataframe.cabal +13/−7
- src/DataFrame.hs +23/−0
- src/DataFrame/Display.hs +13/−9
- src/DataFrame/Display/Terminal/PrettyPrint.hs +50/−24
- src/DataFrame/Functions.hs +41/−176
- src/DataFrame/IO/Parquet.hs +393/−438
- src/DataFrame/IO/Parquet/ColumnStatistics.hs +0/−19
- src/DataFrame/IO/Parquet/Compression.hs +0/−26
- src/DataFrame/IO/Parquet/Decompress.hs +32/−0
- src/DataFrame/IO/Parquet/Dictionary.hs +85/−228
- src/DataFrame/IO/Parquet/Encoding.hs +101/−60
- src/DataFrame/IO/Parquet/Levels.hs +159/−133
- src/DataFrame/IO/Parquet/Page.hs +316/−433
- src/DataFrame/IO/Parquet/Schema.hs +86/−0
- src/DataFrame/IO/Parquet/Seeking.hs +14/−0
- src/DataFrame/IO/Parquet/Thrift.hs +584/−1199
- src/DataFrame/IO/Parquet/Types.hs +0/−314
- src/DataFrame/IO/Parquet/Utils.hs +384/−0
- src/DataFrame/IO/Utils/RandomAccess.hs +78/−0
- src/DataFrame/Internal/DataFrame.hs +130/−40
- src/DataFrame/Internal/Schema.hs +132/−0
- src/DataFrame/Operations/Permutation.hs +38/−10
- src/DataFrame/TH.hs +213/−0
- src/DataFrame/Typed.hs +39/−12
- src/DataFrame/Typed/Aggregate.hs +51/−30
- src/DataFrame/Typed/Expr.hs +0/−13
- src/DataFrame/Typed/Generic.hs +205/−0
- src/DataFrame/Typed/Record.hs +105/−0
- src/DataFrame/Typed/TH.hs +221/−0
- src/DataFrame/Typed/Types.hs +7/−10
- tests/DecisionTree.hs +0/−1
- tests/Main.hs +2/−0
- tests/Operations/Aggregations.hs +3/−3
- tests/Operations/Record.hs +358/−0
- tests/Operations/Sort.hs +72/−0
- tests/Parquet.hs +286/−64
CHANGELOG.md view
@@ -1,5 +1,14 @@ # Revision history for dataframe +## 0.1.3.0.0+### New features+* New `DataFrame.Typed.TH.deriveSchemaFromType` Template Haskell splice generates a typed schema synonym and a `HasSchema` instance from a Haskell record ADT. Pair with `DataFrame.fromRecords` / `DataFrame.toRecords` (or `DataFrame.Typed.fromRecordsTyped` / `toRecordsTyped`) to convert between `[Order]` and `DataFrame`/`TypedDataFrame OrderSchema`. Field names are translated `camelCase → snake_case` by default; the transform is configurable via `SchemaOptions`.+* New `DataFrame.Typed.Generic` exposes `SchemaOf`/`SchemaOfRaw` plus `genericToColumns` / `genericFromColumns`, so users who prefer `GHC.Generics` over a TH splice can derive the same schema and row bridge.+* New `DataFrame.Internal.Schema.deriveSchema` Template Haskell splice generates, from a record ADT, both a runtime `Schema` value (`orderSchema :: Schema`, suitable for `readCsvWithSchema` / `readCsvWithOpts`) and one `Expr` accessor per field (`orderCustomerId :: Expr Int`, etc.), so expression-DSL code can refer to columns by typed name without writing `col @T "snake_case_name"` at each call site. Re-exported from `DataFrame`.++### Refactor+* The untyped Template Haskell splices (`declareColumns`, `declareColumnsFromCsvFile`, `declareColumnsFromCsvWithOpts`, `declareColumnsFromParquetFile`, `declareColumnsWithPrefix`, `declareColumnsWithPrefix'`) have moved from `DataFrame.Functions` to a new `DataFrame.TH` module (re-exported from `DataFrame`). Update imports accordingly; the bundled `dataframe.ghci` already points to the new module.+ ## 1.2.0.0 ### Breaking changes * Remove `Eq` and `Ord` instances for `Expr` — they were buggy under cross-type comparison. Use `eqExpr` / `compareExpr` from `DataFrame.Internal.Expression` directly.
README.md view
@@ -197,7 +197,8 @@ ### Generate column references from a CSV -`declareColumnsFromCsvFile` reads your CSV at compile time and generates typed `Expr` bindings for every column:+`declareColumnsFromCsvFile` (in `DataFrame.TH`, also re-exported from `DataFrame`)+reads your CSV at compile time and generates typed `Expr` bindings for every column: ```haskell {-# LANGUAGE TemplateHaskell #-}@@ -212,7 +213,7 @@ -- total_rooms :: Expr Double -- ocean_proximity :: Expr Text -- ... one binding per column-$(F.declareColumnsFromCsvFile "./data/housing.csv")+$(D.declareColumnsFromCsvFile "./data/housing.csv") main :: IO () main = do@@ -250,6 +251,107 @@ -- , ... -- ] $(T.deriveSchemaFromCsvFile "HousingSchema" "./data/housing.csv")+```++### Generate a schema (and a row bridge) from a record ADT++When the canonical row shape lives in your code as a Haskell record,+`deriveSchemaFromType` produces both the typed schema and a `HasSchema`+instance that converts between `[Order]` and a `DataFrame` (or+`TypedDataFrame OrderSchema`) at runtime:++```haskell+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE TypeFamilies #-}++import Data.Int (Int64)+import qualified Data.Text as T+import qualified DataFrame as D+import qualified DataFrame.Typed as DT++data Order = Order+ { orderId :: Int64+ , region :: T.Text+ , amount :: Double+ } deriving (Show, Eq)++$(DT.deriveSchemaFromType ''Order)+-- expands to:+-- type OrderSchema =+-- '[DT.Column "order_id" Int64, DT.Column "region" T.Text, DT.Column "amount" Double]+-- instance DT.HasSchema Order where+-- type Schema Order = OrderSchema+-- toColumns = ...+-- fromColumns = ...++xs :: [Order]+xs = [Order 1 "us" 10.0, Order 2 "eu" 20.5]++-- Untyped: [Order] <-> DataFrame+df :: D.DataFrame+df = D.fromRecords xs++xs' :: Either T.Text [Order]+xs' = D.toRecords df -- runtime-checked++-- Typed: [Order] <-> TypedDataFrame OrderSchema+tdf :: DT.TypedDataFrame OrderSchema+tdf = DT.fromRecordsTyped xs+```++Field names are translated `camelCase → snake_case` by default; override+the translation with `deriveSchemaFromTypeWith+defaultSchemaOptions{nameTransform = id}` (or any `String -> String`).++If all you need is a runtime `Schema` to drive `readCsvWithSchema` (no+typed-dataframe machinery), there's a companion splice in+`DataFrame.Internal.Schema` (re-exported from `DataFrame`):++```haskell+$(D.deriveSchema ''Order)+-- emits:+-- orderSchema :: Schema+-- orderSchema = makeSchema [("order_id", schemaType @Int64), ...]+-- orderOrderId :: Expr Int64+-- orderOrderId = col "order_id"+-- orderRegion :: Expr Text+-- orderRegion = col "region"+-- orderAmount :: Expr Double+-- orderAmount = col "amount"++orders :: IO D.DataFrame+orders = do+ df <- D.readCsvWithSchema orderSchema "orders.csv"+ pure (D.filter orderAmount (> 100) df)+```++Each record field gets a typed accessor named `<lower-first TyConName><UpperFirst FieldName>`,+so `data Order { customerId :: Int }` yields `orderCustomerId :: Expr Int = col "customer_id"`.+That's the same shape as `$(D.declareColumns df)` produces from a runtime+`DataFrame`, but driven off the ADT instead of an existing frame.++If you'd rather not depend on Template Haskell, the same schema is+available via `GHC.Generics`:++```haskell+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}++import GHC.Generics (Generic)+import DataFrame.Typed (Schema)+import qualified DataFrame.Typed as DT++data Order = Order { … } deriving (Generic)++type OrderSchema = DT.SchemaOf Order++instance DT.HasSchema Order where+ type Schema Order = OrderSchema+ toColumns = DT.genericToColumns+ fromColumns = DT.genericFromColumns ``` ## Typed API
dataframe.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.0 name: dataframe-version: 1.2.0.0+version: 1.3.0.0 synopsis: A fast, safe, and intuitive DataFrame library. @@ -78,20 +78,21 @@ DataFrame.Display.Terminal.Plot, DataFrame.IO.CSV, DataFrame.IO.JSON,+ DataFrame.IO.Utils.RandomAccess, DataFrame.IO.Parquet, DataFrame.IO.Parquet.Binary, DataFrame.IO.Parquet.Dictionary, DataFrame.IO.Parquet.Levels, DataFrame.IO.Parquet.Thrift,- DataFrame.IO.Parquet.ColumnStatistics,- DataFrame.IO.Parquet.Compression,+ DataFrame.IO.Parquet.Decompress, DataFrame.IO.Parquet.Encoding, DataFrame.IO.Parquet.Page,+ DataFrame.IO.Parquet.Utils, DataFrame.IO.Parquet.Seeking, DataFrame.IO.Parquet.Time,- DataFrame.IO.Parquet.Types, DataFrame.Lazy.IO.CSV, DataFrame.Lazy.IO.Binary,+ DataFrame.IO.Parquet.Schema, DataFrame.Lazy.Internal.DataFrame, DataFrame.Lazy.Internal.LogicalPlan, DataFrame.Lazy.Internal.PhysicalPlan,@@ -109,7 +110,10 @@ DataFrame.Typed.TH, DataFrame.Typed.Expr, DataFrame.Typed.Lazy,- DataFrame.Typed+ DataFrame.Typed.Record,+ DataFrame.Typed.Generic,+ DataFrame.Typed,+ DataFrame.TH build-depends: base >= 4 && <5, async >= 2.2 && < 3, deepseq >= 1 && < 2,@@ -141,9 +145,11 @@ stm >= 2.5 && < 3, filepath >= 1.4 && < 2, Glob >= 0.10 && < 1,- http-conduit >= 2.3 && < 3,+ http-conduit >= 2.3 && < 3,+ pinch >= 0.5 && < 1, streamly-core >= 0.2.3 && < 0.4, streamly-bytestring >= 0.2.0 && < 0.4+ hs-source-dirs: src default-language: Haskell2010 @@ -264,6 +270,7 @@ Operations.Statistics, Operations.Take, Operations.Typing,+ Operations.Record, LazyParquet, Parquet, ParquetTestData,@@ -272,7 +279,6 @@ build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.13, dataframe >= 1 && < 2,- directory >= 1.3 && < 2, HUnit ^>= 1.6, QuickCheck >= 2 && < 3, random-shuffle >= 0.0.4 && < 1,
src/DataFrame.hs view
@@ -233,6 +233,12 @@ -- * Errors module Errors, + -- * Record bridge+ module Record,++ -- * Template Haskell column-binding splices+ module TH,+ -- * Plotting module Plot, )@@ -254,6 +260,7 @@ fromCsvBytes, readCsv, readCsvWithOpts,+ readCsvWithSchema, readSeparated, readTsv, writeCsv,@@ -282,6 +289,8 @@ import DataFrame.Internal.DataFrame as Dataframe ( DataFrame, GroupedDataFrame,+ TruncateConfig (..),+ defaultTruncateConfig, empty, null, toCsv,@@ -301,6 +310,7 @@ toRowVector, ) import DataFrame.Internal.Schema as Schema (+ deriveSchema, makeSchema, schemaType, )@@ -395,3 +405,16 @@ parseDefaults, ) import DataFrame.Operators as Operators+import DataFrame.TH as TH (+ declareColumns,+ declareColumnsFromCsvFile,+ declareColumnsFromCsvWithOpts,+ declareColumnsFromParquetFile,+ declareColumnsWithPrefix,+ declareColumnsWithPrefix',+ )+import DataFrame.Typed.Record as Record (+ HasSchema (..),+ fromRecords,+ toRecords,+ )
src/DataFrame/Display.hs view
@@ -4,23 +4,27 @@ module DataFrame.Display where import qualified Data.Text.IO as T+import DataFrame.Display.Terminal.PrettyPrint (RenderFormat (..)) import qualified DataFrame.Internal.DataFrame as D-import qualified DataFrame.Operations.Subset as D -import Data.Function- data DisplayOptions = DisplayOptions- { {-- | Maximum number of rows to render.+ { {-- | Controls truncation on all three axes (rows, columns, cell width). --- -- * If this value is less than or equal to 0, no rows are printed.- -- * If it is greater than the number of rows in the frame, all rows are printed.+ -- * 'Nothing' renders every row and column at full width.+ -- * 'Just cfg' caps each axis whose limit in @cfg@ is positive. --}- displayRows :: Int+ displayTruncate :: Maybe D.TruncateConfig } +{- | Defaults aimed at terminal use: 10 rows, plus the standard column and cell+caps from 'D.defaultTruncateConfig'.+-} defaultDisplayOptions :: DisplayOptions-defaultDisplayOptions = DisplayOptions 10+defaultDisplayOptions =+ DisplayOptions+ { displayTruncate = Just D.defaultTruncateConfig{D.maxRows = 10}+ } -- | Render a 'DataFrame' to stdout according to 'DisplayOptions'. display :: DisplayOptions -> D.DataFrame -> IO ()-display opts df = df & D.take (displayRows opts) & (`D.asText` False) & T.putStrLn+display opts = T.putStrLn . D.asTextWith Plain (displayTruncate opts)
src/DataFrame/Display/Terminal/PrettyPrint.hs view
@@ -3,8 +3,14 @@ module DataFrame.Display.Terminal.PrettyPrint where import qualified Data.Text as T+import qualified Data.Vector as V -import Data.List (transpose)+{- | Output format for 'showTable'. 'Plain' renders a terminal-style table with+ASCII borders; 'Markdown' renders a GitHub-flavoured pipe table suitable for+notebooks.+-}+data RenderFormat = Plain | Markdown+ deriving (Show, Eq) -- Utility functions to show a DataFrame as a Markdown-ish table. @@ -22,16 +28,14 @@ -- functions that fill a string (s) to a given width (n) by adding pad -- character (c) to align left, right, or center fillLeft :: Char -> Int -> T.Text -> T.Text-fillLeft c n s = s `T.append` T.replicate (n - T.length s) (T.singleton c)+fillLeft c n s = s <> T.replicate (n - T.length s) (T.singleton c) fillRight :: Char -> Int -> T.Text -> T.Text-fillRight c n s = T.replicate (n - T.length s) (T.singleton c) `T.append` s+fillRight c n s = T.replicate (n - T.length s) (T.singleton c) <> s fillCenter :: Char -> Int -> T.Text -> T.Text fillCenter c n s =- T.replicate l (T.singleton c)- `T.append` s- `T.append` T.replicate r (T.singleton c)+ T.replicate l (T.singleton c) <> s <> T.replicate r (T.singleton c) where x = n - T.length s l = x `div` 2@@ -47,32 +51,54 @@ center :: Int -> T.Text -> T.Text center = fillCenter ' ' -showTable :: Bool -> [T.Text] -> [T.Text] -> [[T.Text]] -> T.Text-showTable properMarkdown header types rows =- let consolidatedHeader =- if properMarkdown+{- | Render a table from column-major data. @columns@ has one 'V.Vector' per+column; widths are computed in one pass per column (no row-major transpose),+and row lines are built by indexing each column at row @i@.+-}+showTable ::+ RenderFormat ->+ [T.Text] ->+ [T.Text] ->+ [V.Vector T.Text] ->+ T.Text+showTable fmt header types columns =+ let isMarkdown = fmt == Markdown+ consolidatedHeader =+ if isMarkdown then zipWith (\h t -> h <> "<br>" <> t) header types else header cs = map (\h -> ColDesc center h left) consolidatedHeader+ nRows = case columns of+ (c : _) -> V.length c+ [] -> 0+ columnMaxWidth col+ | V.null col = 0+ | otherwise = V.foldl' (\acc x -> max acc (T.length x)) 0 col widths =- [ maximum $ map T.length col- | col <- transpose $ consolidatedHeader : types : rows- ]- border = T.intercalate "---" [T.replicate width (T.singleton '-') | width <- widths]- separator = T.intercalate "-|-" [T.replicate width (T.singleton '-') | width <- widths]- fillCols fill cols =- T.intercalate " | " [fill c width col | (c, width, col) <- zip3 cs widths cols]+ zipWith3+ (\h t col -> T.length h `max` T.length t `max` columnMaxWidth col)+ consolidatedHeader+ types+ columns+ dashesOf w = T.replicate w "-"+ border = T.intercalate "---" (map dashesOf widths)+ separator = T.intercalate "-|-" (map dashesOf widths)+ fillCells fill cells =+ T.intercalate " | " (zipWith3 fill cs widths cells)+ rowCells i = map (V.! i) columns+ rowLines = [fillCells colValueFill (rowCells i) | i <- [0 .. nRows - 1]]+ wrapMd t = T.concat ["| ", t, " |"] outputLines =- if properMarkdown+ if isMarkdown then- T.concat ["| ", fillCols colTitleFill consolidatedHeader, " |"]- : T.concat ["| ", separator, " |"]- : map ((\t -> T.concat ["| ", t, " |"]) . fillCols colValueFill) rows+ wrapMd (fillCells colTitleFill consolidatedHeader)+ : wrapMd separator+ : map wrapMd rowLines else border- : fillCols colTitleFill consolidatedHeader+ : fillCells colTitleFill consolidatedHeader : separator- : fillCols colTitleFill types+ : fillCells colTitleFill types : separator- : map (fillCols colValueFill) rows+ : rowLines in T.unlines outputLines
src/DataFrame/Functions.hs view
@@ -1,15 +1,11 @@-{-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE IncoherentInstances #-}-{-# LANGUAGE InstanceSigs #-}-{-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-}@@ -17,36 +13,22 @@ module DataFrame.Functions (module DataFrame.Functions, module DataFrame.Operators) where import DataFrame.Internal.Column-import DataFrame.Internal.DataFrame (- DataFrame (..),- empty,- unsafeGetColumn,- )-import DataFrame.Internal.Expression hiding (normalize)+import DataFrame.Internal.Expression import DataFrame.Internal.Statistics-import DataFrame.Operations.Core import Control.Applicative-import Control.Monad-import Control.Monad.IO.Class import qualified Data.Char as Char import Data.Either-import Data.Function-import Data.Functor+import Data.Function (on) import Data.Int import qualified Data.List as L import qualified Data.Map as M import qualified Data.Maybe as Maybe-import qualified Data.Set as S import qualified Data.Text as T import Data.Time import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU-import Data.Word-import qualified DataFrame.IO.CSV as CSV-import qualified DataFrame.IO.Parquet as Parquet-import DataFrame.IO.Parquet.Thrift-import DataFrame.IO.Parquet.Types (columnNullCount)+ import DataFrame.Internal.Nullable ( BaseType, NullLift1Op (applyNull1),@@ -55,12 +37,6 @@ NullLift2Result, ) import DataFrame.Operators-import Debug.Trace (trace)-import Language.Haskell.TH-import qualified Language.Haskell.TH.Syntax as TH-import System.Directory (doesDirectoryExist)-import System.FilePath ((</>))-import System.FilePath.Glob (glob) import Text.Regex.TDFA import Prelude hiding (maximum, minimum) import Prelude as P@@ -71,7 +47,10 @@ lift2 :: (Columnable c, Columnable b, Columnable a) =>- (c -> b -> a) -> Expr c -> Expr b -> Expr a+ (c -> b -> a) ->+ Expr c ->+ Expr b ->+ Expr a lift2 f = Binary ( MkBinaryOp@@ -161,7 +140,9 @@ castExpr :: forall b src.- (Columnable b, Columnable src, Read b) => Expr src -> Expr (Maybe b)+ (Columnable b, Columnable src, Read b) =>+ Expr src ->+ Expr (Maybe b) castExpr = CastExprWith @b @(Maybe b) @src "castExpr" (either (const Nothing) Just) castExprWithDefault ::@@ -173,7 +154,9 @@ castExprEither :: forall b src.- (Columnable b, Columnable src, Read b) => Expr src -> Expr (Either T.Text b)+ (Columnable b, Columnable src, Read b) =>+ Expr src ->+ Expr (Either T.Text b) castExprEither = CastExprWith @b @(Either T.Text b) @src "castExprEither"@@ -459,7 +442,11 @@ reduce :: forall a b.- (Columnable a, Columnable b) => Expr b -> a -> (a -> b -> a) -> Expr a+ (Columnable a, Columnable b) =>+ Expr b ->+ a ->+ (a -> b -> a) ->+ Expr a reduce expr start f = Agg (FoldAgg "foldUdf" (Just start) f) expr {-# INLINEABLE reduce #-} @@ -497,21 +484,29 @@ whenPresent :: forall a b.- (Columnable a, Columnable b) => (a -> b) -> Expr (Maybe a) -> Expr (Maybe b)+ (Columnable a, Columnable b) =>+ (a -> b) ->+ Expr (Maybe a) ->+ Expr (Maybe b) whenPresent f = liftDecorated (fmap f) "whenPresent" Nothing {-# INLINEABLE whenPresent #-} whenBothPresent :: forall a b c. (Columnable a, Columnable b, Columnable c) =>- (a -> b -> c) -> Expr (Maybe a) -> Expr (Maybe b) -> Expr (Maybe c)+ (a -> b -> c) ->+ Expr (Maybe a) ->+ Expr (Maybe b) ->+ Expr (Maybe c) whenBothPresent f = lift2Decorated (\l r -> f <$> l <*> r) "whenBothPresent" Nothing False 0 {-# INLINEABLE whenBothPresent #-} recode :: forall a b. (Columnable a, Columnable b, Show (a, b)) =>- [(a, b)] -> Expr a -> Expr (Maybe b)+ [(a, b)] ->+ Expr a ->+ Expr (Maybe b) recode mapping = Unary ( MkUnaryOp@@ -524,13 +519,20 @@ recodeWithCondition :: forall a b. (Columnable a, Columnable b) =>- Expr b -> [(Expr a -> Expr Bool, b)] -> Expr a -> Expr b+ Expr b ->+ [(Expr a -> Expr Bool, b)] ->+ Expr a ->+ Expr b recodeWithCondition fallback [] _val = fallback recodeWithCondition fallback ((cond, val) : rest) expr = ifThenElse (cond expr) (lit val) (recodeWithCondition fallback rest expr) recodeWithDefault :: forall a b.- (Columnable a, Columnable b, Show (a, b)) => b -> [(a, b)] -> Expr a -> Expr b+ (Columnable a, Columnable b, Show (a, b)) =>+ b ->+ [(a, b)] ->+ Expr a ->+ Expr b recodeWithDefault d mapping = Unary ( MkUnaryOp@@ -584,7 +586,9 @@ bind :: forall a b m. (Columnable a, Columnable (m a), Monad m, Columnable b, Columnable (m b)) =>- (a -> m b) -> Expr (m a) -> Expr (m b)+ (a -> m b) ->+ Expr (m a) ->+ Expr (m b) bind f = liftDecorated (>>= f) "bind" Nothing {- | Window function: evaluate an expression partitioned by the given columns.@@ -670,142 +674,3 @@ '[' -> True ']' -> True _ -> False--typeFromString :: [String] -> Q Type-typeFromString [] = fail "No type specified"-typeFromString [t0] = do- let t = normalize t0- case stripBrackets t of- Just inner -> typeFromString [inner] <&> AppT ListT- Nothing- | t == "Text" || t == "Data.Text.Text" || t == "T.Text" ->- pure (ConT ''T.Text)- | otherwise -> do- m <- lookupTypeName t- case m of- Just tyName -> pure (ConT tyName)- Nothing -> fail $ "Unsupported type: " ++ t0-typeFromString [tycon, t1] = AppT <$> typeFromString [tycon] <*> typeFromString [t1]-typeFromString [tycon, t1, t2] =- (\outer a b -> AppT (AppT outer a) b)- <$> typeFromString [tycon]- <*> typeFromString [t1]- <*> typeFromString [t2]-typeFromString s = fail $ "Unsupported types: " ++ unwords s--normalize :: String -> String-normalize = dropWhile (== ' ') . reverse . dropWhile (== ' ') . reverse--stripBrackets :: String -> Maybe String-stripBrackets s =- case s of- ('[' : rest)- | P.not (null rest) && last rest == ']' ->- Just (init rest)- _ -> Nothing--declareColumnsFromCsvFile :: String -> DecsQ-declareColumnsFromCsvFile path = do- df <-- liftIO- (CSV.readSeparated (CSV.defaultReadOptions{CSV.numColumns = Just 100}) path)- declareColumns df--declareColumnsFromParquetFile :: String -> DecsQ-declareColumnsFromParquetFile path = do- isDir <- liftIO $ doesDirectoryExist path- let pat = if isDir then path </> "*.parquet" else path- matches <- liftIO $ glob pat- files <- liftIO $ filterM (fmap Prelude.not . doesDirectoryExist) matches- metas <- liftIO $ mapM (fmap fst . Parquet.readMetadataFromPath) files- let nullableCols :: S.Set T.Text- nullableCols =- S.fromList- [ T.pack (last colPath)- | meta <- metas- , rg <- rowGroups meta- , cc <- rowGroupColumns rg- , let cm = columnMetaData cc- colPath = columnPathInSchema cm- , Prelude.not (null colPath)- , columnNullCount (columnStatistics cm) > 0- ]- let df =- foldl- (\acc meta -> acc <> schemaToEmptyDataFrame nullableCols (schema meta))- DataFrame.Internal.DataFrame.empty- metas- declareColumns df--schemaToEmptyDataFrame :: S.Set T.Text -> [SchemaElement] -> DataFrame-schemaToEmptyDataFrame nullableCols elems =- let leafElems = filter (\e -> numChildren e == 0) elems- in fromNamedColumns (map (schemaElemToColumn nullableCols) leafElems)--schemaElemToColumn :: S.Set T.Text -> SchemaElement -> (T.Text, Column)-schemaElemToColumn nullableCols element =- let colName = elementName element- isNull = colName `S.member` nullableCols- column =- if isNull- then emptyNullableColumnForType (elementType element)- else emptyColumnForType (elementType element)- in (colName, column)--emptyColumnForType :: TType -> Column-emptyColumnForType = \case- BOOL -> fromList @Bool []- BYTE -> fromList @Word8 []- I16 -> fromList @Int16 []- I32 -> fromList @Int32 []- I64 -> fromList @Int64 []- I96 -> fromList @Int64 []- FLOAT -> fromList @Float []- DOUBLE -> fromList @Double []- STRING -> fromList @T.Text []- other -> error $ "Unsupported parquet type for column: " <> show other--emptyNullableColumnForType :: TType -> Column-emptyNullableColumnForType = \case- BOOL -> fromList @(Maybe Bool) []- BYTE -> fromList @(Maybe Word8) []- I16 -> fromList @(Maybe Int16) []- I32 -> fromList @(Maybe Int32) []- I64 -> fromList @(Maybe Int64) []- I96 -> fromList @(Maybe Int64) []- FLOAT -> fromList @(Maybe Float) []- DOUBLE -> fromList @(Maybe Double) []- STRING -> fromList @(Maybe T.Text) []- other -> error $ "Unsupported parquet type for column: " <> show other--declareColumnsFromCsvWithOpts :: CSV.ReadOptions -> String -> DecsQ-declareColumnsFromCsvWithOpts opts path = do- df <- liftIO (CSV.readSeparated opts path)- declareColumns df--declareColumns :: DataFrame -> DecsQ-declareColumns = declareColumnsWithPrefix' Nothing--declareColumnsWithPrefix :: T.Text -> DataFrame -> DecsQ-declareColumnsWithPrefix prefix = declareColumnsWithPrefix' (Just prefix)--declareColumnsWithPrefix' :: Maybe T.Text -> DataFrame -> DecsQ-declareColumnsWithPrefix' prefix df =- let- names = (map fst . L.sortBy (compare `on` snd) . M.toList . columnIndices) df- types = map (columnTypeString . (`unsafeGetColumn` df)) names- specs =- zipWith- ( \colName type_ -> (colName, maybe "" (sanitize . (<> "_")) prefix <> sanitize colName, type_)- )- names- types- in- fmap concat $ forM specs $ \(raw, nm, tyStr) -> do- ty <- typeFromString (words tyStr)- let tyDisplay = if ' ' `elem` tyStr then "(" <> T.pack tyStr <> ")" else T.pack tyStr- trace (T.unpack (nm <> " :: Expr " <> tyDisplay)) pure ()- let n = mkName (T.unpack nm)- sig <- sigD n [t|Expr $(pure ty)|]- val <- valD (varP n) (normalB [|col $(TH.lift raw)|]) []- pure [sig, val]
src/DataFrame/IO/Parquet.hs view
@@ -1,6 +1,8 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE NumericUnderscores #-}+{-# LANGUAGE OverloadedRecordDot #-} {-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} @@ -8,34 +10,73 @@ import Control.Exception (throw, try) import Control.Monad-import qualified Data.ByteString as BSO-import Data.Either-import Data.IORef-import Data.Int+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)+import Data.List (foldl', transpose) import qualified Data.List as L-import qualified Data.Map.Strict as M-import qualified Data.Set as S+import qualified Data.Map as Map import qualified Data.Text as T-import Data.Text.Encoding-import Data.Time+import Data.Text.Encoding (encodeUtf8)+import Data.Time (UTCTime) import Data.Time.Clock.POSIX (posixSecondsToUTCTime)-import qualified Data.Vector as V+import qualified Data.Vector as Vector+import qualified Data.Vector.Unboxed as VU import DataFrame.Errors (DataFrameException (ColumnsNotFoundException))-import DataFrame.Internal.Binary (littleEndianWord32)+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, columns)+import DataFrame.Internal.DataFrame (DataFrame (..)) import DataFrame.Internal.Expression (Expr, getColumns)-import qualified DataFrame.Operations.Core as DI import DataFrame.Operations.Merge () import qualified DataFrame.Operations.Subset as DS-import System.FilePath.Glob (compile, glob, match)--import Data.Aeson (FromJSON (..), eitherDecodeStrict, withObject, (.:))-import DataFrame.IO.Parquet.Dictionary-import DataFrame.IO.Parquet.Levels-import DataFrame.IO.Parquet.Page-import DataFrame.IO.Parquet.Thrift-import DataFrame.IO.Parquet.Types import Network.HTTP.Simple ( getResponseBody, getResponseStatusCode,@@ -43,16 +84,16 @@ parseRequest, setRequestHeader, )+import qualified Pinch+import qualified Streamly.Data.Stream as Stream import System.Directory ( doesDirectoryExist, getHomeDirectory, getTemporaryDirectory, ) import System.Environment (lookupEnv)--import qualified Data.Vector.Unboxed as VU-import DataFrame.IO.Parquet.Seeking import System.FilePath ((</>))+import System.FilePath.Glob (compile, glob, match) import System.IO (IOMode (ReadMode)) -- Options -----------------------------------------------------------------@@ -128,28 +169,6 @@ 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. -}--{- | Strip Parquet encoding artifact names (REPEATED wrappers and their single- list-element children) from a raw column path, leaving user-visible names.--}-cleanColPath :: [SNode] -> [String] -> [String]-cleanColPath nodes path = go nodes path False- where- go _ [] _ = []- go ns (p : ps) skipThis =- case L.find (\n -> sName n == p) ns of- Nothing -> []- Just n- | sRep n == REPEATED && not (null (sChildren n)) ->- let skipChildren = length (sChildren n) == 1- in go (sChildren n) ps skipChildren- | skipThis ->- go (sChildren n) ps False- | null (sChildren n) ->- [p]- | otherwise ->- p : go (sChildren n) ps False- readParquetWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame readParquetWithOpts opts path | isHFUri path = do@@ -159,131 +178,12 @@ pure (applyRowRange opts (mconcat dfs)) | otherwise = _readParquetWithOpts Nothing opts path --- | Internal function to pass testing parameters+-- | 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 -> do- fileMetadata <- readMetadataFromHandle file- let columnPaths = getColumnPaths (drop 1 $ schema fileMetadata)- let columnNames = map fst columnPaths- let leafNames = map (last . T.splitOn ".") columnNames- let availableSelectedColumns = L.nub leafNames- let predicateColumns = maybe [] (L.nub . getColumns) (predicate opts)- let selectedColumnsForRead = case selectedColumns opts of- Nothing -> Nothing- Just selected -> Just (L.nub (selected ++ predicateColumns))- let selectedColumnSet = S.fromList <$> selectedColumnsForRead- let shouldReadColumn colName _ =- case selectedColumnSet of- Nothing -> True- Just selected -> colName `S.member` selected-- case selectedColumnsForRead of- Nothing -> pure ()- Just requested ->- let missing = requested L.\\ availableSelectedColumns- in unless- (L.null missing)- ( throw- ( ColumnsNotFoundException- missing- "readParquetWithOpts"- availableSelectedColumns- )- )-- -- Collect per-column chunk lists; concatenate at the end to preserve bitmaps.- colListMap <- newIORef (M.empty :: M.Map T.Text [DI.Column])- lTypeMap <- newIORef (M.empty :: M.Map T.Text LogicalType)-- let schemaElements = schema fileMetadata- let sNodes = parseAll (drop 1 schemaElements)- let getTypeLength :: [String] -> Maybe Int32- getTypeLength colPath = findTypeLength schemaElements colPath (0 :: Int)- where- findTypeLength [] _ _ = Nothing- findTypeLength (s : ss) targetPath depth- | map T.unpack (pathToElement s ss depth) == targetPath- && elementType s == STRING- && typeLength s > 0 =- Just (typeLength s)- | otherwise =- findTypeLength ss targetPath (if numChildren s > 0 then depth + 1 else depth)-- pathToElement _ _ _ = []-- forM_ (rowGroups fileMetadata) $ \rowGroup -> do- forM_ (zip (rowGroupColumns rowGroup) [(0 :: Int) ..]) $ \(colChunk, colIdx) -> do- let metadata = columnMetaData colChunk- let colPath = columnPathInSchema metadata- let cleanPath = cleanColPath sNodes colPath- let colLeafName =- if null cleanPath- then T.pack $ "col_" ++ show colIdx- else T.pack $ last cleanPath- let colFullName =- if null cleanPath- then colLeafName- else T.intercalate "." $ map T.pack cleanPath-- when (shouldReadColumn colLeafName colPath) $ do- let colDataPageOffset = columnDataPageOffset metadata- let colDictionaryPageOffset = columnDictionaryPageOffset metadata- let colStart =- if colDictionaryPageOffset > 0 && colDataPageOffset > colDictionaryPageOffset- then colDictionaryPageOffset- else colDataPageOffset- let colLength = columnTotalCompressedSize metadata-- columnBytes <-- seekAndReadBytes- (Just (AbsoluteSeek, fromIntegral colStart))- (fromIntegral colLength)- file-- pages <- readAllPages (columnCodec metadata) columnBytes-- let maybeTypeLength =- if columnType metadata == PFIXED_LEN_BYTE_ARRAY- then getTypeLength colPath- else Nothing-- let primaryEncoding = maybe EPLAIN fst (L.uncons (columnEncodings metadata))-- let schemaTail = drop 1 (schema fileMetadata)- let (maxDef, maxRep) = levelsForPath schemaTail colPath- let lType =- maybe- LOGICAL_TYPE_UNKNOWN- logicalType- (findLeafSchema schemaTail colPath)- column <-- processColumnPages- (maxDef, maxRep)- pages- (columnType metadata)- primaryEncoding- maybeTypeLength- lType-- modifyIORef' colListMap (M.insertWith (++) colFullName [column])- modifyIORef' lTypeMap (M.insert colFullName lType)-- finalListMap <- readIORef colListMap- -- Reverse the accumulated lists (they were prepended) and concat columns per-name,- -- preserving bitmaps correctly via concatManyColumns.- let finalColMap = M.map (DI.concatManyColumns . reverse) finalListMap- finalLTypeMap <- readIORef lTypeMap- let orderedColumns =- map- ( \name ->- ( name- , applyLogicalType (finalLTypeMap M.! name) $ finalColMap M.! name- )- )- (filter (`M.member` finalColMap) columnNames)-- pure $ applyReadOptions opts (DI.fromNamedColumns orderedColumns)+_readParquetWithOpts extraConfig opts path =+ withFileBufferedOrSeekable extraConfig path ReadMode $ \file ->+ runReaderIO (parseParquetWithOpts opts) file {- | Read Parquet files from a directory or glob path. @@ -331,6 +231,287 @@ 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@@ -347,7 +528,7 @@ applySafeRead :: ParquetReadOptions -> DataFrame -> DataFrame applySafeRead opts df- | safeColumns opts = df{columns = V.map DI.ensureOptional (columns df)}+ | safeColumns opts = df{columns = Vector.map DI.ensureOptional (columns df)} | otherwise = df applyReadOptions :: ParquetReadOptions -> DataFrame -> DataFrame@@ -357,276 +538,50 @@ . applySelectedColumns opts . applyPredicate opts --- File and metadata parsing --------------------------------------------------- | read the file in memory at once, parse magicString and return the entire file ByteString-readMetadataFromPath :: FilePath -> IO (FileMetadata, BSO.ByteString)-readMetadataFromPath path = do- contents <- BSO.readFile path- let (size, magicString) = readMetadataSizeFromFooter contents- when (magicString /= "PAR1") $ error "Invalid Parquet file"- meta <- readMetadata contents size- pure (meta, contents)---- | read from the end of the file, parse magicString and return the entire file ByteString-readMetadataFromHandle :: FileBufferedOrSeekable -> IO FileMetadata-readMetadataFromHandle sh = do- footerBs <- readLastBytes (fromIntegral footerSize) sh- let (size, magicString) = readMetadataSizeFromFooterSlice footerBs- when (magicString /= "PAR1") $ error "Invalid Parquet file"- readMetadataByHandleMetaSize sh size---- | Takes the last 8 bit of the file to parse metadata size and magic string-readMetadataSizeFromFooterSlice :: BSO.ByteString -> (Int, BSO.ByteString)-readMetadataSizeFromFooterSlice contents =- let- size = fromIntegral (littleEndianWord32 contents)- magicString = BSO.take 4 (BSO.drop 4 contents)- in- (size, magicString)--readMetadataSizeFromFooter :: BSO.ByteString -> (Int, BSO.ByteString)-readMetadataSizeFromFooter = readMetadataSizeFromFooterSlice . BSO.takeEnd 8---- Schema navigation ---------------------------------------------------------getColumnPaths :: [SchemaElement] -> [(T.Text, Int)]-getColumnPaths schemaElements =- let nodes = parseAll schemaElements- in go nodes 0 [] False- where- go [] _ _ _ = []- go (n : ns) idx path skipThis- | null (sChildren n) =- let newPath = if skipThis then path else path ++ [T.pack (sName n)]- fullPath = T.intercalate "." newPath- in (fullPath, idx) : go ns (idx + 1) path skipThis- | sRep n == REPEATED =- let skipChildren = length (sChildren n) == 1- childLeaves = go (sChildren n) idx path skipChildren- in childLeaves ++ go ns (idx + length childLeaves) path skipThis- | skipThis =- let childLeaves = go (sChildren n) idx path False- in childLeaves ++ go ns (idx + length childLeaves) path skipThis- | otherwise =- let subPath = path ++ [T.pack (sName n)]- childLeaves = go (sChildren n) idx subPath False- in childLeaves ++ go ns (idx + length childLeaves) path skipThis--findLeafSchema :: [SchemaElement] -> [String] -> Maybe SchemaElement-findLeafSchema elems path =- case go (parseAll elems) path of- Just node -> L.find (\e -> T.unpack (elementName e) == sName node) elems- Nothing -> Nothing- where- go [] _ = Nothing- go _ [] = Nothing- go nodes [p] = L.find (\n -> sName n == p) nodes- go nodes (p : ps) = L.find (\n -> sName n == p) nodes >>= \n -> go (sChildren n) ps---- Page decoding -------------------------------------------------------------processColumnPages ::- (Int, Int) ->- [Page] ->- ParquetType ->- ParquetEncoding ->- Maybe Int32 ->- LogicalType ->- IO DI.Column-processColumnPages (maxDef, maxRep) pages pType _ maybeTypeLength _lType = do- let dictPages = filter isDictionaryPage pages- let dataPages = filter isDataPage pages-- let dictValsM =- case dictPages of- [] -> Nothing- (dictPage : _) ->- case pageTypeHeader (pageHeader dictPage) of- DictionaryPageHeader{..} ->- let countForBools =- if pType == PBOOLEAN- then Just dictionaryPageHeaderNumValues- else maybeTypeLength- in Just (readDictVals pType (pageBytes dictPage) countForBools)- _ -> Nothing-- cols <- forM dataPages $ \page -> do- let bs0 = pageBytes page- case pageTypeHeader (pageHeader page) of- DataPageHeader{..} -> do- let n = fromIntegral dataPageHeaderNumValues- (defLvls, repLvls, afterLvls) = readLevelsV1 n maxDef maxRep bs0- nPresent = length (filter (== maxDef) defLvls)- decodePageData- dictValsM- (maxDef, maxRep)- pType- maybeTypeLength- dataPageHeaderEncoding- defLvls- repLvls- nPresent- afterLvls- "v1"- DataPageHeaderV2{..} -> do- let n = fromIntegral dataPageHeaderV2NumValues- (defLvls, repLvls, afterLvls) =- readLevelsV2- n- maxDef- maxRep- definitionLevelByteLength- repetitionLevelByteLength- bs0- nPresent- | dataPageHeaderV2NumNulls > 0 =- fromIntegral (dataPageHeaderV2NumValues - dataPageHeaderV2NumNulls)- | otherwise = length (filter (== maxDef) defLvls)- decodePageData- dictValsM- (maxDef, maxRep)- pType- maybeTypeLength- dataPageHeaderV2Encoding- defLvls- repLvls- nPresent- afterLvls- "v2"-- -- Cannot happen as these are filtered out by isDataPage above- DictionaryPageHeader{} -> error "processColumnPages: impossible DictionaryPageHeader"- INDEX_PAGE_HEADER -> error "processColumnPages: impossible INDEX_PAGE_HEADER"- PAGE_TYPE_HEADER_UNKNOWN -> error "processColumnPages: impossible PAGE_TYPE_HEADER_UNKNOWN"- pure $ DI.concatManyColumns cols--decodePageData ::- Maybe DictVals ->- (Int, Int) ->- ParquetType ->- Maybe Int32 ->- ParquetEncoding ->- [Int] ->- [Int] ->- Int ->- BSO.ByteString ->- String ->- IO DI.Column-decodePageData dictValsM (maxDef, maxRep) pType maybeTypeLength encoding defLvls repLvls nPresent afterLvls versionLabel =- case encoding of- EPLAIN ->- case pType of- PBOOLEAN ->- let (vals, _) = readNBool nPresent afterLvls- in pure $- if maxRep > 0- then stitchForRepBool maxRep maxDef repLvls defLvls vals- else toMaybeBool maxDef defLvls vals- PINT32- | maxDef == 0- , maxRep == 0 ->- pure $ DI.fromUnboxedVector (readNInt32Vec nPresent afterLvls)- PINT32 ->- let (vals, _) = readNInt32 nPresent afterLvls- in pure $- if maxRep > 0- then stitchForRepInt32 maxRep maxDef repLvls defLvls vals- else toMaybeInt32 maxDef defLvls vals- PINT64- | maxDef == 0- , maxRep == 0 ->- pure $ DI.fromUnboxedVector (readNInt64Vec nPresent afterLvls)- PINT64 ->- let (vals, _) = readNInt64 nPresent afterLvls- in pure $- if maxRep > 0- then stitchForRepInt64 maxRep maxDef repLvls defLvls vals- else toMaybeInt64 maxDef defLvls vals- PINT96 ->- let (vals, _) = readNInt96Times nPresent afterLvls- in pure $- if maxRep > 0- then stitchForRepUTCTime maxRep maxDef repLvls defLvls vals- else toMaybeUTCTime maxDef defLvls vals- PFLOAT- | maxDef == 0- , maxRep == 0 ->- pure $ DI.fromUnboxedVector (readNFloatVec nPresent afterLvls)- PFLOAT ->- let (vals, _) = readNFloat nPresent afterLvls- in pure $- if maxRep > 0- then stitchForRepFloat maxRep maxDef repLvls defLvls vals- else toMaybeFloat maxDef defLvls vals- PDOUBLE- | maxDef == 0- , maxRep == 0 ->- pure $ DI.fromUnboxedVector (readNDoubleVec nPresent afterLvls)- PDOUBLE ->- let (vals, _) = readNDouble nPresent afterLvls- in pure $- if maxRep > 0- then stitchForRepDouble maxRep maxDef repLvls defLvls vals- else toMaybeDouble maxDef defLvls vals- PBYTE_ARRAY ->- let (raws, _) = readNByteArrays nPresent afterLvls- texts = map decodeUtf8Lenient raws- in pure $- if maxRep > 0- then stitchForRepText maxRep maxDef repLvls defLvls texts- else toMaybeText maxDef defLvls texts- PFIXED_LEN_BYTE_ARRAY ->- case maybeTypeLength of- Just len ->- let (raws, _) = splitFixed nPresent (fromIntegral len) afterLvls- texts = map decodeUtf8Lenient raws- in pure $- if maxRep > 0- then stitchForRepText maxRep maxDef repLvls defLvls texts- else toMaybeText maxDef defLvls texts- Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type length"- PARQUET_TYPE_UNKNOWN -> error "Cannot read unknown Parquet type"- ERLE_DICTIONARY -> decodeDictV1 dictValsM maxDef maxRep repLvls defLvls nPresent afterLvls- EPLAIN_DICTIONARY -> decodeDictV1 dictValsM maxDef maxRep repLvls defLvls nPresent afterLvls- other -> error ("Unsupported " ++ versionLabel ++ " encoding: " ++ show other)- -- Logical type conversion ------------------------------------------------- -applyLogicalType :: LogicalType -> DI.Column -> DI.Column-applyLogicalType (TimestampType _ unit) col =- fromRight col $- DI.mapColumn- (microsecondsToUTCTime . (* (1_000_000 `div` unitDivisor unit)))- col-applyLogicalType (DecimalType precision scale) col- | precision <= 9 = case DI.toVector @Int32 @VU.Vector col of- Right xs ->- DI.fromUnboxedVector $- VU.map (\raw -> fromIntegral @Int32 @Double raw / 10 ^ scale) xs- Left _ -> col- | precision <= 18 = case DI.toVector @Int64 @VU.Vector col of- Right xs ->- DI.fromUnboxedVector $- VU.map (\raw -> fromIntegral @Int64 @Double raw / 10 ^ scale) xs- Left _ -> col- | otherwise = col+{- | 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) -unitDivisor :: TimeUnit -> Int64-unitDivisor MILLISECONDS = 1_000-unitDivisor MICROSECONDS = 1_000_000-unitDivisor NANOSECONDS = 1_000_000_000-unitDivisor TIME_UNIT_UNKNOWN = 1--applyScale :: Int32 -> Int32 -> Double-applyScale scale rawValue =- fromIntegral rawValue / (10 ^ scale)- -- HuggingFace support ----------------------------------------------------- data HFRef = HFRef@@ -670,7 +625,7 @@ _ -> Left $ "Invalid hf:// URI (expected hf://datasets/owner/dataset/glob): " ++ path -getHFToken :: IO (Maybe BSO.ByteString)+getHFToken :: IO (Maybe BS.ByteString) getHFToken = do envToken <- lookupEnv "HF_TOKEN" case envToken of@@ -678,9 +633,9 @@ Nothing -> do home <- getHomeDirectory let tokenPath = home </> ".cache" </> "huggingface" </> "token"- result <- try (BSO.readFile tokenPath) :: IO (Either IOError BSO.ByteString)+ result <- try (BS.readFile tokenPath) :: IO (Either IOError BS.ByteString) case result of- Right bs -> pure (Just (BSO.takeWhile (/= 10) bs))+ Right bs -> pure (Just (BS.takeWhile (/= 10) bs)) Left _ -> pure Nothing {- | Extract the repo-relative path from a HuggingFace download URL.@@ -700,7 +655,7 @@ matchesGlob :: T.Text -> HFParquetFile -> Bool matchesGlob g f = match (compile (T.unpack g)) (hfUrlRepoPath f) -resolveHFUrls :: Maybe BSO.ByteString -> HFRef -> IO [HFParquetFile]+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@@ -721,7 +676,7 @@ Left err -> ioError $ userError $ "Failed to parse HF API response: " ++ err Right hfResp -> pure $ filter (matchesGlob (hfGlob ref)) (hfParquetFiles hfResp) -downloadHFFiles :: Maybe BSO.ByteString -> [HFParquetFile] -> IO [FilePath]+downloadHFFiles :: Maybe BS.ByteString -> [HFParquetFile] -> IO [FilePath] downloadHFFiles mToken files = do tmpDir <- getTemporaryDirectory forM files $ \f -> do@@ -740,7 +695,7 @@ ioError $ userError $ "Failed to download " ++ T.unpack (hfpUrl f) ++ " (HTTP " ++ show status ++ ")"- BSO.writeFile destPath (getResponseBody resp)+ BS.writeFile destPath (getResponseBody resp) pure destPath -- | True when the path contains glob wildcard characters.
− src/DataFrame/IO/Parquet/ColumnStatistics.hs
@@ -1,19 +0,0 @@-module DataFrame.IO.Parquet.ColumnStatistics where--import qualified Data.ByteString as BS-import Data.Int (Int64)--data ColumnStatistics = ColumnStatistics- { columnMin :: BS.ByteString- , columnMax :: BS.ByteString- , columnNullCount :: Int64- , columnDistictCount :: Int64- , columnMinValue :: BS.ByteString- , columnMaxValue :: BS.ByteString- , isColumnMaxValueExact :: Bool- , isColumnMinValueExact :: Bool- }- deriving (Show, Eq)--emptyColumnStatistics :: ColumnStatistics-emptyColumnStatistics = ColumnStatistics BS.empty BS.empty 0 0 BS.empty BS.empty False False
− src/DataFrame/IO/Parquet/Compression.hs
@@ -1,26 +0,0 @@-module DataFrame.IO.Parquet.Compression where--import Data.Int--data CompressionCodec- = UNCOMPRESSED- | SNAPPY- | GZIP- | LZO- | BROTLI- | LZ4- | ZSTD- | LZ4_RAW- | COMPRESSION_CODEC_UNKNOWN- deriving (Show, Eq)--compressionCodecFromInt :: Int32 -> CompressionCodec-compressionCodecFromInt 0 = UNCOMPRESSED-compressionCodecFromInt 1 = SNAPPY-compressionCodecFromInt 2 = GZIP-compressionCodecFromInt 3 = LZO-compressionCodecFromInt 4 = BROTLI-compressionCodecFromInt 5 = LZ4-compressionCodecFromInt 6 = ZSTD-compressionCodecFromInt 7 = LZ4_RAW-compressionCodecFromInt _ = COMPRESSION_CODEC_UNKNOWN
+ 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
@@ -1,54 +1,59 @@ {-# LANGUAGE BangPatterns #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE MonoLocalBinds #-}-{-# LANGUAGE OverloadedStrings #-} -module DataFrame.IO.Parquet.Dictionary where+module DataFrame.IO.Parquet.Dictionary (DictVals (..), readDictVals, decodeRLEBitPackedHybrid) where -import Control.Monad import Data.Bits import qualified Data.ByteString as BS-import Data.IORef-import Data.Int-import Data.Maybe+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+import Data.Time (UTCTime) import qualified Data.Vector as V-import qualified Data.Vector.Mutable as VM-import qualified Data.Vector.Unboxed as VU-import DataFrame.IO.Parquet.Encoding-import DataFrame.IO.Parquet.Levels-import DataFrame.IO.Parquet.Time-import DataFrame.IO.Parquet.Types+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 qualified DataFrame.Internal.Column as DI import GHC.Float -dictCardinality :: DictVals -> Int-dictCardinality (DBool ds) = V.length ds-dictCardinality (DInt32 ds) = V.length ds-dictCardinality (DInt64 ds) = V.length ds-dictCardinality (DInt96 ds) = V.length ds-dictCardinality (DFloat ds) = V.length ds-dictCardinality (DDouble ds) = V.length ds-dictCardinality (DText ds) = V.length ds+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) -readDictVals :: ParquetType -> BS.ByteString -> Maybe Int32 -> DictVals-readDictVals PBOOLEAN bs (Just count) = DBool (V.fromList (take (fromIntegral count) $ readPageBool bs))-readDictVals PINT32 bs _ = DInt32 (V.fromList (readPageInt32 bs))-readDictVals PINT64 bs _ = DInt64 (V.fromList (readPageInt64 bs))-readDictVals PINT96 bs _ = DInt96 (V.fromList (readPageInt96Times bs))-readDictVals PFLOAT bs _ = DFloat (V.fromList (readPageFloat bs))-readDictVals PDOUBLE bs _ = DDouble (V.fromList (readPageWord64 bs))-readDictVals PBYTE_ARRAY bs _ = DText (V.fromList (readPageBytes bs))-readDictVals PFIXED_LEN_BYTE_ARRAY bs (Just len) = DText (V.fromList (readPageFixedBytes bs (fromIntegral len)))-readDictVals t _ _ = error $ "Unsupported dictionary type: " ++ show t+{- | 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 = []@@ -109,199 +114,51 @@ | otherwise = decodeUtf8Lenient (BS.take len xs) : readPageFixedBytes (BS.drop len xs) len -{- | Dispatch to the right multi-level list stitching function.-For maxRep=1 uses stitchList; for 2/3 uses stitchList2/3 with computed thresholds.-Threshold formula: defT_r = maxDef - 2*(maxRep - r).--}-stitchForRepBool :: Int -> Int -> [Int] -> [Int] -> [Bool] -> DI.Column-stitchForRepBool maxRep maxDef rep def vals = case maxRep of- 2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)- 3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)- _ -> DI.fromList (stitchList maxDef rep def vals)--stitchForRepInt32 :: Int -> Int -> [Int] -> [Int] -> [Int32] -> DI.Column-stitchForRepInt32 maxRep maxDef rep def vals = case maxRep of- 2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)- 3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)- _ -> DI.fromList (stitchList maxDef rep def vals)--stitchForRepInt64 :: Int -> Int -> [Int] -> [Int] -> [Int64] -> DI.Column-stitchForRepInt64 maxRep maxDef rep def vals = case maxRep of- 2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)- 3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)- _ -> DI.fromList (stitchList maxDef rep def vals)--stitchForRepUTCTime :: Int -> Int -> [Int] -> [Int] -> [UTCTime] -> DI.Column-stitchForRepUTCTime maxRep maxDef rep def vals = case maxRep of- 2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)- 3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)- _ -> DI.fromList (stitchList maxDef rep def vals)--stitchForRepFloat :: Int -> Int -> [Int] -> [Int] -> [Float] -> DI.Column-stitchForRepFloat maxRep maxDef rep def vals = case maxRep of- 2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)- 3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)- _ -> DI.fromList (stitchList maxDef rep def vals)--stitchForRepDouble :: Int -> Int -> [Int] -> [Int] -> [Double] -> DI.Column-stitchForRepDouble maxRep maxDef rep def vals = case maxRep of- 2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)- 3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)- _ -> DI.fromList (stitchList maxDef rep def vals)--stitchForRepText :: Int -> Int -> [Int] -> [Int] -> [T.Text] -> DI.Column-stitchForRepText maxRep maxDef rep def vals = case maxRep of- 2 -> DI.fromList (stitchList2 (maxDef - 2) maxDef rep def vals)- 3 -> DI.fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef rep def vals)- _ -> DI.fromList (stitchList maxDef rep def vals)--{- | Build a Column from a dictionary + index vector + def levels in a single-mutable-vector pass, avoiding the intermediate [a] and [Maybe a] lists.-For maxRep > 0 (list columns) the caller must use the rep-stitching path instead.--}-applyDictToColumn ::- (DI.Columnable a, DI.Columnable (Maybe a)) =>- V.Vector a ->- VU.Vector Int ->- Int -> -- maxDef- [Int] -> -- defLvls- IO DI.Column-applyDictToColumn dict idxs maxDef defLvls- | maxDef == 0 = do- -- All rows are required; no nullability to check.- let n = VU.length idxs- pure $ DI.fromVector (V.generate n (\i -> dict V.! (idxs VU.! i)))- | otherwise = do- let n = length defLvls- mv <- VM.new n- hasNullRef <- newIORef False- let go _ _ [] = pure ()- go !i !j (d : ds)- | d == maxDef = do- VM.write mv i (Just (dict V.! (idxs VU.! j)))- go (i + 1) (j + 1) ds- | otherwise = do- writeIORef hasNullRef True- VM.write mv i Nothing- go (i + 1) j ds- go 0 0 defLvls- vec <- V.freeze mv- hasNull <- readIORef hasNullRef- pure $- if hasNull- then DI.fromVector vec -- VB.Vector (Maybe a) → OptionalColumn- else DI.fromVector (V.map fromJust vec) -- VB.Vector a → BoxedColumn/UnboxedColumn--decodeDictV1 ::- Maybe DictVals ->- Int ->- Int ->- [Int] ->- [Int] ->- Int ->- BS.ByteString ->- IO DI.Column-decodeDictV1 dictValsM maxDef maxRep repLvls defLvls nPresent bytes =- case dictValsM of- Nothing -> error "Dictionary-encoded page but dictionary is missing"- Just dictVals ->- let (idxs, _rest) = decodeDictIndicesV1 nPresent (dictCardinality dictVals) bytes- in do- when (VU.length idxs /= nPresent) $- error $- "dict index count mismatch: got "- ++ show (VU.length idxs)- ++ ", expected "- ++ show nPresent- if maxRep > 0- then do- case dictVals of- DBool ds ->- pure $- stitchForRepBool maxRep maxDef repLvls defLvls (map (ds V.!) (VU.toList idxs))- DInt32 ds ->- pure $- stitchForRepInt32 maxRep maxDef repLvls defLvls (map (ds V.!) (VU.toList idxs))- DInt64 ds ->- pure $- stitchForRepInt64 maxRep maxDef repLvls defLvls (map (ds V.!) (VU.toList idxs))- DInt96 ds ->- pure $- stitchForRepUTCTime- maxRep- maxDef- repLvls- defLvls- (map (ds V.!) (VU.toList idxs))- DFloat ds ->- pure $- stitchForRepFloat maxRep maxDef repLvls defLvls (map (ds V.!) (VU.toList idxs))- DDouble ds ->- pure $- stitchForRepDouble maxRep maxDef repLvls defLvls (map (ds V.!) (VU.toList idxs))- DText ds ->- pure $- stitchForRepText maxRep maxDef repLvls defLvls (map (ds V.!) (VU.toList idxs))- else case dictVals of- -- Fast path: unboxable types, no nulls — one allocation via VU.map- DInt32 ds | maxDef == 0 -> pure $ DI.fromUnboxedVector (VU.map (ds V.!) idxs)- DInt64 ds | maxDef == 0 -> pure $ DI.fromUnboxedVector (VU.map (ds V.!) idxs)- DFloat ds | maxDef == 0 -> pure $ DI.fromUnboxedVector (VU.map (ds V.!) idxs)- DDouble ds | maxDef == 0 -> pure $ DI.fromUnboxedVector (VU.map (ds V.!) idxs)- DBool ds -> applyDictToColumn ds idxs maxDef defLvls- DInt32 ds -> applyDictToColumn ds idxs maxDef defLvls- DInt64 ds -> applyDictToColumn ds idxs maxDef defLvls- DInt96 ds -> applyDictToColumn ds idxs maxDef defLvls- DFloat ds -> applyDictToColumn ds idxs maxDef defLvls- DDouble ds -> applyDictToColumn ds idxs maxDef defLvls- DText ds -> applyDictToColumn ds idxs maxDef defLvls--toMaybeInt32 :: Int -> [Int] -> [Int32] -> DI.Column-toMaybeInt32 maxDef def xs =- let filled = stitchNullable maxDef def xs- in if all isJust filled- then DI.fromList (map (fromMaybe 0) filled)- else DI.fromList filled--toMaybeDouble :: Int -> [Int] -> [Double] -> DI.Column-toMaybeDouble maxDef def xs =- let filled = stitchNullable maxDef def xs- in if all isJust filled- then DI.fromList (map (fromMaybe 0) filled)- else DI.fromList filled--toMaybeText :: Int -> [Int] -> [T.Text] -> DI.Column-toMaybeText maxDef def xs =- let filled = stitchNullable maxDef def xs- in if all isJust filled- then DI.fromList (map (fromMaybe "") filled)- else DI.fromList filled--toMaybeBool :: Int -> [Int] -> [Bool] -> DI.Column-toMaybeBool maxDef def xs =- let filled = stitchNullable maxDef def xs- in if all isJust filled- then DI.fromList (map (fromMaybe False) filled)- else DI.fromList filled--toMaybeInt64 :: Int -> [Int] -> [Int64] -> DI.Column-toMaybeInt64 maxDef def xs =- let filled = stitchNullable maxDef def xs- in if all isJust filled- then DI.fromList (map (fromMaybe 0) filled)- else DI.fromList filled+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) -toMaybeFloat :: Int -> [Int] -> [Float] -> DI.Column-toMaybeFloat maxDef def xs =- let filled = stitchNullable maxDef def xs- in if all isJust filled- then DI.fromList (map (fromMaybe 0.0) filled)- else DI.fromList filled+-- | 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 -toMaybeUTCTime :: Int -> [Int] -> [UTCTime] -> DI.Column-toMaybeUTCTime maxDef def times =- let filled = stitchNullable maxDef def times- defaultTime = UTCTime (fromGregorian 1970 1 1) (secondsToDiffTime 0)- in if all isJust filled- then DI.fromList (map (fromMaybe defaultTime) filled)- else DI.fromList filled+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
@@ -1,8 +1,18 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} -module DataFrame.IO.Parquet.Encoding where+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@@ -10,10 +20,15 @@ import Data.List (foldl') #endif 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@@ -22,73 +37,99 @@ bitWidthForMaxLevel :: Int -> Int bitWidthForMaxLevel maxLevel = ceilLog2 (maxLevel + 1) -bytesForBW :: Int -> Int-bytesForBW bw = (bw + 7) `div` 8--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+-- ---------------------------------------------------------------------------+-- Vector-based RLE / bit-packed hybrid decoder+-- --------------------------------------------------------------------------- decodeRLEBitPackedHybrid ::- Int -> Int -> BS.ByteString -> ([Word32], BS.ByteString)+ -- | 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 = (replicate need 0, bs)- | otherwise = go 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 :: Word32- mask = if bw == 32 then maxBound else (1 `shiftL` bw) - 1- go :: Int -> BS.ByteString -> [Word32] -> ([Word32], BS.ByteString)- go 0 rest acc = (reverse acc, rest)- go n rest acc- | BS.null rest = (reverse acc, rest)+ !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 rest+ let (hdr64, afterHdr) = readUVarInt buf isPacked = (hdr64 .&. 1) == 1 in if isPacked- then+ then do let groups = fromIntegral (hdr64 `shiftR` 1) :: Int totalVals = groups * 8- (valsAll, afterRun) = unpackBitPacked bw totalVals afterHdr- takeN = min n totalVals- actualTaken = take takeN valsAll- in go (n - takeN) afterRun (reverse actualTaken ++ acc)- else+ 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 = bytesForBW bw- word32 = littleEndianWord32 (BS.take 4 afterHdr)- afterV = BS.drop nbytes afterHdr- val = word32 .&. mask- takeN = min n runLen- in go (n - takeN) afterV (replicate takeN val ++ acc)+ 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 #-} -decodeDictIndicesV1 ::- Int -> Int -> BS.ByteString -> (VU.Vector Int, BS.ByteString)-decodeDictIndicesV1 need _dictCard bs =- case BS.uncons bs of- Nothing -> error "empty dictionary index stream"- Just (w0, rest0) ->- let bw = fromIntegral w0 :: Int- (u32s, rest1) = decodeRLEBitPackedHybrid bw need rest0- in (VU.fromList (map fromIntegral u32s), rest1)+-- | 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
@@ -1,145 +1,142 @@-module DataFrame.IO.Parquet.Levels where+{-# LANGUAGE BangPatterns #-} -import qualified Data.ByteString as BS-import Data.Int-import Data.List-import qualified Data.Text as T+module DataFrame.IO.Parquet.Levels (+ -- Level readers+ readLevelsV1,+ readLevelsV2,+ -- Stitch functions+ stitchList,+ stitchList2,+ stitchList3,+) where -import DataFrame.IO.Parquet.Encoding-import DataFrame.IO.Parquet.Thrift-import DataFrame.IO.Parquet.Types+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) -readLevelsV1 ::- Int -> Int -> Int -> BS.ByteString -> ([Int], [Int], BS.ByteString)-readLevelsV1 n maxDef maxRep bs =- let bwDef = bitWidthForMaxLevel maxDef- bwRep = bitWidthForMaxLevel maxRep+-- ---------------------------------------------------------------------------+-- Level readers+-- --------------------------------------------------------------------------- - (repLvls, afterRep) =- if bwRep == 0- then (replicate n 0, bs)- else- let repLength = littleEndianWord32 (BS.take 4 bs)- repData = BS.take (fromIntegral repLength) (BS.drop 4 bs)- afterRepData = BS.drop (4 + fromIntegral repLength) bs- (repVals, _) = decodeRLEBitPackedHybrid bwRep n repData- in (map fromIntegral repVals, afterRepData)+{- | 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) - (defLvls, afterDef) =- if bwDef == 0- then (replicate n 0, afterRep)- else- let defLength = littleEndianWord32 (BS.take 4 afterRep)- defData = BS.take (fromIntegral defLength) (BS.drop 4 afterRep)- afterDefData = BS.drop (4 + fromIntegral defLength) afterRep- (defVals, _) = decodeRLEBitPackedHybrid bwDef n defData- in (map fromIntegral defVals, afterDefData)- in (defLvls, repLvls, afterDef)+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 ->- ([Int], [Int], BS.ByteString)-readLevelsV2 n maxDef maxRep defLen repLen bs =+ (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- bwDef = bitWidthForMaxLevel maxDef bwRep = bitWidthForMaxLevel maxRep- (repLvlsRaw, _) =- if bwRep == 0- then (replicate n 0, repBytes)- else decodeRLEBitPackedHybrid bwRep n repBytes- (defLvlsRaw, _) =- if bwDef == 0- then (replicate n 0, defBytes)- else decodeRLEBitPackedHybrid bwDef n defBytes- in (map fromIntegral defLvlsRaw, map fromIntegral repLvlsRaw, afterDefBytes)--stitchNullable :: Int -> [Int] -> [a] -> [Maybe a]-stitchNullable maxDef = go- where- go [] _ = []- go (d : ds) vs- | d == maxDef = case vs of- (v : vs') -> Just v : go ds vs'- [] -> error "value stream exhausted"- | otherwise = Nothing : go ds vs--data SNode = SNode- { sName :: String- , sRep :: RepetitionType- , sChildren :: [SNode]- }- deriving (Show, Eq)--parseOne :: [SchemaElement] -> (SNode, [SchemaElement])-parseOne [] = error "parseOne: empty schema list"-parseOne (se : rest) =- let childCount = fromIntegral (numChildren se)- (kids, rest') = parseMany childCount rest- in ( SNode- { sName = T.unpack (elementName se)- , sRep = repetitionType se- , sChildren = kids- }- , rest'- )--parseMany :: Int -> [SchemaElement] -> ([SNode], [SchemaElement])-parseMany 0 xs = ([], xs)-parseMany n xs =- let (node, xs') = parseOne xs- (nodes, xs'') = parseMany (n - 1) xs'- in (node : nodes, xs'')--parseAll :: [SchemaElement] -> [SNode]-parseAll [] = []-parseAll xs = let (n, xs') = parseOne xs in n : parseAll xs'---- | Tag leaf values as Just/Nothing according to maxDef.-pairWithVals :: Int -> [(Int, Int)] -> [a] -> [(Int, Int, Maybe a)]-pairWithVals _ [] _ = []-pairWithVals maxDef ((r, d) : rds) vs- | d == maxDef = case vs of- (v : vs') -> (r, d, Just v) : pairWithVals maxDef rds vs'- [] -> error "pairWithVals: value stream exhausted"- | otherwise = (r, d, Nothing) : pairWithVals maxDef rds vs---- | Split triplets into groups; 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+ 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) -{- | Reconstruct a list column from Dremel encoding levels.-rep=0 starts a new top-level row; def=0 means the entire list slot is null.-Returns one Maybe [Maybe a] per row.+{- | 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 -> [Int] -> [Int] -> [a] -> [Maybe [Maybe a]]-stitchList maxDef repLvls defLvls vals =- let triplets = pairWithVals maxDef (zip repLvls defLvls) vals- rows = splitAtRepBound 0 triplets- in map toRow rows+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] -{- | Reconstruct a 2-level nested list (maxRep=2) from Dremel triplets.-defT1: def threshold at which the depth-1 element is present (not null).-maxDef: def threshold at which the leaf is present.+{- | Stitch a doubly-nested list column (@maxRep == 2@).+@defT1@ is the def threshold at which the depth-1 element is present. -}-stitchList2 :: Int -> Int -> [Int] -> [Int] -> [a] -> [Maybe [Maybe [Maybe a]]]-stitchList2 defT1 maxDef repLvls defLvls vals =- let triplets = pairWithVals maxDef (zip repLvls defLvls) vals- in map toRow (splitAtRepBound 0 triplets)+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))@@ -149,16 +146,22 @@ toLeaf [] = Nothing toLeaf ((_, _, v) : _) = v -{- | Reconstruct a 3-level nested list (maxRep=3) from Dremel triplets.-defT1, defT2: def thresholds at which depth-1 and depth-2 elements are present.-maxDef: def threshold at which the leaf is present.+{- | 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 -> [Int] -> [Int] -> [a] -> [Maybe [Maybe [Maybe [Maybe a]]]]-stitchList3 defT1 defT2 maxDef repLvls defLvls vals =- let triplets = pairWithVals maxDef (zip repLvls defLvls) vals- in map toRow (splitAtRepBound 0 triplets)+ 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))@@ -171,14 +174,37 @@ toLeaf [] = Nothing toLeaf ((_, _, v) : _) = v -levelsForPath :: [SchemaElement] -> [String] -> (Int, Int)-levelsForPath schemaTail = go 0 0 (parseAll schemaTail)+-- ---------------------------------------------------------------------------+-- 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- go defC repC _ [] = (defC, repC)- go defC repC nodes (p : ps) =- case find (\n -> sName n == p) nodes of- Nothing -> (defC, repC)- Just n ->- let defC' = defC + (if sRep n == OPTIONAL || sRep n == REPEATED then 1 else 0)- repC' = repC + (if sRep n == REPEATED then 1 else 0)- in go defC' repC' (sChildren n) ps+ 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
@@ -1,469 +1,352 @@-{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE OverloadedRecordDot #-}+{-# LANGUAGE ScopedTypeVariables #-} -module DataFrame.IO.Parquet.Page where+module DataFrame.IO.Parquet.Page (+ -- Types+ PageDecoder,+ UnboxedPageDecoder,+ -- Per-type decoders+ boolDecoder,+ int32Decoder,+ int64Decoder,+ int96Decoder,+ floatDecoder,+ doubleDecoder,+ byteArrayDecoder,+ fixedLenByteArrayDecoder,+ -- Page iteration+ readPages,+) where -import qualified Codec.Compression.GZip as GZip-import qualified Codec.Compression.Zstd.Streaming as Zstd-import Data.Bits+import Control.Monad.IO.Class (MonadIO (liftIO))+import Data.Bits (shiftR, (.&.)) import qualified Data.ByteString as BS-import qualified Data.ByteString.Lazy as LB-import Data.Int-import Data.Maybe (fromMaybe)+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.Binary-import DataFrame.IO.Parquet.Thrift-import DataFrame.IO.Parquet.Types+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-import qualified Snappy+import GHC.Float (castWord32ToFloat, castWord64ToDouble)+import Pinch (decodeWithLeftovers)+import qualified Pinch+import Streamly.Internal.Data.Unfold (Step (..), Unfold, mkUnfoldM) -isDataPage :: Page -> Bool-isDataPage page = case pageTypeHeader (pageHeader page) of- DataPageHeader{} -> True- DataPageHeaderV2{} -> True- _ -> False+-- ---------------------------------------------------------------------------+-- Types+-- --------------------------------------------------------------------------- -isDictionaryPage :: Page -> Bool-isDictionaryPage page = case pageTypeHeader (pageHeader page) of- DictionaryPageHeader{} -> True- _ -> False+{- | 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 -readPage :: CompressionCodec -> BS.ByteString -> IO (Maybe Page, BS.ByteString)-readPage c columnBytes =- if BS.null columnBytes- then pure (Nothing, BS.empty)- else do- let (hdr, remainder) = readPageHeader emptyPageHeader columnBytes 0+type UnboxedPageDecoder a =+ Maybe DictVals -> Encoding -> Int -> BS.ByteString -> VU.Vector a - let compressed = BS.take (fromIntegral $ compressedPageSize hdr) remainder+-- ---------------------------------------------------------------------------+-- Per-type decoders+-- --------------------------------------------------------------------------- - fullData <- case c of- ZSTD -> do- result <- Zstd.decompress- drainZstd result compressed []- where- drainZstd (Zstd.Consume f) input acc = do- result <- f input- drainZstd result BS.empty acc- drainZstd (Zstd.Produce chunk next) _ acc = do- result <- next- drainZstd result BS.empty (chunk : acc)- drainZstd (Zstd.Done final) _ acc =- pure $ BS.concat (reverse (final : acc))- drainZstd (Zstd.Error msg msg2) _ _ =- error ("ZSTD error: " ++ msg ++ " " ++ msg2)- 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)- pure- ( Just $ Page hdr fullData- , BS.drop (fromIntegral $ compressedPageSize hdr) remainder- )+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) -readPageHeader ::- PageHeader -> BS.ByteString -> Int16 -> (PageHeader, BS.ByteString)-readPageHeader hdr xs lastFieldId =- if BS.null xs- then (hdr, BS.empty)- else- let- fieldContents = readField' xs lastFieldId- in- case fieldContents of- Nothing -> (hdr, BS.drop 1 xs)- Just (remainder, _elemType, identifier) -> case identifier of- 1 ->- let- (pType, remainder') = readInt32FromBytes remainder- in- readPageHeader- (hdr{pageHeaderPageType = pageTypeFromInt pType})- remainder'- identifier- 2 ->- let- (parsedUncompressedPageSize, remainder') = readInt32FromBytes remainder- in- readPageHeader- (hdr{uncompressedPageSize = parsedUncompressedPageSize})- remainder'- identifier- 3 ->- let- (parsedCompressedPageSize, remainder') = readInt32FromBytes remainder- in- readPageHeader- (hdr{compressedPageSize = parsedCompressedPageSize})- remainder'- identifier- 4 ->- let- (crc, remainder') = readInt32FromBytes remainder- in- readPageHeader (hdr{pageHeaderCrcChecksum = crc}) remainder' identifier- 5 ->- let- (dataPageHeader, remainder') = readPageTypeHeader emptyDataPageHeader remainder 0- in- readPageHeader (hdr{pageTypeHeader = dataPageHeader}) remainder' identifier- 6 -> error "Index page header not supported"- 7 ->- let- (dictionaryPageHeader, remainder') = readPageTypeHeader emptyDictionaryPageHeader remainder 0- in- readPageHeader- (hdr{pageTypeHeader = dictionaryPageHeader})- remainder'- identifier- 8 ->- let- (dataPageHeaderV2, remainder') = readPageTypeHeader emptyDataPageHeaderV2 remainder 0- in- readPageHeader (hdr{pageTypeHeader = dataPageHeaderV2}) remainder' identifier- n -> error $ "Unknown page header field " ++ show n+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) -readPageTypeHeader ::- PageTypeHeader -> BS.ByteString -> Int16 -> (PageTypeHeader, BS.ByteString)-readPageTypeHeader INDEX_PAGE_HEADER _ _ = error "readPageTypeHeader: unsupported INDEX_PAGE_HEADER"-readPageTypeHeader PAGE_TYPE_HEADER_UNKNOWN _ _ = error "readPageTypeHeader: unsupported PAGE_TYPE_HEADER_UNKNOWN"-readPageTypeHeader hdr@(DictionaryPageHeader{}) xs lastFieldId =- if BS.null xs- then (hdr, BS.empty)- else- let- fieldContents = readField' xs lastFieldId- in- case fieldContents of- Nothing -> (hdr, BS.drop 1 xs)- Just (remainder, _elemType, identifier) -> case identifier of- 1 ->- let- (numValues, remainder') = readInt32FromBytes remainder- in- readPageTypeHeader- (hdr{dictionaryPageHeaderNumValues = numValues})- remainder'- identifier- 2 ->- let- (enc, remainder') = readInt32FromBytes remainder- in- readPageTypeHeader- (hdr{dictionaryPageHeaderEncoding = parquetEncodingFromInt enc})- remainder'- identifier- 3 ->- let- isSorted = fromMaybe (error "readPageTypeHeader: not enough bytes") (remainder BS.!? 0)- in- readPageTypeHeader- (hdr{dictionaryPageIsSorted = isSorted == compactBooleanTrue})- -- TODO(mchavinda): The bool logic here is a little tricky.- -- If the field is a bool then you can get the value- -- from the byte (and you don't have to drop a field).- -- But in other cases you do.- -- This might become a problem later but in the mean- -- time I'm not dropping (this assumes this is the common case).- remainder- identifier- n ->- error $ "readPageTypeHeader: unsupported identifier " ++ show n-readPageTypeHeader hdr@(DataPageHeader{}) xs lastFieldId =- if BS.null xs- then (hdr, BS.empty)- else- let- fieldContents = readField' xs lastFieldId- in- case fieldContents of- Nothing -> (hdr, BS.drop 1 xs)- Just (remainder, _elemType, identifier) -> case identifier of- 1 ->- let- (numValues, remainder') = readInt32FromBytes remainder- in- readPageTypeHeader- (hdr{dataPageHeaderNumValues = numValues})- remainder'- identifier- 2 ->- let- (enc, remainder') = readInt32FromBytes remainder- in- readPageTypeHeader- (hdr{dataPageHeaderEncoding = parquetEncodingFromInt enc})- remainder'- identifier- 3 ->- let- (enc, remainder') = readInt32FromBytes remainder- in- readPageTypeHeader- (hdr{definitionLevelEncoding = parquetEncodingFromInt enc})- remainder'- identifier- 4 ->- let- (enc, remainder') = readInt32FromBytes remainder- in- readPageTypeHeader- (hdr{repetitionLevelEncoding = parquetEncodingFromInt enc})- remainder'- identifier- 5 ->- let- (stats, remainder') = readStatisticsFromBytes emptyColumnStatistics remainder 0- in- readPageTypeHeader (hdr{dataPageHeaderStatistics = stats}) remainder' identifier- n -> error $ show n-readPageTypeHeader hdr@(DataPageHeaderV2{}) xs lastFieldId =- if BS.null xs- then (hdr, BS.empty)- else- let- fieldContents = readField' xs lastFieldId- in- case fieldContents of- Nothing -> (hdr, BS.drop 1 xs)- Just (remainder, _elemType, identifier) -> case identifier of- 1 ->- let- (numValues, remainder') = readInt32FromBytes remainder- in- readPageTypeHeader- (hdr{dataPageHeaderV2NumValues = numValues})- remainder'- identifier- 2 ->- let- (numNulls, remainder') = readInt32FromBytes remainder- in- readPageTypeHeader- (hdr{dataPageHeaderV2NumNulls = numNulls})- remainder'- identifier- 3 ->- let- (parsedNumRows, remainder') = readInt32FromBytes remainder- in- readPageTypeHeader- (hdr{dataPageHeaderV2NumRows = parsedNumRows})- remainder'- identifier- 4 ->- let- (enc, remainder') = readInt32FromBytes remainder- in- readPageTypeHeader- (hdr{dataPageHeaderV2Encoding = parquetEncodingFromInt enc})- remainder'- identifier- 5 ->- let- (n, remainder') = readInt32FromBytes remainder- in- readPageTypeHeader (hdr{definitionLevelByteLength = n}) remainder' identifier- 6 ->- let- (n, remainder') = readInt32FromBytes remainder- in- readPageTypeHeader (hdr{repetitionLevelByteLength = n}) remainder' identifier- 7 ->- let- (isCompressed, remainder') = case BS.uncons remainder of- Just (b, bytes) -> ((b .&. 0x0f) == compactBooleanTrue, bytes)- Nothing -> (True, BS.empty)- in- readPageTypeHeader- (hdr{dataPageHeaderV2IsCompressed = isCompressed})- remainder'- identifier- 8 ->- let- (stats, remainder') = readStatisticsFromBytes emptyColumnStatistics remainder 0- in- readPageTypeHeader- (hdr{dataPageHeaderV2Statistics = stats})- remainder'- identifier- n -> error $ show n+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) -readField' :: BS.ByteString -> Int16 -> Maybe (BS.ByteString, TType, Int16)-readField' bs lastFieldId = case BS.uncons bs of- Nothing -> Nothing- Just (x, xs) ->- if x .&. 0x0f == 0- then Nothing- else- let modifier = fromIntegral ((x .&. 0xf0) `shiftR` 4) :: Int16- (identifier, remainder) =- if modifier == 0- then readIntFromBytes @Int16 xs- else (lastFieldId + modifier, xs)- elemType = toTType (x .&. 0x0f)- in Just (remainder, elemType, identifier)+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) -readAllPages :: CompressionCodec -> BS.ByteString -> IO [Page]-readAllPages codec bytes = go bytes []+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- go bs acc =- if BS.null bs- then return (reverse acc)- else do- (maybePage, remainderaining) <- readPage codec bs- case maybePage of- Nothing -> return (reverse acc)- Just page -> go remainderaining (page : acc)+ getFloat (DFloat ds) i = ds VB.! i+ getFloat d _ = error ("floatDecoder: wrong dict type, got " ++ show d) --- | Read n Int32 values directly into an unboxed vector (no intermediate list).-readNInt32Vec :: Int -> BS.ByteString -> VU.Vector Int32-readNInt32Vec n bs = VU.generate n (\i -> littleEndianInt32 (BS.drop (4 * i) bs))+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) --- | Read n Int64 values directly into an unboxed vector.-readNInt64Vec :: Int -> BS.ByteString -> VU.Vector Int64-readNInt64Vec n bs = VU.generate n (\i -> fromIntegral (littleEndianWord64 (BS.drop (8 * i) bs)))+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) --- | Read n Float values directly into an unboxed vector.-readNFloatVec :: Int -> BS.ByteString -> VU.Vector Float-readNFloatVec n bs =- VU.generate- n- (\i -> castWord32ToFloat (littleEndianWord32 (BS.drop (4 * i) bs)))+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) --- | Read n Double values directly into an unboxed vector.-readNDoubleVec :: Int -> BS.ByteString -> VU.Vector Double-readNDoubleVec n bs =- VU.generate- n- (\i -> castWord64ToDouble (littleEndianWord64 (BS.drop (8 * i) bs)))+{- | 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) -readNInt32 :: Int -> BS.ByteString -> ([Int32], BS.ByteString)-readNInt32 0 bs = ([], bs)-readNInt32 k bs =- let x = littleEndianInt32 (BS.take 4 bs)- bs' = BS.drop 4 bs- (xs, rest) = readNInt32 (k - 1) bs'- in (x : xs, rest)+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) -readNDouble :: Int -> BS.ByteString -> ([Double], BS.ByteString)-readNDouble 0 bs = ([], bs)-readNDouble k bs =- let x = castWord64ToDouble (littleEndianWord64 (BS.take 8 bs))- bs' = BS.drop 8 bs- (xs, rest) = readNDouble (k - 1) bs'- in (x : xs, rest)+-- ---------------------------------------------------------------------------+-- Core page-iteration loop+-- --------------------------------------------------------------------------- -readNByteArrays :: Int -> BS.ByteString -> ([BS.ByteString], BS.ByteString)-readNByteArrays 0 bs = ([], bs)-readNByteArrays k bs =- let len = fromIntegral (littleEndianInt32 (BS.take 4 bs)) :: Int- body = BS.take len (BS.drop 4 bs)- bs' = BS.drop (4 + len) bs- (xs, rest) = readNByteArrays (k - 1) bs'- in (body : xs, rest)+-- | 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) -readNBool :: Int -> BS.ByteString -> ([Bool], BS.ByteString)-readNBool 0 bs = ([], bs)+{- | 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- chunk = BS.take totalBytes bs- rest = BS.drop totalBytes bs bits = concatMap (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7])- (BS.unpack chunk)- bools = take count bits- in (bools, rest)+ (BS.unpack (BS.take totalBytes bs))+ in take count bits -readNInt64 :: Int -> BS.ByteString -> ([Int64], BS.ByteString)-readNInt64 0 bs = ([], bs)-readNInt64 k bs =- let x = fromIntegral (littleEndianWord64 (BS.take 8 bs))- bs' = BS.drop 8 bs- (xs, rest) = readNInt64 (k - 1) bs'- in (x : xs, rest)+readNInt32 :: Int -> BS.ByteString -> VU.Vector Int32+readNInt32 n bs = VU.generate n $ \i -> littleEndianInt32 (BS.drop (4 * i) bs) -readNFloat :: Int -> BS.ByteString -> ([Float], BS.ByteString)-readNFloat 0 bs = ([], bs)-readNFloat k bs =- let x = castWord32ToFloat (littleEndianWord32 (BS.take 4 bs))- bs' = BS.drop 4 bs- (xs, rest) = readNFloat (k - 1) bs'- in (x : xs, rest)+readNInt64 :: Int -> BS.ByteString -> VU.Vector Int64+readNInt64 n bs = VU.generate n $ \i ->+ fromIntegral (littleEndianWord64 (BS.drop (8 * i) bs)) -splitFixed :: Int -> Int -> BS.ByteString -> ([BS.ByteString], BS.ByteString)-splitFixed 0 _ bs = ([], bs)-splitFixed k len bs =- let body = BS.take len bs- bs' = BS.drop len bs- (xs, rest) = splitFixed (k - 1) len bs'- in (body : xs, rest)+readNInt96 :: Int -> BS.ByteString -> [UTCTime]+readNInt96 0 _ = []+readNInt96 n bs = int96ToUTCTime (BS.take 12 bs) : readNInt96 (n - 1) (BS.drop 12 bs) -readStatisticsFromBytes ::- ColumnStatistics -> BS.ByteString -> Int16 -> (ColumnStatistics, BS.ByteString)-readStatisticsFromBytes cs xs lastFieldId =- let- fieldContents = readField' xs lastFieldId- in- case fieldContents of- Nothing -> (cs, BS.drop 1 xs)- Just (remainder, _elemType, identifier) -> case identifier of- 1 ->- let- (maxInBytes, remainder') = readByteStringFromBytes remainder- in- readStatisticsFromBytes (cs{columnMax = maxInBytes}) remainder' identifier- 2 ->- let- (minInBytes, remainder') = readByteStringFromBytes remainder- in- readStatisticsFromBytes (cs{columnMin = minInBytes}) remainder' identifier- 3 ->- let- (nullCount, remainder') = readIntFromBytes @Int64 remainder- in- readStatisticsFromBytes (cs{columnNullCount = nullCount}) remainder' identifier- 4 ->- let- (distinctCount, remainder') = readIntFromBytes @Int64 remainder- in- readStatisticsFromBytes- (cs{columnDistictCount = distinctCount})- remainder'- identifier- 5 ->- let- (maxInBytes, remainder') = readByteStringFromBytes remainder- in- readStatisticsFromBytes (cs{columnMaxValue = maxInBytes}) remainder' identifier- 6 ->- let- (minInBytes, remainder') = readByteStringFromBytes remainder- in- readStatisticsFromBytes (cs{columnMinValue = minInBytes}) remainder' identifier- 7 ->- case BS.uncons remainder of- Nothing ->- error "readStatisticsFromBytes: not enough bytes"- Just (isMaxValueExact, remainder') ->- readStatisticsFromBytes- (cs{isColumnMaxValueExact = isMaxValueExact == compactBooleanTrue})- remainder'- identifier- 8 ->- case BS.uncons remainder of- Nothing ->- error "readStatisticsFromBytes: not enough bytes"- Just (isMinValueExact, remainder') ->- readStatisticsFromBytes- (cs{isColumnMinValueExact = isMinValueExact == compactBooleanTrue})- remainder'- identifier- n -> error $ show n+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,86 @@+{-# 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)+import DataFrame.Operations.Core (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
@@ -16,11 +16,14 @@ 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@@ -131,6 +134,17 @@ 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
src/DataFrame/IO/Parquet/Thrift.hs view
@@ -1,1199 +1,584 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE StrictData #-}-{-# LANGUAGE TypeApplications #-}--module DataFrame.IO.Parquet.Thrift where--import Control.Monad-import Data.Bits-import qualified Data.ByteString as BS-import Data.Char-import Data.IORef-import Data.Int-import qualified Data.Map as M-import Data.Maybe-import qualified Data.Text as T-import Data.Typeable (Typeable)-import qualified Data.Vector as V-import qualified Data.Vector.Unboxed as VU-import Data.Word-import DataFrame.IO.Parquet.Binary-import DataFrame.IO.Parquet.Seeking-import DataFrame.IO.Parquet.Types-import qualified DataFrame.Internal.Column as DI-import DataFrame.Internal.DataFrame (DataFrame, unsafeGetColumn)-import qualified DataFrame.Operations.Core as DI-import Type.Reflection (- eqTypeRep,- typeRep,- (:~~:) (HRefl),- )--data SchemaElement = SchemaElement- { elementName :: T.Text- , elementType :: TType- , typeLength :: Int32- , numChildren :: Int32- , fieldId :: Int32- , repetitionType :: RepetitionType- , convertedType :: Int32- , scale :: Int32- , precision :: Int32- , logicalType :: LogicalType- }- deriving (Show, Eq)--createParquetSchema :: DataFrame -> [SchemaElement]-createParquetSchema df = schemaDef : map toSchemaElement (DI.columnNames df)- where- -- The schema always contains an initial element- -- indicating the group of fields.- schemaDef =- SchemaElement- { elementName = "schema"- , elementType = STOP- , typeLength = 0- , numChildren = fromIntegral (snd (DI.dimensions df))- , fieldId = -1- , repetitionType = UNKNOWN_REPETITION_TYPE- , convertedType = 0- , scale = 0- , precision = 0- , logicalType = LOGICAL_TYPE_UNKNOWN- }- toSchemaElement colName =- let- colType :: TType- colType = case unsafeGetColumn colName df of- (DI.BoxedColumn _ (_col :: V.Vector a)) -> haskellToTType @a- (DI.UnboxedColumn _ (_col :: VU.Vector a)) -> haskellToTType @a- lType =- if DI.hasElemType @T.Text (unsafeGetColumn colName df)- || DI.hasElemType @(Maybe T.Text) (unsafeGetColumn colName df)- then STRING_TYPE- else LOGICAL_TYPE_UNKNOWN- in- SchemaElement colName colType 0 0 (-1) OPTIONAL 0 0 0 lType--data KeyValue = KeyValue- { key :: String- , value :: String- }- deriving (Show, Eq)--data FileMetadata = FileMetaData- { version :: Int32- , schema :: [SchemaElement]- , numRows :: Integer- , rowGroups :: [RowGroup]- , keyValueMetadata :: [KeyValue]- , createdBy :: Maybe String- , columnOrders :: [ColumnOrder]- , encryptionAlgorithm :: EncryptionAlgorithm- , footerSigningKeyMetadata :: BS.ByteString- }- deriving (Show, Eq)--data TType- = STOP- | BOOL- | BYTE- | I16- | I32- | I64- | I96- | FLOAT- | DOUBLE- | STRING- | LIST- | SET- | MAP- | STRUCT- | UUID- deriving (Show, Eq)--haskellToTType :: forall a. (Typeable a) => TType-haskellToTType- | is @Bool = BOOL- | is @Int8 = BYTE- | is @Word8 = BYTE- | is @Int16 = I16- | is @Word16 = I16- | is @Int32 = I32- | is @Word32 = I32- | is @Int64 = I64- | is @Word64 = I64- | is @Float = FLOAT- | is @Double = DOUBLE- | is @String = STRING- | is @T.Text = STRING- | is @BS.ByteString = STRING- | otherwise = STOP- where- is :: forall x. (Typeable x) => Bool- is = case eqTypeRep (typeRep @a) (typeRep @x) of- Just HRefl -> True- Nothing -> False--defaultMetadata :: FileMetadata-defaultMetadata =- FileMetaData- { version = 0- , schema = []- , numRows = 0- , rowGroups = []- , keyValueMetadata = []- , createdBy = Nothing- , columnOrders = []- , encryptionAlgorithm = ENCRYPTION_ALGORITHM_UNKNOWN- , footerSigningKeyMetadata = BS.empty- }--data ColumnMetaData = ColumnMetaData- { columnType :: ParquetType- , columnEncodings :: [ParquetEncoding]- , columnPathInSchema :: [String]- , columnCodec :: CompressionCodec- , columnNumValues :: Int64- , columnTotalUncompressedSize :: Int64- , columnTotalCompressedSize :: Int64- , columnKeyValueMetadata :: [KeyValue]- , columnDataPageOffset :: Int64- , columnIndexPageOffset :: Int64- , columnDictionaryPageOffset :: Int64- , columnStatistics :: ColumnStatistics- , columnEncodingStats :: [PageEncodingStats]- , bloomFilterOffset :: Int64- , bloomFilterLength :: Int32- , columnSizeStatistics :: SizeStatistics- , columnGeospatialStatistics :: GeospatialStatistics- }- deriving (Show, Eq)--data ColumnChunk = ColumnChunk- { columnChunkFilePath :: String- , columnChunkMetadataFileOffset :: Int64- , columnMetaData :: ColumnMetaData- , columnChunkOffsetIndexOffset :: Int64- , columnChunkOffsetIndexLength :: Int32- , columnChunkColumnIndexOffset :: Int64- , columnChunkColumnIndexLength :: Int32- , cryptoMetadata :: ColumnCryptoMetadata- , encryptedColumnMetadata :: BS.ByteString- }- deriving (Show, Eq)--data RowGroup = RowGroup- { rowGroupColumns :: [ColumnChunk]- , totalByteSize :: Int64- , rowGroupNumRows :: Int64- , rowGroupSortingColumns :: [SortingColumn]- , fileOffset :: Int64- , totalCompressedSize :: Int64- , ordinal :: Int16- }- deriving (Show, Eq)--defaultSchemaElement :: SchemaElement-defaultSchemaElement =- SchemaElement- ""- STOP- 0- 0- (-1)- UNKNOWN_REPETITION_TYPE- 0- 0- 0- LOGICAL_TYPE_UNKNOWN--emptyColumnMetadata :: ColumnMetaData-emptyColumnMetadata =- ColumnMetaData- PARQUET_TYPE_UNKNOWN- []- []- COMPRESSION_CODEC_UNKNOWN- 0- 0- 0- []- 0- 0- 0- emptyColumnStatistics- []- 0- 0- emptySizeStatistics- emptyGeospatialStatistics--emptyColumnChunk :: ColumnChunk-emptyColumnChunk =- ColumnChunk- ""- 0- emptyColumnMetadata- 0- 0- 0- 0- COLUMN_CRYPTO_METADATA_UNKNOWN- BS.empty--emptyKeyValue :: KeyValue-emptyKeyValue = KeyValue{key = "", value = ""}--emptyRowGroup :: RowGroup-emptyRowGroup = RowGroup [] 0 0 [] 0 0 0--compactBooleanTrue- , compactI32- , compactI64- , compactDouble- , compactBinary- , compactList- , compactStruct ::- Word8-compactBooleanTrue = 0x01-compactI32 = 0x05-compactI64 = 0x06-compactDouble = 0x07-compactBinary = 0x08-compactList = 0x09-compactStruct = 0x0C--toTType :: Word8 -> TType-toTType t =- fromMaybe STOP $- M.lookup (t .&. 0x0f) $- M.fromList- [ (compactBooleanTrue, BOOL)- , (compactI32, I32)- , (compactI64, I64)- , (compactDouble, DOUBLE)- , (compactBinary, STRING)- , (compactList, LIST)- , (compactStruct, STRUCT)- ]--readField ::- BS.ByteString -> IORef Int -> Int16 -> IO (Maybe (TType, Int16))-readField buf pos lastFieldId = do- t <- readAndAdvance pos buf- if t .&. 0x0f == 0- then return Nothing- else do- let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16- identifier <-- if modifier == 0- then readIntFromBuffer @Int16 buf pos- else return (lastFieldId + modifier)- let elemType = toTType (t .&. 0x0f)- pure $ Just (elemType, identifier)--skipToStructEnd :: BS.ByteString -> IORef Int -> IO ()-skipToStructEnd buf pos = do- t <- readAndAdvance pos buf- if t .&. 0x0f == 0- then return ()- else do- let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16- _identifier <-- if modifier == 0- then readIntFromBuffer @Int16 buf pos- else return 0- let elemType = toTType (t .&. 0x0f)- skipFieldData elemType buf pos- skipToStructEnd buf pos--skipFieldData :: TType -> BS.ByteString -> IORef Int -> IO ()-skipFieldData fieldType buf pos = case fieldType of- BOOL -> return ()- I32 -> void (readIntFromBuffer @Int32 buf pos)- I64 -> void (readIntFromBuffer @Int64 buf pos)- DOUBLE -> void (readIntFromBuffer @Int64 buf pos)- STRING -> void (readByteString buf pos)- LIST -> skipList buf pos- STRUCT -> skipToStructEnd buf pos- _ -> error $ "Unknown field type" ++ show fieldType--skipList :: BS.ByteString -> IORef Int -> IO ()-skipList buf pos = do- sizeAndType <- readAndAdvance pos buf- let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int- let elemType = toTType sizeAndType- replicateM_ sizeOnly (skipFieldData elemType buf pos)--{- | This avoids reading entire bytestring at once: it uses the seekable handle- seeks it to the end of the file to read the metadata--}-readMetadataByHandleMetaSize :: FileBufferedOrSeekable -> Int -> IO FileMetadata-readMetadataByHandleMetaSize sh metaSize = do- let lastFieldId = 0- bs <- readLastBytes (fromIntegral $ metaSize + footerSize) sh- bufferPos <- newIORef 0- readFileMetaData defaultMetadata bs bufferPos lastFieldId---- | metadata starts from (L - 8 - meta_size) to L - 8 - 1.-readMetadata :: BS.ByteString -> Int -> IO FileMetadata-readMetadata contents size = do- let metadataStartPos = BS.length contents - footerSize - size- let metadataBytes =- BS.pack $- map (BS.index contents) [metadataStartPos .. (metadataStartPos + size - 1)]- let lastFieldId = 0- bufferPos <- newIORef (0 :: Int)- readFileMetaData defaultMetadata metadataBytes bufferPos lastFieldId--readFileMetaData ::- FileMetadata ->- BS.ByteString ->- IORef Int ->- Int16 ->- IO FileMetadata-readFileMetaData metadata metaDataBuf bufferPos lastFieldId = do- fieldContents <- readField metaDataBuf bufferPos lastFieldId- case fieldContents of- Nothing -> return metadata- Just (_elemType, identifier) -> case identifier of- 1 -> do- parsedVersion <- readIntFromBuffer @Int32 metaDataBuf bufferPos- readFileMetaData- (metadata{version = parsedVersion})- metaDataBuf- bufferPos- identifier- 2 -> do- sizeAndType <- readAndAdvance bufferPos metaDataBuf- listSize <-- if (sizeAndType `shiftR` 4) .&. 0x0f == 15- then readVarIntFromBuffer @Int metaDataBuf bufferPos- else return $ fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f)- let _elemType = toTType sizeAndType- schemaElements <-- replicateM- listSize- (readSchemaElement defaultSchemaElement metaDataBuf bufferPos 0)- readFileMetaData- (metadata{schema = schemaElements})- metaDataBuf- bufferPos- identifier- 3 -> do- parsedNumRows <- readIntFromBuffer @Int64 metaDataBuf bufferPos- readFileMetaData- (metadata{numRows = fromIntegral parsedNumRows})- metaDataBuf- bufferPos- identifier- 4 -> do- sizeAndType <- readAndAdvance bufferPos metaDataBuf- listSize <-- if (sizeAndType `shiftR` 4) .&. 0x0f == 15- then readVarIntFromBuffer @Int metaDataBuf bufferPos- else return $ fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f)-- -- TODO actually check elemType agrees (also for all the other underscored _elemType in this module)- let _elemType = toTType sizeAndType- parsedRowGroups <-- replicateM listSize (readRowGroup emptyRowGroup metaDataBuf bufferPos 0)- readFileMetaData- (metadata{rowGroups = parsedRowGroups})- metaDataBuf- bufferPos- identifier- 5 -> do- sizeAndType <- readAndAdvance bufferPos metaDataBuf- listSize <-- if (sizeAndType `shiftR` 4) .&. 0x0f == 15- then readVarIntFromBuffer @Int metaDataBuf bufferPos- else return $ fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f)-- let _elemType = toTType sizeAndType- parsedKeyValueMetadata <-- replicateM listSize (readKeyValue emptyKeyValue metaDataBuf bufferPos 0)- readFileMetaData- (metadata{keyValueMetadata = parsedKeyValueMetadata})- metaDataBuf- bufferPos- identifier- 6 -> do- parsedCreatedBy <- readString metaDataBuf bufferPos- readFileMetaData- (metadata{createdBy = Just parsedCreatedBy})- metaDataBuf- bufferPos- identifier- 7 -> do- sizeAndType <- readAndAdvance bufferPos metaDataBuf- listSize <-- if (sizeAndType `shiftR` 4) .&. 0x0f == 15- then readVarIntFromBuffer @Int metaDataBuf bufferPos- else return $ fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f)-- let _elemType = toTType sizeAndType- parsedColumnOrders <-- replicateM listSize (readColumnOrder metaDataBuf bufferPos 0)- readFileMetaData- (metadata{columnOrders = parsedColumnOrders})- metaDataBuf- bufferPos- identifier- 8 -> do- parsedEncryptionAlgorithm <- readEncryptionAlgorithm metaDataBuf bufferPos 0- readFileMetaData- (metadata{encryptionAlgorithm = parsedEncryptionAlgorithm})- metaDataBuf- bufferPos- identifier- 9 -> do- parsedFooterSigningKeyMetadata <- readByteString metaDataBuf bufferPos- readFileMetaData- (metadata{footerSigningKeyMetadata = parsedFooterSigningKeyMetadata})- metaDataBuf- bufferPos- identifier- n -> return $ error $ "UNIMPLEMENTED " ++ show n--readSchemaElement ::- SchemaElement ->- BS.ByteString ->- IORef Int ->- Int16 ->- IO SchemaElement-readSchemaElement schemaElement buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return schemaElement- Just (_elemType, identifier) -> case identifier of- 1 -> do- schemaElemType <- toIntegralType <$> readInt32FromBuffer buf pos- readSchemaElement- (schemaElement{elementType = schemaElemType})- buf- pos- identifier- 2 -> do- parsedTypeLength <- readInt32FromBuffer buf pos- readSchemaElement- (schemaElement{typeLength = parsedTypeLength})- buf- pos- identifier- 3 -> do- fieldRepetitionType <- readInt32FromBuffer buf pos- readSchemaElement- (schemaElement{repetitionType = repetitionTypeFromInt fieldRepetitionType})- buf- pos- identifier- 4 -> do- nameSize <- readVarIntFromBuffer @Int buf pos- if nameSize <= 0- then readSchemaElement schemaElement buf pos identifier- else do- contents <- replicateM nameSize (readAndAdvance pos buf)- readSchemaElement- (schemaElement{elementName = T.pack (map (chr . fromIntegral) contents)})- buf- pos- identifier- 5 -> do- parsedNumChildren <- readInt32FromBuffer buf pos- readSchemaElement- (schemaElement{numChildren = parsedNumChildren})- buf- pos- identifier- 6 -> do- parsedConvertedType <- readInt32FromBuffer buf pos- readSchemaElement- (schemaElement{convertedType = parsedConvertedType})- buf- pos- identifier- 7 -> do- parsedScale <- readInt32FromBuffer buf pos- readSchemaElement (schemaElement{scale = parsedScale}) buf pos identifier- 8 -> do- parsedPrecision <- readInt32FromBuffer buf pos- readSchemaElement- (schemaElement{precision = parsedPrecision})- buf- pos- identifier- 9 -> do- parsedFieldId <- readInt32FromBuffer buf pos- readSchemaElement- (schemaElement{fieldId = parsedFieldId})- buf- pos- identifier- 10 -> do- parsedLogicalType <- readLogicalType LOGICAL_TYPE_UNKNOWN buf pos 0- readSchemaElement- (schemaElement{logicalType = parsedLogicalType})- buf- pos- identifier- n -> error ("Uknown schema element: " ++ show n)--readRowGroup ::- RowGroup -> BS.ByteString -> IORef Int -> Int16 -> IO RowGroup-readRowGroup r buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return r- Just (_elemType, identifier) -> case identifier of- 1 -> do- sizeAndType <- readAndAdvance pos buf- listSize <-- if (sizeAndType `shiftR` 4) .&. 0x0f == 15- then readVarIntFromBuffer @Int buf pos- else return $ fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f)- let _elemType = toTType sizeAndType- columnChunks <-- replicateM listSize (readColumnChunk emptyColumnChunk buf pos 0)- readRowGroup (r{rowGroupColumns = columnChunks}) buf pos identifier- 2 -> do- totalBytes <- readIntFromBuffer @Int64 buf pos- readRowGroup (r{totalByteSize = totalBytes}) buf pos identifier- 3 -> do- nRows <- readIntFromBuffer @Int64 buf pos- readRowGroup (r{rowGroupNumRows = nRows}) buf pos identifier- 4 -> return r- 5 -> do- offset <- readIntFromBuffer @Int64 buf pos- readRowGroup (r{fileOffset = offset}) buf pos identifier- 6 -> do- compressedSize <- readIntFromBuffer @Int64 buf pos- readRowGroup- (r{totalCompressedSize = compressedSize})- buf- pos- identifier- 7 -> do- parsedOrdinal <- readIntFromBuffer @Int16 buf pos- readRowGroup (r{ordinal = parsedOrdinal}) buf pos identifier- _ -> error $ "Unknown row group field: " ++ show identifier--readColumnChunk ::- ColumnChunk -> BS.ByteString -> IORef Int -> Int16 -> IO ColumnChunk-readColumnChunk c buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return c- Just (_elemType, identifier) -> case identifier of- 1 -> do- stringSize <- readVarIntFromBuffer @Int buf pos- contents <-- map (chr . fromIntegral) <$> replicateM stringSize (readAndAdvance pos buf)- readColumnChunk- (c{columnChunkFilePath = contents})- buf- pos- identifier- 2 -> do- parsedMetadataFileOffset <- readIntFromBuffer @Int64 buf pos- readColumnChunk- (c{columnChunkMetadataFileOffset = parsedMetadataFileOffset})- buf- pos- identifier- 3 -> do- columnMetadata <- readColumnMetadata emptyColumnMetadata buf pos 0- readColumnChunk- (c{columnMetaData = columnMetadata})- buf- pos- identifier- 4 -> do- columnOffsetIndexOffset <- readIntFromBuffer @Int64 buf pos- readColumnChunk- (c{columnChunkOffsetIndexOffset = columnOffsetIndexOffset})- buf- pos- identifier- 5 -> do- columnOffsetIndexLength <- readInt32FromBuffer buf pos- readColumnChunk- (c{columnChunkOffsetIndexLength = columnOffsetIndexLength})- buf- pos- identifier- 6 -> do- parsedColumnIndexOffset <- readIntFromBuffer @Int64 buf pos- readColumnChunk- (c{columnChunkColumnIndexOffset = parsedColumnIndexOffset})- buf- pos- identifier- 7 -> do- parsedColumnIndexLength <- readInt32FromBuffer buf pos- readColumnChunk- (c{columnChunkColumnIndexLength = parsedColumnIndexLength})- buf- pos- identifier- _ -> error "Unknown column chunk"--readColumnMetadata ::- ColumnMetaData ->- BS.ByteString ->- IORef Int ->- Int16 ->- IO ColumnMetaData-readColumnMetadata cm buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return cm- Just (_elemType, identifier) -> case identifier of- 1 -> do- cType <- parquetTypeFromInt <$> readInt32FromBuffer buf pos- readColumnMetadata (cm{columnType = cType}) buf pos identifier- 2 -> do- sizeAndType <- readAndAdvance pos buf- let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int- let _elemType = toTType sizeAndType- encodings <- replicateM sizeOnly (readParquetEncoding buf pos 0)- readColumnMetadata- (cm{columnEncodings = encodings})- buf- pos- identifier- 3 -> do- sizeAndType <- readAndAdvance pos buf- let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int- let _elemType = toTType sizeAndType- paths <- replicateM sizeOnly (readString buf pos)- readColumnMetadata- (cm{columnPathInSchema = paths})- buf- pos- identifier- 4 -> do- cType <- compressionCodecFromInt <$> readInt32FromBuffer buf pos- readColumnMetadata (cm{columnCodec = cType}) buf pos identifier- 5 -> do- numValues <- readIntFromBuffer @Int64 buf pos- readColumnMetadata (cm{columnNumValues = numValues}) buf pos identifier- 6 -> do- parsedTotalUncompressedSize <- readIntFromBuffer @Int64 buf pos- readColumnMetadata- (cm{columnTotalUncompressedSize = parsedTotalUncompressedSize})- buf- pos- identifier- 7 -> do- parsedTotalCompressedSize <- readIntFromBuffer @Int64 buf pos- readColumnMetadata- (cm{columnTotalCompressedSize = parsedTotalCompressedSize})- buf- pos- identifier- 8 -> do- sizeAndType <- readAndAdvance pos buf- let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int- let _elemType = toTType sizeAndType- parsedKeyValueMeta <-- replicateM sizeOnly (readKeyValue emptyKeyValue buf pos 0)- readColumnMetadata- (cm{columnKeyValueMetadata = parsedKeyValueMeta})- buf- pos- identifier- 9 -> do- parsedDataPageOffset <- readIntFromBuffer @Int64 buf pos- readColumnMetadata- (cm{columnDataPageOffset = parsedDataPageOffset})- buf- pos- identifier- 10 -> do- parsedIndexPageOffset <- readIntFromBuffer @Int64 buf pos- readColumnMetadata- (cm{columnIndexPageOffset = parsedIndexPageOffset})- buf- pos- identifier- 11 -> do- parsedDictionaryPageOffset <- readIntFromBuffer @Int64 buf pos- readColumnMetadata- (cm{columnDictionaryPageOffset = parsedDictionaryPageOffset})- buf- pos- identifier- 12 -> do- stats <- readStatistics emptyColumnStatistics buf pos 0- readColumnMetadata (cm{columnStatistics = stats}) buf pos identifier- 13 -> do- sizeAndType <- readAndAdvance pos buf- let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int- let _elemType = toTType sizeAndType- pageEncodingStats <-- replicateM sizeOnly (readPageEncodingStats emptyPageEncodingStats buf pos 0)- readColumnMetadata- (cm{columnEncodingStats = pageEncodingStats})- buf- pos- identifier- 14 -> do- parsedBloomFilterOffset <- readIntFromBuffer @Int64 buf pos- readColumnMetadata- (cm{bloomFilterOffset = parsedBloomFilterOffset})- buf- pos- identifier- 15 -> do- parsedBloomFilterLength <- readInt32FromBuffer buf pos- readColumnMetadata- (cm{bloomFilterLength = parsedBloomFilterLength})- buf- pos- identifier- 16 -> do- stats <- readSizeStatistics emptySizeStatistics buf pos 0- readColumnMetadata- (cm{columnSizeStatistics = stats})- buf- pos- identifier- 17 -> return $ error "UNIMPLEMENTED"- _ -> error $ "Unknown column metadata " ++ show identifier--readEncryptionAlgorithm ::- BS.ByteString -> IORef Int -> Int16 -> IO EncryptionAlgorithm-readEncryptionAlgorithm buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return ENCRYPTION_ALGORITHM_UNKNOWN- Just (_elemType, identifier) -> case identifier of- 1 -> do- readAesGcmV1- ( AesGcmV1- { aadPrefix = BS.empty- , aadFileUnique = BS.empty- , supplyAadPrefix = False- }- )- buf- pos- 0- 2 -> do- readAesGcmCtrV1- ( AesGcmCtrV1- { aadPrefix = BS.empty- , aadFileUnique = BS.empty- , supplyAadPrefix = False- }- )- buf- pos- 0- _n -> return ENCRYPTION_ALGORITHM_UNKNOWN--readColumnOrder ::- BS.ByteString -> IORef Int -> Int16 -> IO ColumnOrder-readColumnOrder buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return COLUMN_ORDER_UNKNOWN- Just (_elemType, identifier) -> case identifier of- 1 -> do- -- Read begin struct and stop since this an empty struct.- replicateM_ 2 (readTypeOrder buf pos 0)- return TYPE_ORDER- _ -> return COLUMN_ORDER_UNKNOWN--readAesGcmCtrV1 ::- EncryptionAlgorithm ->- BS.ByteString ->- IORef Int ->- Int16 ->- IO EncryptionAlgorithm-readAesGcmCtrV1 v@(AesGcmCtrV1 _aadPrefix _aadFileUnique _supplyAadPrefix) buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return v- Just (_elemType, identifier) -> case identifier of- 1 -> do- parsedAadPrefix <- readByteString buf pos- readAesGcmCtrV1 (v{aadPrefix = parsedAadPrefix}) buf pos identifier- 2 -> do- parsedAadFileUnique <- readByteString buf pos- readAesGcmCtrV1- (v{aadFileUnique = parsedAadFileUnique})- buf- pos- identifier- 3 -> do- parsedSupplyAadPrefix <- readAndAdvance pos buf- readAesGcmCtrV1- (v{supplyAadPrefix = parsedSupplyAadPrefix == compactBooleanTrue})- buf- pos- identifier- _ -> return ENCRYPTION_ALGORITHM_UNKNOWN-readAesGcmCtrV1 _ _ _ _ =- error "readAesGcmCtrV1 called with non AesGcmCtrV1"--readAesGcmV1 ::- EncryptionAlgorithm ->- BS.ByteString ->- IORef Int ->- Int16 ->- IO EncryptionAlgorithm-readAesGcmV1 v@(AesGcmV1 _aadPrefix _aadFileUnique _supplyAadPrefix) buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return v- Just (_elemType, identifier) -> case identifier of- 1 -> do- parsedAadPrefix <- readByteString buf pos- readAesGcmV1 (v{aadPrefix = parsedAadPrefix}) buf pos identifier- 2 -> do- parsedAadFileUnique <- readByteString buf pos- readAesGcmV1 (v{aadFileUnique = parsedAadFileUnique}) buf pos identifier- 3 -> do- parsedSupplyAadPrefix <- readAndAdvance pos buf- readAesGcmV1- (v{supplyAadPrefix = parsedSupplyAadPrefix == compactBooleanTrue})- buf- pos- identifier- _ -> return ENCRYPTION_ALGORITHM_UNKNOWN-readAesGcmV1 _ _ _ _ =- error "readAesGcmV1 called with non AesGcmV1"--readTypeOrder ::- BS.ByteString -> IORef Int -> Int16 -> IO ColumnOrder-readTypeOrder buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return TYPE_ORDER- Just (elemType, identifier) ->- if elemType == STOP- then return TYPE_ORDER- else readTypeOrder buf pos identifier--readKeyValue ::- KeyValue -> BS.ByteString -> IORef Int -> Int16 -> IO KeyValue-readKeyValue kv buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return kv- Just (_elemType, identifier) -> case identifier of- 1 -> do- k <- readString buf pos- readKeyValue (kv{key = k}) buf pos identifier- 2 -> do- v <- readString buf pos- readKeyValue (kv{value = v}) buf pos identifier- _ -> error "Unknown kv"--readPageEncodingStats ::- PageEncodingStats ->- BS.ByteString ->- IORef Int ->- Int16 ->- IO PageEncodingStats-readPageEncodingStats pes buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return pes- Just (_elemType, identifier) -> case identifier of- 1 -> do- pType <- pageTypeFromInt <$> readInt32FromBuffer buf pos- readPageEncodingStats (pes{pageEncodingPageType = pType}) buf pos identifier- 2 -> do- pEnc <- parquetEncodingFromInt <$> readInt32FromBuffer buf pos- readPageEncodingStats (pes{pageEncoding = pEnc}) buf pos identifier- 3 -> do- encodedCount <- readInt32FromBuffer buf pos- readPageEncodingStats- (pes{pagesWithEncoding = encodedCount})- buf- pos- identifier- _ -> error "Unknown page encoding stats"--readParquetEncoding ::- BS.ByteString -> IORef Int -> Int16 -> IO ParquetEncoding-readParquetEncoding buf pos _lastFieldId = parquetEncodingFromInt <$> readInt32FromBuffer buf pos--readStatistics ::- ColumnStatistics ->- BS.ByteString ->- IORef Int ->- Int16 ->- IO ColumnStatistics-readStatistics cs buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return cs- Just (_elemType, identifier) -> case identifier of- 1 -> do- maxInBytes <- readByteString buf pos- readStatistics (cs{columnMax = maxInBytes}) buf pos identifier- 2 -> do- minInBytes <- readByteString buf pos- readStatistics (cs{columnMin = minInBytes}) buf pos identifier- 3 -> do- nullCount <- readIntFromBuffer @Int64 buf pos- readStatistics (cs{columnNullCount = nullCount}) buf pos identifier- 4 -> do- distinctCount <- readIntFromBuffer @Int64 buf pos- readStatistics- (cs{columnDistictCount = distinctCount})- buf- pos- identifier- 5 -> do- maxInBytes <- readByteString buf pos- readStatistics (cs{columnMaxValue = maxInBytes}) buf pos identifier- 6 -> do- minInBytes <- readByteString buf pos- readStatistics (cs{columnMinValue = minInBytes}) buf pos identifier- 7 -> do- isMaxValueExact <- readAndAdvance pos buf- readStatistics- (cs{isColumnMaxValueExact = isMaxValueExact == compactBooleanTrue})- buf- pos- identifier- 8 -> do- isMinValueExact <- readAndAdvance pos buf- readStatistics- (cs{isColumnMinValueExact = isMinValueExact == compactBooleanTrue})- buf- pos- identifier- _ -> error "Unknown statistics"--readSizeStatistics ::- SizeStatistics ->- BS.ByteString ->- IORef Int ->- Int16 ->- IO SizeStatistics-readSizeStatistics ss buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return ss- Just (_elemType, identifier) -> case identifier of- 1 -> do- parsedUnencodedByteArrayDataTypes <- readIntFromBuffer @Int64 buf pos- readSizeStatistics- (ss{unencodedByteArrayDataTypes = parsedUnencodedByteArrayDataTypes})- buf- pos- identifier- 2 -> do- sizeAndType <- readAndAdvance pos buf- let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int- let _elemType = toTType sizeAndType- parsedRepetitionLevelHistogram <-- replicateM sizeOnly (readIntFromBuffer @Int64 buf pos)- readSizeStatistics- (ss{repetitionLevelHistogram = parsedRepetitionLevelHistogram})- buf- pos- identifier- 3 -> do- sizeAndType <- readAndAdvance pos buf- let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int- let _elemType = toTType sizeAndType- parsedDefinitionLevelHistogram <-- replicateM sizeOnly (readIntFromBuffer @Int64 buf pos)- readSizeStatistics- (ss{definitionLevelHistogram = parsedDefinitionLevelHistogram})- buf- pos- identifier- _ -> error "Unknown size statistics"--footerSize :: Int-footerSize = 8--toIntegralType :: Int32 -> TType-toIntegralType n- | n == 0 = BOOL- | n == 1 = I32- | n == 2 = I64- | n == 3 = I96- | n == 4 = FLOAT- | n == 5 = DOUBLE- | n == 6 = STRING- | n == 7 = STRING- | otherwise = error ("Unknown type in schema: " ++ show n)--readLogicalType ::- LogicalType -> BS.ByteString -> IORef Int -> Int16 -> IO LogicalType-readLogicalType parsedLogicalType buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> pure parsedLogicalType- Just (_elemType, identifier) -> case identifier of- 1 -> do- -- This is an empty enum and is read as a field.- _ <- readField buf pos 0- readLogicalType STRING_TYPE buf pos identifier- 2 -> do- _ <- readField buf pos 0- readLogicalType MAP_TYPE buf pos identifier- 3 -> do- _ <- readField buf pos 0- readLogicalType LIST_TYPE buf pos identifier- 4 -> do- _ <- readField buf pos 0- readLogicalType ENUM_TYPE buf pos identifier- 5 -> do- decimal <- readDecimalType 0 0 buf pos 0- readLogicalType decimal buf pos identifier- 6 -> do- _ <- readField buf pos 0- readLogicalType DATE_TYPE buf pos identifier- 7 -> do- time <- readTimeType False MILLISECONDS buf pos 0- readLogicalType time buf pos identifier- 8 -> do- timestamp <- readTimestampType False MILLISECONDS buf pos 0- readLogicalType timestamp buf pos identifier- -- Apparently reserved for interval types- 9 -> do- _ <- readField buf pos 0- readLogicalType LOGICAL_TYPE_UNKNOWN buf pos identifier- 10 -> do- intType <- readIntType 0 False buf pos 0- readLogicalType intType buf pos identifier- 11 -> do- _ <- readField buf pos 0- readLogicalType LOGICAL_TYPE_UNKNOWN buf pos identifier- 12 -> do- _ <- readField buf pos 0- readLogicalType JSON_TYPE buf pos identifier- 13 -> do- _ <- readField buf pos 0- readLogicalType BSON_TYPE buf pos identifier- 14 -> do- _ <- readField buf pos 0- readLogicalType UUID_TYPE buf pos identifier- 15 -> do- _ <- readField buf pos 0- readLogicalType FLOAT16_TYPE buf pos identifier- 16 -> error "Variant fields are unsupported"- 17 -> error "Geometry fields are unsupported"- 18 -> error "Geography fields are unsupported"- n -> error $ "Unknown logical type field: " ++ show n--readIntType ::- Int8 ->- Bool ->- BS.ByteString ->- IORef Int ->- Int16 ->- IO LogicalType-readIntType parsedBitWidth parsedIntIsSigned buf pos lastFieldId = do- t <- readAndAdvance pos buf- if t .&. 0x0f == 0- then return (IntType parsedBitWidth parsedIntIsSigned)- else do- let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16- identifier <-- if modifier == 0- then readIntFromBuffer @Int16 buf pos- else return (lastFieldId + modifier)-- case identifier of- 1 -> do- bitWidth' <- readAndAdvance pos buf- readIntType (fromIntegral bitWidth') parsedIntIsSigned buf pos identifier- 2 -> do- let intIsSigned' = (t .&. 0x0f) == compactBooleanTrue- readIntType parsedBitWidth intIsSigned' buf pos identifier- _ -> error $ "UNKNOWN field ID for IntType: " ++ show identifier--readDecimalType ::- Int32 ->- Int32 ->- BS.ByteString ->- IORef Int ->- Int16 ->- IO LogicalType-readDecimalType parsedPrecision parsedScale buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return (DecimalType parsedPrecision parsedScale)- Just (_elemType, identifier) -> case identifier of- 1 -> do- scale' <- readInt32FromBuffer buf pos- readDecimalType parsedPrecision scale' buf pos identifier- 2 -> do- precision' <- readInt32FromBuffer buf pos- readDecimalType precision' parsedScale buf pos identifier- _ -> error $ "UNKNOWN field ID for DecimalType" ++ show identifier--readTimeType ::- Bool ->- TimeUnit ->- BS.ByteString ->- IORef Int ->- Int16 ->- IO LogicalType-readTimeType parsedIsAdjustedToUTC parsedUnit buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing ->- return (TimeType{isAdjustedToUTC = parsedIsAdjustedToUTC, unit = parsedUnit})- Just (elemType, identifier) -> case identifier of- 1 -> do- let isAdjustedToUTC' = elemType == toTType compactBooleanTrue- readTimeType isAdjustedToUTC' parsedUnit buf pos identifier- 2 -> do- unit' <- readUnit TIME_UNIT_UNKNOWN buf pos 0- readTimeType parsedIsAdjustedToUTC unit' buf pos identifier- _ -> error $ "UNKNOWN field ID for TimeType" ++ show identifier--readTimestampType ::- Bool ->- TimeUnit ->- BS.ByteString ->- IORef Int ->- Int16 ->- IO LogicalType-readTimestampType parsedIsAdjustedToUTC parsedUnit buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing ->- return- (TimestampType{isAdjustedToUTC = parsedIsAdjustedToUTC, unit = parsedUnit})- Just (elemType, identifier) -> case identifier of- 1 -> do- let isAdjustedToUTC' = elemType == toTType compactBooleanTrue- readTimestampType isAdjustedToUTC' parsedUnit buf pos identifier- 2 -> do- unit' <- readUnit TIME_UNIT_UNKNOWN buf pos 0- readTimestampType parsedIsAdjustedToUTC unit' buf pos identifier- _ -> error $ "UNKNOWN field ID for TimestampType " ++ show identifier--readUnit :: TimeUnit -> BS.ByteString -> IORef Int -> Int16 -> IO TimeUnit-readUnit parsedUnit buf pos lastFieldId = do- fieldContents <- readField buf pos lastFieldId- case fieldContents of- Nothing -> return parsedUnit- Just (_elemType, identifier) -> case identifier of- 1 -> do- _ <- readField buf pos 0- readUnit MILLISECONDS buf pos identifier- 2 -> do- _ <- readField buf pos 0- readUnit MICROSECONDS buf pos identifier- 3 -> do- _ <- readField buf pos 0- readUnit NANOSECONDS buf pos identifier- n -> error $ "Unknown time unit: " ++ show n+{-# 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/Types.hs
@@ -1,314 +0,0 @@-module DataFrame.IO.Parquet.Types where--import qualified Data.ByteString as BS-import Data.Int-import qualified Data.Text as T-import Data.Time-import qualified Data.Vector as V--data ParquetType- = PBOOLEAN- | PINT32- | PINT64- | PINT96- | PFLOAT- | PDOUBLE- | PBYTE_ARRAY- | PFIXED_LEN_BYTE_ARRAY- | PARQUET_TYPE_UNKNOWN- deriving (Show, Eq)--parquetTypeFromInt :: Int32 -> ParquetType-parquetTypeFromInt 0 = PBOOLEAN-parquetTypeFromInt 1 = PINT32-parquetTypeFromInt 2 = PINT64-parquetTypeFromInt 3 = PINT96-parquetTypeFromInt 4 = PFLOAT-parquetTypeFromInt 5 = PDOUBLE-parquetTypeFromInt 6 = PBYTE_ARRAY-parquetTypeFromInt 7 = PFIXED_LEN_BYTE_ARRAY-parquetTypeFromInt _ = PARQUET_TYPE_UNKNOWN--data PageType- = DATA_PAGE- | INDEX_PAGE- | DICTIONARY_PAGE- | DATA_PAGE_V2- | PAGE_TYPE_UNKNOWN- deriving (Show, Eq)--pageTypeFromInt :: Int32 -> PageType-pageTypeFromInt 0 = DATA_PAGE-pageTypeFromInt 1 = INDEX_PAGE-pageTypeFromInt 2 = DICTIONARY_PAGE-pageTypeFromInt 3 = DATA_PAGE_V2-pageTypeFromInt _ = PAGE_TYPE_UNKNOWN--data ParquetEncoding- = EPLAIN- | EPLAIN_DICTIONARY- | ERLE- | EBIT_PACKED- | EDELTA_BINARY_PACKED- | EDELTA_LENGTH_BYTE_ARRAY- | EDELTA_BYTE_ARRAY- | ERLE_DICTIONARY- | EBYTE_STREAM_SPLIT- | PARQUET_ENCODING_UNKNOWN- deriving (Show, Eq)--parquetEncodingFromInt :: Int32 -> ParquetEncoding-parquetEncodingFromInt 0 = EPLAIN-parquetEncodingFromInt 2 = EPLAIN_DICTIONARY-parquetEncodingFromInt 3 = ERLE-parquetEncodingFromInt 4 = EBIT_PACKED-parquetEncodingFromInt 5 = EDELTA_BINARY_PACKED-parquetEncodingFromInt 6 = EDELTA_LENGTH_BYTE_ARRAY-parquetEncodingFromInt 7 = EDELTA_BYTE_ARRAY-parquetEncodingFromInt 8 = ERLE_DICTIONARY-parquetEncodingFromInt 9 = EBYTE_STREAM_SPLIT-parquetEncodingFromInt _ = PARQUET_ENCODING_UNKNOWN--data CompressionCodec- = UNCOMPRESSED- | SNAPPY- | GZIP- | LZO- | BROTLI- | LZ4- | ZSTD- | LZ4_RAW- | COMPRESSION_CODEC_UNKNOWN- deriving (Show, Eq)--data PageEncodingStats = PageEncodingStats- { pageEncodingPageType :: PageType- , pageEncoding :: ParquetEncoding- , pagesWithEncoding :: Int32- }- deriving (Show, Eq)--emptyPageEncodingStats :: PageEncodingStats-emptyPageEncodingStats = PageEncodingStats PAGE_TYPE_UNKNOWN PARQUET_ENCODING_UNKNOWN 0--data SizeStatistics = SizeStatisics- { unencodedByteArrayDataTypes :: Int64- , repetitionLevelHistogram :: [Int64]- , definitionLevelHistogram :: [Int64]- }- deriving (Show, Eq)--emptySizeStatistics :: SizeStatistics-emptySizeStatistics = SizeStatisics 0 [] []--data BoundingBox = BoundingBox- { xmin :: Double- , xmax :: Double- , ymin :: Double- , ymax :: Double- , zmin :: Double- , zmax :: Double- , mmin :: Double- , mmax :: Double- }- deriving (Show, Eq)--emptyBoundingBox :: BoundingBox-emptyBoundingBox = BoundingBox 0 0 0 0 0 0 0 0--data GeospatialStatistics = GeospatialStatistics- { bbox :: BoundingBox- , geospatialTypes :: [Int32]- }- deriving (Show, Eq)--emptyGeospatialStatistics :: GeospatialStatistics-emptyGeospatialStatistics = GeospatialStatistics emptyBoundingBox []--data ColumnStatistics = ColumnStatistics- { columnMin :: BS.ByteString- , columnMax :: BS.ByteString- , columnNullCount :: Int64- , columnDistictCount :: Int64- , columnMinValue :: BS.ByteString- , columnMaxValue :: BS.ByteString- , isColumnMaxValueExact :: Bool- , isColumnMinValueExact :: Bool- }- deriving (Show, Eq)--emptyColumnStatistics :: ColumnStatistics-emptyColumnStatistics = ColumnStatistics BS.empty BS.empty 0 0 BS.empty BS.empty False False--data ColumnCryptoMetadata- = COLUMN_CRYPTO_METADATA_UNKNOWN- | ENCRYPTION_WITH_FOOTER_KEY- | EncryptionWithColumnKey- { columnCryptPathInSchema :: [String]- , columnKeyMetadata :: BS.ByteString- }- deriving (Show, Eq)--data SortingColumn = SortingColumn- { columnIndex :: Int32- , columnOrderDescending :: Bool- , nullFirst :: Bool- }- deriving (Show, Eq)--emptySortingColumn :: SortingColumn-emptySortingColumn = SortingColumn 0 False False--data ColumnOrder- = TYPE_ORDER- | COLUMN_ORDER_UNKNOWN- deriving (Show, Eq)--data EncryptionAlgorithm- = ENCRYPTION_ALGORITHM_UNKNOWN- | AesGcmV1- { aadPrefix :: BS.ByteString- , aadFileUnique :: BS.ByteString- , supplyAadPrefix :: Bool- }- | AesGcmCtrV1- { aadPrefix :: BS.ByteString- , aadFileUnique :: BS.ByteString- , supplyAadPrefix :: Bool- }- deriving (Show, Eq)--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)--data Page = Page- { pageHeader :: PageHeader- , pageBytes :: BS.ByteString- }- deriving (Show, Eq)--data PageHeader = PageHeader- { pageHeaderPageType :: PageType- , uncompressedPageSize :: Int32- , compressedPageSize :: Int32- , pageHeaderCrcChecksum :: Int32- , pageTypeHeader :: PageTypeHeader- }- deriving (Show, Eq)--emptyPageHeader :: PageHeader-emptyPageHeader = PageHeader PAGE_TYPE_UNKNOWN 0 0 0 PAGE_TYPE_HEADER_UNKNOWN--data PageTypeHeader- = DataPageHeader- { dataPageHeaderNumValues :: Int32- , dataPageHeaderEncoding :: ParquetEncoding- , definitionLevelEncoding :: ParquetEncoding- , repetitionLevelEncoding :: ParquetEncoding- , dataPageHeaderStatistics :: ColumnStatistics- }- | DataPageHeaderV2- { dataPageHeaderV2NumValues :: Int32- , dataPageHeaderV2NumNulls :: Int32- , dataPageHeaderV2NumRows :: Int32- , dataPageHeaderV2Encoding :: ParquetEncoding- , definitionLevelByteLength :: Int32- , repetitionLevelByteLength :: Int32- , dataPageHeaderV2IsCompressed :: Bool- , dataPageHeaderV2Statistics :: ColumnStatistics- }- | DictionaryPageHeader- { dictionaryPageHeaderNumValues :: Int32- , dictionaryPageHeaderEncoding :: ParquetEncoding- , dictionaryPageIsSorted :: Bool- }- | INDEX_PAGE_HEADER- | PAGE_TYPE_HEADER_UNKNOWN- deriving (Show, Eq)--emptyDictionaryPageHeader :: PageTypeHeader-emptyDictionaryPageHeader = DictionaryPageHeader 0 PARQUET_ENCODING_UNKNOWN False--emptyDataPageHeader :: PageTypeHeader-emptyDataPageHeader =- DataPageHeader- 0- PARQUET_ENCODING_UNKNOWN- PARQUET_ENCODING_UNKNOWN- PARQUET_ENCODING_UNKNOWN- emptyColumnStatistics-emptyDataPageHeaderV2 :: PageTypeHeader-emptyDataPageHeaderV2 =- DataPageHeaderV2- 0- 0- 0- PARQUET_ENCODING_UNKNOWN- 0- 0 {- default for v2 is compressed -}- True- emptyColumnStatistics--data RepetitionType = REQUIRED | OPTIONAL | REPEATED | UNKNOWN_REPETITION_TYPE- deriving (Eq, Show)--data LogicalType- = STRING_TYPE- | MAP_TYPE- | LIST_TYPE- | ENUM_TYPE- | DECIMAL_TYPE- | DATE_TYPE- | DecimalType {decimalTypePrecision :: Int32, decimalTypeScale :: Int32}- | TimeType {isAdjustedToUTC :: Bool, unit :: TimeUnit}- | -- This should probably have a different, more constrained TimeUnit type.- TimestampType {isAdjustedToUTC :: Bool, unit :: TimeUnit}- | IntType {bitWidth :: Int8, intIsSigned :: Bool}- | LOGICAL_TYPE_UNKNOWN- | JSON_TYPE- | BSON_TYPE- | UUID_TYPE- | FLOAT16_TYPE- | VariantType {specificationVersion :: Int8}- | GeometryType {crs :: T.Text}- | GeographyType {crs :: T.Text, algorithm :: EdgeInterpolationAlgorithm}- deriving (Eq, Show)--data TimeUnit- = MILLISECONDS- | MICROSECONDS- | NANOSECONDS- | TIME_UNIT_UNKNOWN- deriving (Eq, Show)--data EdgeInterpolationAlgorithm- = SPHERICAL- | VINCENTY- | THOMAS- | ANDOYER- | KARNEY- deriving (Eq, Show)--repetitionTypeFromInt :: Int32 -> RepetitionType-repetitionTypeFromInt 0 = REQUIRED-repetitionTypeFromInt 1 = OPTIONAL-repetitionTypeFromInt 2 = REPEATED-repetitionTypeFromInt _ = UNKNOWN_REPETITION_TYPE--compressionCodecFromInt :: Int32 -> CompressionCodec-compressionCodecFromInt 0 = UNCOMPRESSED-compressionCodecFromInt 1 = SNAPPY-compressionCodecFromInt 2 = GZIP-compressionCodecFromInt 3 = LZO-compressionCodecFromInt 4 = BROTLI-compressionCodecFromInt 5 = LZ4-compressionCodecFromInt 6 = ZSTD-compressionCodecFromInt 7 = LZ4_RAW-compressionCodecFromInt _ = COMPRESSION_CODEC_UNKNOWN
+ src/DataFrame/IO/Parquet/Utils.hs view
@@ -0,0 +1,384 @@+{-# 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 DataFrame.Internal.Types (SBool (..), sUnbox)+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
src/DataFrame/Internal/DataFrame.hs view
@@ -18,7 +18,7 @@ import Control.DeepSeq (NFData (..), rnf) import Control.Exception (throw) import Data.Function (on)-import Data.List (sortBy, transpose, (\\))+import Data.List (sortBy, (\\)) import Data.Maybe (fromMaybe) import Data.Type.Equality ( TestEquality (testEquality),@@ -87,35 +87,89 @@ instance Show DataFrame where show :: DataFrame -> String show d =- let- rows = 20- (r, c) = dataframeDimensions d- d' =- d- { columns = V.map (takeColumn rows) (columns d)- , dataframeDimensions = (min rows r, c)- }- truncationInfo =- "\n"- ++ "Showing "- ++ show (min rows r)- ++ " rows out of "- ++ show r- in- T.unpack (asText d' False) ++ (if r > rows then truncationInfo else "")+ let (r, _) = dataframeDimensions d+ cfg = defaultTruncateConfig+ shown = if maxRows cfg > 0 then min (maxRows cfg) r else r+ body = asTextWith Plain (Just cfg) d+ footer+ | shown < r =+ "\nShowing "+ <> T.pack (show shown)+ <> " rows out of "+ <> T.pack (show r)+ | otherwise = T.empty+ in T.unpack (body <> footer) +{- | Configures how a 'DataFrame' is rendered as text. A non-positive value on+any field means \"no limit\" on that axis.++* 'maxRows' — render at most this many rows from the top of the frame.+* 'maxColumns' — when the frame has more columns than this, the middle columns+ are collapsed into a single ellipsis column.+* 'maxCellWidth' — text in any individual cell (including headers and type+ rows) longer than this is truncated with a trailing ellipsis.+-}+data TruncateConfig = TruncateConfig+ { maxRows :: Int+ , maxColumns :: Int+ , maxCellWidth :: Int+ }+ deriving (Show, Eq)++-- | Sensible defaults for GHCi: 20 rows, 10 columns, 30 characters per cell.+defaultTruncateConfig :: TruncateConfig+defaultTruncateConfig =+ TruncateConfig{maxRows = 20, maxColumns = 10, maxCellWidth = 30}++-- | Ellipsis character used to mark elided columns and clipped cells.+ellipsisText :: T.Text+ellipsisText = "\x2026"+ -- | For showing the dataframe as markdown in notebooks. toMarkdown :: DataFrame -> T.Text-toMarkdown df = asText df True+toMarkdown = asText Markdown -- | For showing the dataframe as a string markdown in notebooks. toMarkdown' :: DataFrame -> String toMarkdown' = T.unpack . toMarkdown -asText :: DataFrame -> Bool -> T.Text-asText d properMarkdown =- let header = map fst (sortBy (compare `on` snd) $ M.toList (columnIndices d))- types = V.toList $ V.filter (/= "") $ V.map getType (columns d)+asText :: RenderFormat -> DataFrame -> T.Text+asText fmt = asTextWith fmt Nothing++asTextWith :: RenderFormat -> Maybe TruncateConfig -> DataFrame -> T.Text+asTextWith fmt mTrunc d =+ let allHeaders =+ map fst (sortBy (compare `on` snd) (M.toList (columnIndices d)))+ nCols = length allHeaders+ (totalRows, _) = dataframeDimensions d++ rowCap = case mTrunc of+ Just cfg | maxRows cfg > 0 -> min totalRows (maxRows cfg)+ _ -> totalRows++ (visibleHeaders, ellipsisAt) = pickColumns mTrunc nCols allHeaders++ lookupCol name =+ fmap+ (takeColumn rowCap)+ ((V.!?) (columns d) ((M.!) (columnIndices d) name))+ survivingCols = map lookupCol visibleHeaders+ survivingTypes = map (maybe "" getType) survivingCols+ survivingData = map get survivingCols++ clipCell = case mTrunc of+ Just cfg | maxCellWidth cfg > 0 -> truncateCell (maxCellWidth cfg)+ _ -> id++ (finalHeaders, finalTypes, finalCols) = case ellipsisAt of+ Nothing -> (visibleHeaders, survivingTypes, survivingData)+ Just i ->+ let ellipsisCol = V.replicate rowCap ellipsisText+ in ( insertAt i ellipsisText visibleHeaders+ , insertAt i ellipsisText survivingTypes+ , insertAt i ellipsisCol survivingData+ )+ getType :: Column -> T.Text showMaybeType :: forall a. (Typeable a) => String showMaybeType =@@ -125,32 +179,68 @@ getType (BoxedColumn (Just _) (_ :: V.Vector a)) = T.pack $ showMaybeType @a getType (UnboxedColumn Nothing (_ :: VU.Vector a)) = T.pack $ show (typeRep @a) getType (UnboxedColumn (Just _) (_ :: VU.Vector a)) = T.pack $ showMaybeType @a- -- Separate out cases dynamically so we don't end up making round trip string- -- copies.++ -- Separate out cases dynamically so we don't end up making round trip+ -- string copies. get :: Maybe Column -> V.Vector T.Text get (Just (BoxedColumn (Just bm) (column :: V.Vector a))) = V.generate (V.length column) $ \i -> if bitmapTestBit bm i then T.pack (show (Just (V.unsafeIndex column i))) else "Nothing"- get (Just (BoxedColumn Nothing (column :: V.Vector a))) = case testEquality (typeRep @a) (typeRep @T.Text) of- Just Refl -> column- Nothing -> case testEquality (typeRep @a) (typeRep @String) of- Just Refl -> V.map T.pack column- Nothing -> V.map (T.pack . show) column+ get (Just (BoxedColumn Nothing (column :: V.Vector a))) =+ case testEquality (typeRep @a) (typeRep @T.Text) of+ Just Refl -> column+ Nothing -> case testEquality (typeRep @a) (typeRep @String) of+ Just Refl -> V.map T.pack column+ Nothing -> V.map (T.pack . show) column get (Just (UnboxedColumn (Just bm) column)) =- let col = V.convert column- in V.generate (V.length col) $ \i ->- if bitmapTestBit bm i- then T.pack (show (Just (V.unsafeIndex col i)))- else "Nothing"- get (Just (UnboxedColumn Nothing column)) = V.map (T.pack . show) (V.convert column)+ V.generate (VU.length column) $ \i ->+ if bitmapTestBit bm i+ then T.pack (show (Just (VU.unsafeIndex column i)))+ else "Nothing"+ get (Just (UnboxedColumn Nothing column)) =+ V.generate (VU.length column) (T.pack . show . VU.unsafeIndex column) get Nothing = V.empty- getTextColumnFromFrame _df (_i, name) = get $ (V.!?) (columns d) ((M.!) (columnIndices d) name)- rows =- transpose $- zipWith (curry (V.toList . getTextColumnFromFrame d)) [(0 :: Int) ..] header- in showTable properMarkdown header types rows+ in showTable+ fmt+ (map clipCell finalHeaders)+ (map clipCell finalTypes)+ (map (V.map clipCell) finalCols)++{- | Decide which columns survive horizontal truncation and where (if anywhere)+to splice in the ellipsis column. The split puts the extra column on the+left for odd 'maxColumns'; the ellipsis is only inserted when it actually+saves space (i.e. the frame has more than 'maxColumns' + 1 columns).+-}+pickColumns ::+ Maybe TruncateConfig ->+ Int ->+ [a] ->+ ([a], Maybe Int)+pickColumns mTrunc nCols xs = case mTrunc of+ Just cfg+ | let c = maxColumns cfg+ , c > 0+ , nCols > c + 1 ->+ let leftN = (c + 1) `div` 2+ rightN = c - leftN+ in ( Prelude.take leftN xs ++ Prelude.drop (nCols - rightN) xs+ , Just leftN+ )+ _ -> (xs, Nothing)++-- | Splice @x@ into @xs@ at index @i@ (0-based), shifting later elements right.+insertAt :: Int -> a -> [a] -> [a]+insertAt i x xs = let (l, r) = splitAt i xs in l ++ x : r++-- | Cap a single cell's rendered length, appending an ellipsis when shortened.+truncateCell :: Int -> T.Text -> T.Text+truncateCell n t+ | n <= 0 = t+ | T.compareLength t n /= GT = t+ | n == 1 = ellipsisText+ | otherwise = T.take (n - 1) t <> ellipsisText -- | O(1) Creates an empty dataframe empty :: DataFrame
src/DataFrame/Internal/Schema.hs view
@@ -4,11 +4,13 @@ {-# LANGUAGE InstanceSigs #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TemplateHaskellQuotes #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} module DataFrame.Internal.Schema where +import Data.Char (isUpper, toLower, toUpper) import qualified Data.Map as M import qualified Data.Proxy as P import qualified Data.Text as T@@ -16,6 +18,9 @@ import Data.Maybe (isJust) import Data.Type.Equality (TestEquality (..)) import DataFrame.Internal.Column (Columnable)+import DataFrame.Internal.Expression (Expr)+import DataFrame.Operators (col)+import Language.Haskell.TH import Type.Reflection (typeRep) -- | A runtime tag for a column’s element type.@@ -108,3 +113,130 @@ -} makeSchema :: [(T.Text, SchemaType)] -> Schema makeSchema = Schema . M.fromList++{- | Auto-generate a runtime 'Schema' (and per-column @'Expr'@ accessors)+from a record ADT.++The splice reifies the record, applies @camelCase -> snake_case@ to each+record-selector name, and emits:++* a top-level @\<lower-first TyConName\>Schema :: 'Schema'@ binding suitable+ for passing to 'DataFrame.IO.CSV.readCsvWithSchema' /+ 'DataFrame.IO.CSV.readCsvWithOpts'.+* one @\<lower-first TyConName\>\<UpperFirst FieldName\> :: 'Expr' /ty/@ binding+ per field, so you can refer to columns in expression DSL code by name+ without writing @col \@/ty/ "snake_case_name"@ at every call site.++@+data Order = Order { customerId :: Int, region :: Text, amount :: Double }++\$(deriveSchema ''Order)+-- expands to:+-- orderSchema :: Schema+-- orderSchema = makeSchema+-- [ ("customer_id", schemaType \@Int)+-- , ("region", schemaType \@Text)+-- , ("amount", schemaType \@Double)+-- ]+-- orderCustomerId :: Expr Int+-- orderCustomerId = col "customer_id"+-- orderRegion :: Expr Text+-- orderRegion = col "region"+-- orderAmount :: Expr Double+-- orderAmount = col "amount"++main = do+ df <- D.readCsvWithSchema orderSchema "orders.csv"+ let bigOrders = D.filterWhere (orderAmount .>. 100) df+ ...+@++The data type must have exactly one record constructor; sum types or+positional constructors fail the splice with a descriptive error. Field+types must satisfy @('Columnable' a, 'Read' a)@ — the same constraints+'schemaType' already requires.+-}+deriveSchema :: Name -> DecsQ+deriveSchema tyName = do+ info <- reify tyName+ fields <- extractRecordFields tyName info+ let entries =+ [ (camelToSnake fieldBase, fieldBase, fTy)+ | (fName, _bang, fTy) <- fields+ , let fieldBase = nameBase fName+ ]+ schemaName = mkName (lowerFirst (nameBase tyName) ++ "Schema")+ prefix = lowerFirst (nameBase tyName)+ tupleE (colName, _, fTy) =+ TupE+ [ Just (AppE (VarE 'T.pack) (LitE (StringL colName)))+ , Just (AppTypeE (VarE 'schemaType) fTy)+ ]+ schemaBody =+ AppE (VarE 'makeSchema) (ListE (map tupleE entries))+ schemaDecls =+ [ SigD schemaName (ConT ''Schema)+ , ValD (VarP schemaName) (NormalB schemaBody) []+ ]+ accessorDecls =+ concat+ [ [ SigD accName (AppT (ConT ''Expr) fTy)+ , ValD+ (VarP accName)+ ( NormalB+ ( AppE+ (VarE 'col)+ ( AppE+ (VarE 'T.pack)+ (LitE (StringL colName))+ )+ )+ )+ []+ ]+ | (colName, fieldBase, fTy) <- entries+ , let accName = mkName (prefix ++ upperFirst fieldBase)+ ]+ pure (schemaDecls ++ accessorDecls)++extractRecordFields :: Name -> Info -> Q [VarBangType]+extractRecordFields _ (TyConI dec) = case dec of+ DataD _ _ _ _ [RecC _ fs] _ -> pure fs+ NewtypeD _ _ _ _ (RecC _ fs) _ -> pure fs+ DataD _ n _ _ _ _ ->+ fail $+ "deriveSchema: "+ ++ show n+ ++ " must have exactly one record constructor"+ NewtypeD _ n _ _ _ _ ->+ fail $+ "deriveSchema: " ++ show n ++ " newtype must use record syntax"+ other ->+ fail $+ "deriveSchema: unsupported declaration: " ++ show other+extractRecordFields tyName _ =+ fail $+ "deriveSchema: "+ ++ show tyName+ ++ " is not a data/newtype declaration"++-- Local @camelCase -> snake_case@: lowercase the first char, then prefix+-- @\'_\'@ before any uppercase character (lowercased). Duplicated from+-- 'DataFrame.Typed.TH.camelToSnake' to keep this module free of any+-- @DataFrame.Typed.*@ imports.+camelToSnake :: String -> String+camelToSnake [] = []+camelToSnake (c : cs) = toLower c : go cs+ where+ go [] = []+ go (x : xs)+ | isUpper x = '_' : toLower x : go xs+ | otherwise = x : go xs++lowerFirst :: String -> String+lowerFirst [] = []+lowerFirst (c : cs) = toLower c : cs++upperFirst :: String -> String+upperFirst [] = []+upperFirst (c : cs) = toUpper c : cs
src/DataFrame/Operations/Permutation.hs view
@@ -22,8 +22,9 @@ import DataFrame.Errors (DataFrameException (..)) import DataFrame.Internal.Column (Column (..), Columnable, atIndicesStable) import DataFrame.Internal.DataFrame (DataFrame (..), unsafeGetColumn)-import DataFrame.Internal.Expression (Expr (Col))+import DataFrame.Internal.Expression (Expr (Col), getColumns) import DataFrame.Operations.Core (columnNames, dimensions)+import DataFrame.Operations.Transformations (derive) import System.Random (Random (randomR), RandomGen) import Type.Reflection (typeRep) @@ -38,15 +39,40 @@ (==) (Desc _) (Desc _) = True (==) _ _ = False -getSortColumnName :: SortOrder -> T.Text-getSortColumnName (Asc (Col n)) = n-getSortColumnName (Desc (Col n)) = n-getSortColumnName _ = error "Sorting on compound column"+sortOrderColumns :: SortOrder -> [T.Text]+sortOrderColumns (Asc e) = getColumns e+sortOrderColumns (Desc e) = getColumns e mustFlipCompare :: SortOrder -> Bool mustFlipCompare (Asc _) = True mustFlipCompare (Desc _) = False +{- | Materialize any compound sort expressions into synthetic columns on+a working dataframe, returning rewritten 'SortOrder's that reference+those columns by name.+-}+prepareSortColumns :: [SortOrder] -> DataFrame -> ([SortOrder], DataFrame)+prepareSortColumns = go 0+ where+ go _ [] acc = ([], acc)+ go i (ord : rest) acc =+ let (ord', acc') = materializeSortOrder i ord acc+ (rest', acc'') = go (i + 1) rest acc'+ in (ord' : rest', acc'')++materializeSortOrder :: Int -> SortOrder -> DataFrame -> (SortOrder, DataFrame)+materializeSortOrder _ ord@(Asc (Col _)) df = (ord, df)+materializeSortOrder _ ord@(Desc (Col _)) df = (ord, df)+materializeSortOrder i (Asc (e :: Expr a)) df =+ let name = syntheticName i+ in (Asc (Col name :: Expr a), derive name e df)+materializeSortOrder i (Desc (e :: Expr a)) df =+ let name = syntheticName i+ in (Desc (Col name :: Expr a), derive name e df)++syntheticName :: Int -> T.Text+syntheticName i = "__sortBy_synthetic_" <> T.pack (show i) <> "__"+ {- | O(k log n) Sorts the dataframe by a given row. > sortBy Ascending ["Age"] df@@ -56,22 +82,24 @@ DataFrame -> DataFrame sortBy sortOrds df- | any (`notElem` columnNames df) names =+ | not (null missing) = throw $ ColumnsNotFoundException- (names L.\\ columnNames df)+ missing "sortBy" (columnNames df) | otherwise = let- comparators = map (`sortOrderComparator` df) sortOrds+ (sortOrds', df') = prepareSortColumns sortOrds df+ comparators = map (`sortOrderComparator` df') sortOrds' compositeCompare i j = mconcat [c i j | c <- comparators]- nRows = fst (dataframeDimensions df)+ nRows = fst (dataframeDimensions df') indexes = sortIndices compositeCompare nRows in df{columns = V.map (atIndicesStable indexes) (columns df)} where- names = map getSortColumnName sortOrds+ referenced = L.nub (concatMap sortOrderColumns sortOrds)+ missing = referenced L.\\ columnNames df {- | Build a row-index comparator from a SortOrder and a DataFrame. The Ord dictionary is recovered from the SortOrder GADT.
+ src/DataFrame/TH.hs view
@@ -0,0 +1,213 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE TypeApplications #-}++{- |+Module : DataFrame.TH+License : MIT++Template Haskell splices for the untyped 'DataFrame' API.++These splices generate top-level @Expr a@ bindings — one per column of a+'DataFrame' — so you can refer to columns by name in GHCi or in code+without writing @F.col \@T \"name\"@ at every use site.++Most users will reach for these via the @:declareColumns@ GHCi macro+provided by @dataframe.ghci@.++@+ghci> :set -XTemplateHaskell+ghci> :declareColumns df+ghci> :type passengers+passengers :: Expr Int+@++The typed-API equivalents (which generate a schema type synonym) live in+"DataFrame.Typed.TH".+-}+module DataFrame.TH (+ -- * Declare one binding per column+ declareColumns,+ declareColumnsWithPrefix,+ declareColumnsWithPrefix',++ -- * From a file+ declareColumnsFromCsvFile,+ declareColumnsFromCsvWithOpts,+ declareColumnsFromParquetFile,++ -- * Type-string parser (exposed for testing)+ typeFromString,+) where++import Control.Monad (filterM, forM)+import Control.Monad.IO.Class (liftIO)+import Data.Function (on)+import Data.Functor ((<&>))+import Data.Int (Int64)+import qualified Data.List as L+import qualified Data.Map as M+import qualified Data.Maybe as Maybe+import qualified Data.Set as S+import qualified Data.Text as T++import Language.Haskell.TH+import qualified Language.Haskell.TH.Syntax as TH++import System.Directory (doesDirectoryExist)+import System.FilePath ((</>))+import System.FilePath.Glob (glob)++import DataFrame.Functions (sanitize)+import qualified DataFrame.IO.CSV as CSV+import qualified DataFrame.IO.Parquet as Parquet+import DataFrame.IO.Parquet.Schema (schemaToEmptyDataFrame)+import DataFrame.IO.Parquet.Thrift (+ cc_meta_data,+ cmd_path_in_schema,+ cmd_statistics,+ rg_columns,+ row_groups,+ schema,+ stats_null_count,+ unField,+ )+import DataFrame.Internal.Column (columnTypeString)+import DataFrame.Internal.DataFrame (+ DataFrame (..),+ unsafeGetColumn,+ )+import qualified DataFrame.Internal.DataFrame as DI+import DataFrame.Internal.Expression (Expr)+import DataFrame.Operators (col)+import Prelude as P++typeFromString :: [String] -> Q Type+typeFromString [] = fail "No type specified"+typeFromString [t0] = do+ let t = trim t0+ case stripBrackets t of+ Just inner -> typeFromString [inner] <&> AppT ListT+ Nothing+ | t == "Text" || t == "Data.Text.Text" || t == "T.Text" ->+ pure (ConT ''T.Text)+ | otherwise -> do+ m <- lookupTypeName t+ case m of+ Just tyName -> pure (ConT tyName)+ Nothing -> fail $ "Unsupported type: " ++ t0+typeFromString [tycon, t1] = AppT <$> typeFromString [tycon] <*> typeFromString [t1]+typeFromString [tycon, t1, t2] =+ (\outer a b -> AppT (AppT outer a) b)+ <$> typeFromString [tycon]+ <*> typeFromString [t1]+ <*> typeFromString [t2]+typeFromString s = fail $ "Unsupported types: " ++ unwords s++trim :: String -> String+trim = dropWhile (== ' ') . reverse . dropWhile (== ' ') . reverse++stripBrackets :: String -> Maybe String+stripBrackets s =+ case s of+ ('[' : rest)+ | P.not (null rest) && last rest == ']' ->+ Just (init rest)+ _ -> Nothing++{- | Splice a binding for every column of the 'DataFrame' read from a CSV+file. Each binding has type @Expr T@ where @T@ is the inferred column+type.+-}+declareColumnsFromCsvFile :: String -> DecsQ+declareColumnsFromCsvFile path = do+ df <-+ liftIO+ ( CSV.readSeparated+ (CSV.defaultReadOptions{CSV.numColumns = Just 100})+ path+ )+ declareColumns df++-- | Like 'declareColumnsFromCsvFile' but with custom 'CSV.ReadOptions'.+declareColumnsFromCsvWithOpts :: CSV.ReadOptions -> String -> DecsQ+declareColumnsFromCsvWithOpts opts path = do+ df <- liftIO (CSV.readSeparated opts path)+ declareColumns df++{- | Splice a binding for every column of a parquet file (or directory of+parquet files). The schema is read from each file's metadata and merged.+-}+declareColumnsFromParquetFile :: String -> DecsQ+declareColumnsFromParquetFile path = do+ isDir <- liftIO $ doesDirectoryExist path+ let pat = if isDir then path </> "*.parquet" else path+ matches <- liftIO $ glob pat+ files <- liftIO $ filterM (fmap P.not . doesDirectoryExist) matches+ metas <- liftIO $ mapM Parquet.readMetadataFromPath files+ let nullableCols :: S.Set T.Text+ nullableCols =+ S.fromList+ [ T.pack (last colPath)+ | meta <- metas+ , rg <- unField (row_groups meta)+ , cc <- unField (rg_columns rg)+ , Just cm <- [unField (cc_meta_data cc)]+ , let colPath = map T.unpack (unField (cmd_path_in_schema cm))+ , P.not (null colPath)+ , let nc :: Int64+ nc = case unField (cmd_statistics cm) of+ Nothing -> 0+ Just stats ->+ Maybe.fromMaybe 0 (unField $ stats_null_count stats)+ , nc > 0+ ]+ let df =+ foldl+ ( \acc meta ->+ acc+ <> schemaToEmptyDataFrame+ nullableCols+ (unField (schema meta))+ )+ DI.empty+ metas++ declareColumns df++{- | Splice a binding for every column of @df@, named after the column.+Column names that are not valid Haskell identifiers are sanitized+(see 'DataFrame.Functions.sanitize').+-}+declareColumns :: DataFrame -> DecsQ+declareColumns = declareColumnsWithPrefix' Nothing++-- | Like 'declareColumns' but prefixes every binding name with @prefix_@.+declareColumnsWithPrefix :: T.Text -> DataFrame -> DecsQ+declareColumnsWithPrefix prefix = declareColumnsWithPrefix' (Just prefix)++-- | Like 'declareColumnsWithPrefix' but takes an optional prefix.+declareColumnsWithPrefix' :: Maybe T.Text -> DataFrame -> DecsQ+declareColumnsWithPrefix' prefix df =+ let+ names = (map fst . L.sortBy (compare `on` snd) . M.toList . columnIndices) df+ types = map (columnTypeString . (`unsafeGetColumn` df)) names+ specs =+ zipWith+ ( \colName type_ ->+ ( colName+ , maybe "" (sanitize . (<> "_")) prefix <> sanitize colName+ , type_+ )+ )+ names+ types+ in+ fmap concat $ forM specs $ \(raw, nm, tyStr) -> do+ ty <- typeFromString (words tyStr)+ let n = mkName (T.unpack nm)+ sig <- sigD n [t|Expr $(pure ty)|]+ val <- valD (varP n) (normalB [|col $(TH.lift raw)|]) []+ pure [sig, val]
src/DataFrame/Typed.hs view
@@ -46,13 +46,13 @@ T.filterAllJust df :: TypedDataFrame '[Column \"x\" Double, Column \"y\" Int] @ -== Typed aggregation (Option B)+== Typed aggregation @ result = T.aggregate- (T.agg \@\"total\" (T.tsum (T.col \@\"salary\"))- $ T.agg \@\"count\" (T.tcount (T.col \@\"salary\"))- $ T.aggNil)+ ( T.as \@\"total\" (T.sum (T.col \@\"salary\"))+ . T.as \@\"count\" (T.count (T.col \@\"salary\"))+ ) (T.groupBy \@'[\"dept\"] employees) @ -}@@ -122,9 +122,6 @@ asc, desc, - -- * Named expression helper- DataFrame.Typed.Expr.as,- -- * Freeze / thaw boundary freeze, freezeWithError,@@ -182,10 +179,9 @@ rightJoin, fullOuterJoin, - -- * GroupBy and Aggregation (Option B)+ -- * GroupBy and Aggregation groupBy,- agg,- aggNil,+ as, aggregate, aggregateUntyped, @@ -193,7 +189,23 @@ deriveSchema, deriveSchemaFromCsvFile, deriveSchemaFromCsvFileWith,+ deriveSchemaFromType,+ deriveSchemaFromTypeWith,+ SchemaOptions (..),+ defaultSchemaOptions, + -- * Record bridge (ADT <-> TypedDataFrame)+ HasSchema (..),+ fromRecordsTyped,+ toRecordsTyped,++ -- * Generics opt-in for schema derivation+ SchemaOf,+ SchemaOfRaw,+ NameCase (..),+ genericToColumns,+ genericFromColumns,+ -- * Schema type families (for advanced use) Lookup, SafeLookup,@@ -226,21 +238,36 @@ import DataFrame.Typed.Access (columnAsList, columnAsVector) import DataFrame.Typed.Aggregate (- agg,- aggNil, aggregate, aggregateUntyped,+ as, groupBy, ) import DataFrame.Typed.Expr import DataFrame.Typed.Freeze (freeze, freezeWithError, thaw, unsafeFreeze)+import DataFrame.Typed.Generic (+ NameCase (..),+ SchemaOf,+ SchemaOfRaw,+ genericFromColumns,+ genericToColumns,+ ) import DataFrame.Typed.Join (fullOuterJoin, innerJoin, leftJoin, rightJoin) import DataFrame.Typed.Operations+import DataFrame.Typed.Record (+ HasSchema (..),+ fromRecordsTyped,+ toRecordsTyped,+ ) import DataFrame.Typed.Schema import DataFrame.Typed.TH (+ SchemaOptions (..),+ defaultSchemaOptions, deriveSchema, deriveSchemaFromCsvFile, deriveSchemaFromCsvFileWith,+ deriveSchemaFromType,+ deriveSchemaFromTypeWith, ) import DataFrame.Typed.Types ( Column,
src/DataFrame/Typed/Aggregate.hs view
@@ -11,9 +11,8 @@ -- * Typed groupBy groupBy, - -- * Typed aggregation builder (Option B)- agg,- aggNil,+ -- * Naming an aggregation+ as, -- * Running aggregations aggregate,@@ -47,50 +46,72 @@ TypedDataFrame cols -> TypedGrouped keys cols groupBy (TDF df) = TGD (DA.groupBy (symbolVals @keys) df) --- | The empty aggregation — no output columns beyond the group keys.-aggNil :: TAgg keys cols '[]-aggNil = TAggNil+{- | Build a named aggregation entry. The result column name is supplied via+@TypeApplications@; the underlying expression is validated against the+source schema at compile time. -{- | Add one aggregation to the builder.+@as@ produces a /transformer/ on the aggregation chain — entries compose+with plain @(.)@ from Prelude (or via @(|>)@ for SQL-like postfix+reading). 'aggregate' applies the composed transformer to the empty chain+internally, so no terminator is needed. -Each call prepends a @Column name a@ to the result schema and records-the runtime 'NamedExpr'. The expression is validated against the-source schema @cols@ at compile time.+==== __Prefix form__ @-agg \@\"total_sales\" (tsum (col \@\"salary\"))- $ agg \@\"avg_price\" (tmean (col \@\"price\"))- $ aggNil+result = grouped |> aggregate+ ( as \@\"total\" (sum (col \@\"amount\"))+ . as \@\"orders\" (count (col \@\"order_id\"))+ . as \@\"avg\" (mean (col \@\"amount\"))+ ) @++==== __Postfix form (SQL-like)__++@+result = grouped |> aggregate+ ( (sum (col \@\"amount\") |> as \@\"total\")+ . (count (col \@\"order_id\") |> as \@\"orders\")+ . (mean (col \@\"amount\") |> as \@\"avg\")+ )+@++Per-entry parentheses are required in the postfix form because+@(.)@ binds tighter than @(|>)@. -}-agg ::+as :: forall name a keys cols aggs.- ( KnownSymbol name- , Columnable a- ) =>- TExpr cols a -> TAgg keys cols aggs -> TAgg keys cols (Column name a ': aggs)-agg = TAggCons colName- where- colName = T.pack (symbolVal (Proxy @name))+ (KnownSymbol name, Columnable a) =>+ TExpr cols a ->+ TAgg keys cols aggs ->+ TAgg keys cols (Column name a ': aggs)+as = TAggCons (T.pack (symbolVal (Proxy @name))) -{- | Run a typed aggregation.+{- | Run a typed aggregation against a grouped DataFrame. -Result schema = grouping key columns ++ aggregated columns (in declaration order).+The first argument is a chain of 'as' entries composed with @(.)@. The+empty composition (@id@) yields just the group keys. The result schema is+the group-key columns followed by the aggregation columns in declaration+order. @-result = aggregate- (agg \@\"total\" (tsum (col @"salary")) $ agg \@\"count\" (tcount (col @"salary") $ aggNil)- (groupBy \@'[\"dept\"] employees)--- result :: TDF '[Column \"dept\" Text, Column \"total\" Double, Column \"count\" Int]+result = grouped |> aggregate+ ( as \@\"total\" (sum (col \@\"amount\"))+ . as \@\"orders\" (count (col \@\"order_id\"))+ )+-- result :: TypedDataFrame+-- '[ Column \"region\" Text+-- , Column \"total\" Double+-- , Column \"orders\" Int+-- ] @ -} aggregate :: forall keys cols aggs.- TAgg keys cols aggs ->+ (TAgg keys cols '[] -> TAgg keys cols aggs) -> TypedGrouped keys cols -> TypedDataFrame (Append (GroupKeyColumns keys cols) (Reverse aggs))-aggregate tagg (TGD gdf) =- unsafeFreeze (DA.aggregate (taggToNamedExprs tagg) gdf)+aggregate build (TGD gdf) =+ unsafeFreeze (DA.aggregate (taggToNamedExprs (build TAggNil)) gdf) -- | Escape hatch: run an untyped aggregation and return a raw 'DataFrame'. aggregateUntyped :: [NamedExpr] -> TypedGrouped keys cols -> D.DataFrame
src/DataFrame/Typed/Expr.hs view
@@ -119,9 +119,6 @@ unsafeCastExpr, toDouble, - -- * Named expression helper- as,- -- * Sort helpers asc, desc,@@ -138,8 +135,6 @@ import DataFrame.Internal.Expression ( BinaryOp (..), Expr (..),- NamedExpr,- UExpr (..), UnaryOp (..), ) import DataFrame.Internal.Nullable (@@ -639,14 +634,6 @@ toDouble :: (Columnable a, Real a) => TExpr cols a -> TExpr cols Double toDouble (TExpr e) = TExpr (F.toDouble e)------------------------------------------------------------------------------------ Named expression helper------------------------------------------------------------------------------------ | Create a 'NamedExpr' for use with 'aggregateUntyped'.-as :: (Columnable a) => TExpr cols a -> T.Text -> NamedExpr-as (TExpr e) name = (name, UExpr e) -- | Create an ascending sort order from a typed expression. asc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols
+ src/DataFrame/Typed/Generic.hs view
@@ -0,0 +1,205 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE UndecidableInstances #-}++{- |+Module : DataFrame.Typed.Generic+License : MIT++Generic-based opt-in for record-to-schema derivation. Mirrors the Template+Haskell splice in "DataFrame.Typed.TH" but builds the schema type from a+@GHC.Generics.Generic@ instance instead of @reify@.++Use it like this:++@+data Order = Order+ { orderId :: Int64+ , region :: Text+ , amount :: Double+ } deriving (Show, Eq, Generic)++type OrderSchema = SchemaOf Order++instance HasSchema Order OrderSchema where+ toColumns = genericToColumns+ fromColumns = genericFromColumns+@++Field names are translated with the @CamelCase -> snake_case@ rule+(matching 'DataFrame.Typed.TH.camelToSnake'); use 'SchemaOfRaw' if you+want the schema to keep the record selector names verbatim — in that+case you cannot use 'genericToColumns' \/ 'genericFromColumns' and must+either hand-roll the instance or use the TH splice with a custom name+transform.+-}+module DataFrame.Typed.Generic (+ -- * Type-level schema derivation+ NameCase (..),+ SchemaOf,+ SchemaOfRaw,+ RepToSchema,+ CamelToSnake,++ -- * Value-level default methods+ genericToColumns,+ genericFromColumns,+ GHasColumns,+) where++import Data.Kind (Type)+import Data.Proxy (Proxy (..))+import qualified Data.Text as T+import qualified Data.Vector as VB+import GHC.Generics (+ C,+ D,+ Generic (..),+ K1 (..),+ M1 (..),+ Meta (..),+ S,+ type (:*:) (..),+ )+import GHC.TypeLits (+ CharToNat,+ ConsSymbol,+ KnownSymbol,+ NatToChar,+ Symbol,+ UnconsSymbol,+ symbolVal,+ type (+),+ )++import Data.Type.Bool (If, type (&&))+import Data.Type.Ord (type (<=?))++import qualified DataFrame.Internal.Column as C+import qualified DataFrame.Internal.DataFrame as D+import DataFrame.Typed.Record (requireColumn)+import DataFrame.Typed.Schema (Append)+import DataFrame.Typed.TH (camelToSnake)+import DataFrame.Typed.Types (Column)++{- | Field-name policy applied to record selectors when computing+'RepToSchema'.++* 'SnakeCase' — translate @camelCaseField@ to @\"camel_case_field\"@.+* 'IdentityCase' — keep the selector name verbatim.+-}+data NameCase = SnakeCase | IdentityCase++{- | The schema type @[Column name ty, ...]@ derived from the 'Rep' of a+record type, with the given 'NameCase' applied to each field name.+-}+type family RepToSchema (nc :: NameCase) (r :: Type -> Type) :: [Type] where+ RepToSchema nc (M1 D _ f) = RepToSchema nc f+ RepToSchema nc (M1 C _ f) = RepToSchema nc f+ RepToSchema nc (a :*: b) = Append (RepToSchema nc a) (RepToSchema nc b)+ RepToSchema nc (M1 S ('MetaSel ('Just name) _ _ _) (K1 _ a)) =+ '[Column (TransformName nc name) a]++type family TransformName (nc :: NameCase) (name :: Symbol) :: Symbol where+ TransformName 'SnakeCase s = CamelToSnake s+ TransformName 'IdentityCase s = s++-- | Type-level camelCase -> snake_case. Matches 'camelToSnake' at the value level.+type family CamelToSnake (s :: Symbol) :: Symbol where+ CamelToSnake s = SnakeStart (UnconsSymbol s)++type family SnakeStart (mu :: Maybe (Char, Symbol)) :: Symbol where+ SnakeStart 'Nothing = ""+ SnakeStart ('Just '(c, r)) =+ ConsSymbol (ToLowerChar c) (SnakeRest (UnconsSymbol r))++type family SnakeRest (mu :: Maybe (Char, Symbol)) :: Symbol where+ SnakeRest 'Nothing = ""+ SnakeRest ('Just '(c, r)) =+ SnakeStep (IsUpperChar c) c (SnakeRest (UnconsSymbol r))++type family SnakeStep (up :: Bool) (c :: Char) (rest :: Symbol) :: Symbol where+ SnakeStep 'True c rest = ConsSymbol '_' (ConsSymbol (ToLowerChar c) rest)+ SnakeStep 'False c rest = ConsSymbol c rest++type family IsUpperChar (c :: Char) :: Bool where+ IsUpperChar c =+ (CharToNat 'A' <=? CharToNat c) && (CharToNat c <=? CharToNat 'Z')++type family ToLowerChar (c :: Char) :: Char where+ ToLowerChar c = If (IsUpperChar c) (NatToChar (CharToNat c + 32)) c++-- | Snake_case schema derived from @a@'s 'Generic' representation.+type SchemaOf a = RepToSchema 'SnakeCase (Rep a)++-- | Identity-cased schema derived from @a@'s 'Generic' representation.+type SchemaOfRaw a = RepToSchema 'IdentityCase (Rep a)++{- | Walks the 'Rep' tree of a record, producing or consuming a list of+named columns. Used by 'genericToColumns' \/ 'genericFromColumns'.+-}+class GHasColumns (r :: Type -> Type) where+ gToColumns :: [r p] -> [(T.Text, C.Column)]+ gFromColumns :: D.DataFrame -> Either T.Text [r p]++instance (GHasColumns f) => GHasColumns (M1 D meta f) where+ gToColumns rs = gToColumns (map unM1 rs)+ gFromColumns df = fmap (map M1) (gFromColumns df)++instance (GHasColumns f) => GHasColumns (M1 C meta f) where+ gToColumns rs = gToColumns (map unM1 rs)+ gFromColumns df = fmap (map M1) (gFromColumns df)++instance (GHasColumns a, GHasColumns b) => GHasColumns (a :*: b) where+ gToColumns rs =+ gToColumns (map (\(x :*: _) -> x) rs)+ ++ gToColumns (map (\(_ :*: y) -> y) rs)+ gFromColumns df = do+ as <- gFromColumns df+ bs <- gFromColumns df+ pure (zipWith (:*:) as bs)++instance+ (KnownSymbol name, C.Columnable a) =>+ GHasColumns+ ( M1+ S+ ('MetaSel ('Just name) su ss ds)+ (K1 i a)+ )+ where+ gToColumns rs =+ let colName = T.pack (camelToSnake (symbolVal (Proxy @name)))+ vals = map (unK1 . unM1) rs+ in [(colName, C.fromList vals)]+ gFromColumns df = do+ let colName = T.pack (camelToSnake (symbolVal (Proxy @name)))+ v <- requireColumn @a colName df+ pure (map (M1 . K1) (VB.toList v))++{- | Default implementation of 'DataFrame.Typed.Record.toColumns' for any+@Generic@ record. Field names are translated with @camelCase -> snake_case@.++@+instance HasSchema Order (SchemaOf Order) where+ toColumns = genericToColumns+ fromColumns = genericFromColumns+@+-}+genericToColumns ::+ forall a. (Generic a, GHasColumns (Rep a)) => [a] -> [(T.Text, C.Column)]+genericToColumns = gToColumns . map from++{- | Default implementation of 'DataFrame.Typed.Record.fromColumns' for any+@Generic@ record.+-}+genericFromColumns ::+ forall a. (Generic a, GHasColumns (Rep a)) => D.DataFrame -> Either T.Text [a]+genericFromColumns df = fmap (map to) (gFromColumns df)
+ src/DataFrame/Typed/Record.hs view
@@ -0,0 +1,105 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeFamilies #-}++{- |+Module : DataFrame.Typed.Record+License : MIT++Bridge a Haskell record type to a typed dataframe schema. Instances are+typically generated by the @deriveSchemaFromType@ Template Haskell splice+(see "DataFrame.Typed.TH"), but can be written by hand as well, or+plugged into the @Generic@-derived defaults from "DataFrame.Typed.Generic".+-}+module DataFrame.Typed.Record (+ -- * Class+ HasSchema (..),++ -- * Untyped helpers+ fromRecords,+ toRecords,++ -- * Typed helpers+ fromRecordsTyped,+ toRecordsTyped,++ -- * Helpers used by generated code+ requireColumn,+) where++import Data.Kind (Type)+import qualified Data.Text as T+import qualified Data.Vector as VB++import qualified DataFrame.Internal.Column as C+import qualified DataFrame.Internal.DataFrame as D+import DataFrame.Operations.Core (fromNamedColumns)+import DataFrame.Typed.Types (TypedDataFrame (..))++{- | Bridge a Haskell record type @a@ to a typed-dataframe schema.++The schema is exposed as an associated type family 'Schema' so that+instances can pick it up from a 'GHC.Generics.Rep' computation (see+'DataFrame.Typed.Generic.SchemaOf') or from an explicit list emitted by+'DataFrame.Typed.TH.deriveSchemaFromType'.++@toColumns@ explodes a list of records into a list of named columns.+@fromColumns@ reconstructs the records from a 'D.DataFrame', returning+@Left err@ if a column is missing or has the wrong type.+-}+class HasSchema a where+ type Schema a :: [Type]+ toColumns :: [a] -> [(T.Text, C.Column)]+ fromColumns :: D.DataFrame -> Either T.Text [a]++{- | Build an untyped 'D.DataFrame' from a list of records.++@+data Order = Order { orderId :: Int64, region :: Text, amount :: Double }+\$(deriveSchemaFromType ''Order)++xs :: [Order]+xs = [Order 1 "us" 10.0, Order 2 "eu" 20.0]++df :: DataFrame+df = fromRecords xs+@+-}+fromRecords :: (HasSchema a) => [a] -> D.DataFrame+fromRecords = fromNamedColumns . toColumns++{- | Parse a list of records out of an untyped 'D.DataFrame'.++Returns @Left err@ on schema mismatch (missing column, wrong type).+-}+toRecords :: (HasSchema a) => D.DataFrame -> Either T.Text [a]+toRecords = fromColumns++-- | Like 'fromRecords' but returns a 'TypedDataFrame' tagged with the schema.+fromRecordsTyped :: forall a. (HasSchema a) => [a] -> TypedDataFrame (Schema a)+fromRecordsTyped = TDF . fromRecords++-- | Like 'toRecords' but accepts a 'TypedDataFrame'.+toRecordsTyped ::+ forall a. (HasSchema a) => TypedDataFrame (Schema a) -> Either T.Text [a]+toRecordsTyped (TDF df) = fromColumns df++{- | Extract a column as a boxed vector by name, returning a 'T.Text' error+on missing column or type mismatch.++Used by code generated by 'DataFrame.Typed.TH.deriveSchemaFromType'.+-}+requireColumn ::+ forall a.+ (C.Columnable a) => T.Text -> D.DataFrame -> Either T.Text (VB.Vector a)+requireColumn name df = case D.getColumn name df of+ Nothing ->+ Left $ "Column '" <> name <> "' not found in DataFrame"+ Just col -> case C.toVector col of+ Right v -> Right v+ Left e ->+ Left $+ "Column '" <> name <> "': " <> T.pack (show e)
src/DataFrame/Typed/TH.hs view
@@ -10,21 +10,38 @@ deriveSchemaFromCsvFile, deriveSchemaFromCsvFileWith, + -- * ADT-based schema derivation+ deriveSchemaFromType,+ deriveSchemaFromTypeWith,+ SchemaOptions (..),+ defaultSchemaOptions,+ camelToSnake,+ -- * Re-export for TH splices TypedDataFrame, Column, ) where +import Control.Monad (when) import Control.Monad.IO.Class+import Data.Char (isUpper, toLower) import qualified Data.List as L import qualified Data.Map as M import qualified Data.Text as T+import qualified Data.Vector as VB import Language.Haskell.TH import qualified DataFrame.IO.CSV as D import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.DataFrame as D+import DataFrame.Typed.Record (+ HasSchema,+ Schema,+ fromColumns,+ requireColumn,+ toColumns,+ ) import DataFrame.Typed.Types (Column, TypedDataFrame) {- | Generate a type synonym for a schema based on an existing 'DataFrame'.@@ -96,3 +113,207 @@ findDuplicate (x : xs) | x `elem` xs = Just x | otherwise = findDuplicate xs++-- | Options controlling 'deriveSchemaFromTypeWith'.+data SchemaOptions = SchemaOptions+ { nameTransform :: String -> String+ -- ^ Map each record selector name to a column name. Default: 'camelToSnake'.+ , schemaTypeName :: Maybe String+ -- ^ Override the generated type synonym name. Default: @\<TypeName\>Schema@.+ , generateInstance :: Bool+ {- ^ When @True@ (default), also generate a 'HasSchema' instance so the+ record can be turned into a 'DataFrame' via 'fromRecords' / 'toRecords'.+ -}+ }++defaultSchemaOptions :: SchemaOptions+defaultSchemaOptions =+ SchemaOptions+ { nameTransform = camelToSnake+ , schemaTypeName = Nothing+ , generateInstance = True+ }++{- | Convert a camelCase identifier to snake_case.++>>> camelToSnake "orderId"+"order_id"+>>> camelToSnake "amountUS"+"amount_u_s"+>>> camelToSnake "region"+"region"+-}+camelToSnake :: String -> String+camelToSnake [] = []+camelToSnake (c : cs) = toLower c : go cs+ where+ go [] = []+ go (x : xs)+ | isUpper x = '_' : toLower x : go xs+ | otherwise = x : go xs++{- | Derive a schema type synonym and a 'HasSchema' instance from a single-record+ADT, using 'defaultSchemaOptions'.++@+data Order = Order+ { orderId :: Int64+ , region :: Text+ , amount :: Double+ } deriving (Show, Eq)++\$(deriveSchemaFromType ''Order)++-- expands to:+-- type OrderSchema =+-- '[Column "order_id" Int64, Column "region" Text, Column "amount" Double]+--+-- instance HasSchema Order OrderSchema where+-- toColumns rs = ...+-- fromColumns df = ...+@+-}+deriveSchemaFromType :: Name -> DecsQ+deriveSchemaFromType = deriveSchemaFromTypeWith defaultSchemaOptions++-- | Like 'deriveSchemaFromType' but accepts custom 'SchemaOptions'.+deriveSchemaFromTypeWith :: SchemaOptions -> Name -> DecsQ+deriveSchemaFromTypeWith opts tyName = do+ info <- reify tyName+ (conName, vbts) <- extractRecord tyName info+ when (Prelude.null vbts) $+ fail $+ "deriveSchemaFromType: record "+ ++ show tyName+ ++ " has no fields"+ let fields =+ [ (nameTransform opts (nameBase fName), fName, ty)+ | (fName, _bang, ty) <- vbts+ ]+ colNames = [c | (c, _, _) <- fields]+ case findDuplicate colNames of+ Just dup ->+ fail $+ "deriveSchemaFromType: duplicate transformed column name "+ ++ show dup+ ++ " (consider customizing nameTransform via deriveSchemaFromTypeWith)"+ Nothing -> pure ()+ let synName = case schemaTypeName opts of+ Just s -> mkName s+ Nothing -> mkName (nameBase tyName ++ "Schema")+ let columnTypes =+ [ ConT ''Column `AppT` LitT (StrTyLit colName) `AppT` ty+ | (colName, _, ty) <- fields+ ]+ schemaType =+ foldr+ (\t acc -> PromotedConsT `AppT` t `AppT` acc)+ PromotedNilT+ columnTypes+ synDec = TySynD synName [] schemaType+ if generateInstance opts+ then do+ inst <- mkHasSchemaInstance tyName schemaType conName fields+ pure [synDec, inst]+ else pure [synDec]++extractRecord :: Name -> Info -> Q (Name, [VarBangType])+extractRecord _ (TyConI dec) = case dec of+ DataD _ _ _ _ [RecC conName fs] _ -> pure (conName, fs)+ NewtypeD _ _ _ _ (RecC conName fs) _ -> pure (conName, fs)+ DataD _ name _ _ _ _ ->+ fail $+ "deriveSchemaFromType: "+ ++ show name+ ++ " must have exactly one record constructor"+ NewtypeD _ name _ _ _ _ ->+ fail $+ "deriveSchemaFromType: "+ ++ show name+ ++ " newtype must use record syntax"+ other ->+ fail $+ "deriveSchemaFromType: unsupported declaration: " ++ show other+extractRecord tyName _ =+ fail $+ "deriveSchemaFromType: " ++ show tyName ++ " is not a data/newtype declaration"++mkHasSchemaInstance ::+ Name -> Type -> Name -> [(String, Name, Type)] -> Q Dec+mkHasSchemaInstance tyName schemaType conName fields = do+ toClause <- mkToColumnsClause fields+ fromClause <- mkFromColumnsClause conName fields+ let instType = ConT ''HasSchema `AppT` ConT tyName+ schemaInst =+ TySynInstD+ (TySynEqn Nothing (ConT ''Schema `AppT` ConT tyName) schemaType)+ pure $+ InstanceD+ Nothing+ []+ instType+ [ schemaInst+ , FunD 'toColumns [toClause]+ , FunD 'fromColumns [fromClause]+ ]++mkToColumnsClause :: [(String, Name, Type)] -> Q Clause+mkToColumnsClause fields = do+ rs <- newName "rs"+ let mkPair (colName, fieldFn, _ty) =+ let nameE = AppE (VarE 'T.pack) (LitE (StringL colName))+ colE =+ AppE+ (VarE 'C.fromList)+ ( AppE+ (AppE (VarE 'map) (VarE fieldFn))+ (VarE rs)+ )+ in TupE [Just nameE, Just colE]+ listExp = ListE (map mkPair fields)+ pure $ Clause [VarP rs] (NormalB listExp) []++mkFromColumnsClause :: Name -> [(String, Name, Type)] -> Q Clause+mkFromColumnsClause conName fields = do+ df <- newName "df"+ iN <- newName "i"+ nN <- newName "n"+ vNames <- mapM (\k -> newName ("v" ++ show (k :: Int))) [0 .. length fields - 1]+ let mkBind v (colName, _, ty) =+ let nameE = AppE (VarE 'T.pack) (LitE (StringL colName))+ callE =+ AppE+ ( AppE+ (AppTypeE (VarE 'requireColumn) ty)+ nameE+ )+ (VarE df)+ in BindS (VarP v) callE+ binds = zipWith mkBind vNames fields+ firstV = case vNames of+ (v0 : _) -> v0+ [] -> error "mkFromColumnsClause: empty fields (should have failed earlier)"+ lengthBind =+ LetS+ [ ValD+ (VarP nN)+ (NormalB (AppE (VarE 'VB.length) (VarE firstV)))+ []+ ]+ indexE v =+ AppE (AppE (VarE 'VB.unsafeIndex) (VarE v)) (VarE iN)+ elemE =+ foldl+ (\acc v -> AppE acc (indexE v))+ (ConE conName)+ vNames+ nMinus1E =+ InfixE+ (Just (VarE nN))+ (VarE '(-))+ (Just (LitE (IntegerL 1)))+ rangeE = ArithSeqE (FromToR (LitE (IntegerL 0)) nMinus1E)+ compExp = CompE [BindS (VarP iN) rangeE, NoBindS elemE]+ rightE = AppE (ConE 'Right) compExp+ body = DoE Nothing (binds ++ [lengthBind, NoBindS rightE])+ pure $ Clause [VarP df] (NormalB body) []
src/DataFrame/Typed/Types.hs view
@@ -80,18 +80,15 @@ newtype TypedGrouped (keys :: [Symbol]) (cols :: [Type]) = TGD {unTGD :: D.GroupedDataFrame} -{- | A typed aggregation builder (Option B).--Accumulates 'NamedExpr' values at the term level while building-the result schema at the type level. Each @agg@ call prepends a-'Column' to the @aggs@ phantom list.--Usage:+{- | Internal aggregation chain. Each cons prepends a 'Column' to the+@aggs@ phantom list. End users never construct this directly — they+compose 'DataFrame.Typed.Aggregate.as' entries with @(.)@ and let+'DataFrame.Typed.Aggregate.aggregate' apply the composition to+'TAggNil'. @-agg \@\"total\" (F.sum salary)- $ agg \@\"avg_age\" (F.mean age)- $ aggNil+as \@\"total\" (F.sum salary)+ . as \@\"avg_age\" (F.mean age) @ -} data TAgg (keys :: [Symbol]) (cols :: [Type]) (aggs :: [Type]) where
tests/DecisionTree.hs view
@@ -1,7 +1,6 @@ {-# LANGUAGE LambdaCase #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TypeApplications #-}-{-# OPTIONS_GHC -Wno-x-partial #-} module DecisionTree where
tests/Main.hs view
@@ -27,6 +27,7 @@ import qualified Operations.Nullable import qualified Operations.Provenance import qualified Operations.ReadCsv+import qualified Operations.Record import qualified Operations.Shuffle import qualified Operations.Sort import qualified Operations.Statistics@@ -55,6 +56,7 @@ ++ Operations.Nullable.tests ++ Operations.Provenance.tests ++ Operations.ReadCsv.tests+ ++ Operations.Record.tests ++ Operations.WriteCsv.tests ++ Operations.Shuffle.tests ++ Operations.Sort.tests
tests/Operations/Aggregations.hs view
@@ -82,7 +82,7 @@ , DT.Column "test6" Integer ] & DT.groupBy @'["test1"]- & DT.aggregate (DT.agg @"n" DT.countAll DT.aggNil)+ & DT.aggregate (DT.as @"n" DT.countAll) & DT.sortBy [DT.asc (DT.col @"test1")] & DT.thaw )@@ -109,7 +109,7 @@ , DT.Column "test6" Integer ] & DT.groupBy @'["test1"]- & DT.aggregate (DT.agg @"test2_count" (DT.count (DT.col @"test2")) DT.aggNil)+ & DT.aggregate (DT.as @"test2_count" (DT.count (DT.col @"test2"))) & DT.sortBy [DT.asc (DT.col @"test1")] & DT.thaw )@@ -153,7 +153,7 @@ , DT.Column "test6" Integer ] & DT.groupBy @'["test1"]- & DT.aggregate (DT.agg @"test2_mean" (DT.mean (DT.col @"test2")) DT.aggNil)+ & DT.aggregate (DT.as @"test2_mean" (DT.mean (DT.col @"test2"))) & DT.sortBy [DT.asc (DT.col @"test1")] & DT.thaw )
+ tests/Operations/Record.hs view
@@ -0,0 +1,358 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE UndecidableInstances #-}++module Operations.Record where++import Data.Int (Int64)+import qualified Data.Map.Strict as M+import qualified Data.Text as T+import qualified Data.Text.IO as TIO+import GHC.Generics (Generic)++import qualified DataFrame as D+import qualified DataFrame.Functions as F+import qualified DataFrame.Internal.Column as DI+import qualified DataFrame.Internal.Schema as IS+import DataFrame.Operators+import DataFrame.Typed (Schema)+import qualified DataFrame.Typed as DT++import Test.HUnit++-- Basic record with three different-type fields (Int64 / Text / Double).+data Order = Order+ { orderId :: Int64+ , region :: T.Text+ , amount :: Double+ }+ deriving (Show, Eq)++$(DT.deriveSchemaFromType ''Order)+$(D.deriveSchema ''Order)++-- Nullable fields (Maybe Text -> RNullableBoxed; Maybe Int -> RNullableUnboxed).+data User = User+ { userId :: Int64+ , userName :: Maybe T.Text+ , userAge :: Maybe Int+ }+ deriving (Show, Eq)++$(DT.deriveSchemaFromType ''User)+$(D.deriveSchema ''User)++-- Identity-cased: keep the record selector names verbatim.+data Account = Account+ { accountId :: Int64+ , accountName :: T.Text+ }+ deriving (Show, Eq)++$( DT.deriveSchemaFromTypeWith+ DT.defaultSchemaOptions{DT.nameTransform = id}+ ''Account+ )++-- Schema-name override.+data Item = Item+ { itemId :: Int64+ , itemPrice :: Double+ }+ deriving (Show, Eq)++$( DT.deriveSchemaFromTypeWith+ DT.defaultSchemaOptions{DT.schemaTypeName = Just "ItemsSchema"}+ ''Item+ )++-- Wide record (more than zipWith3 can handle).+data Wide = Wide+ { f1 :: Int+ , f2 :: Int+ , f3 :: Int+ , f4 :: Int+ , f5 :: Int+ , f6 :: Int+ , f7 :: Int+ , f8 :: Int+ }+ deriving (Show, Eq)++$(DT.deriveSchemaFromType ''Wide)+$(D.deriveSchema ''Wide)++-- Generics opt-in: derive the schema via Generic, not TH.+data Foo = Foo+ { fooId :: Int64+ , fooName :: T.Text+ , fooValue :: Double+ }+ deriving (Show, Eq, Generic)++type FooSchema = DT.SchemaOf Foo++instance DT.HasSchema Foo where+ type Schema Foo = FooSchema+ toColumns = DT.genericToColumns+ fromColumns = DT.genericFromColumns++orderSample :: [Order]+orderSample =+ [ Order 1 "us" 10.0+ , Order 2 "eu" 20.5+ , Order 3 "ap" 30.0+ ]++basicTypedRoundTrip :: Test+basicTypedRoundTrip = TestCase $ do+ let df :: DT.TypedDataFrame OrderSchema+ df = DT.fromRecordsTyped orderSample+ case DT.toRecordsTyped df of+ Left e -> assertFailure (T.unpack e)+ Right xs -> assertEqual "typed round-trip" orderSample xs++basicUntypedRoundTrip :: Test+basicUntypedRoundTrip = TestCase $ do+ let df = D.fromRecords orderSample+ case D.toRecords df :: Either T.Text [Order] of+ Left e -> assertFailure (T.unpack e)+ Right xs -> assertEqual "untyped round-trip" orderSample xs++emptyRoundTrip :: Test+emptyRoundTrip = TestCase $ do+ let df = D.fromRecords ([] :: [Order])+ case D.toRecords df :: Either T.Text [Order] of+ Left e -> assertFailure (T.unpack e)+ Right xs -> assertEqual "empty round-trip" [] xs++nullableRoundTrip :: Test+nullableRoundTrip = TestCase $ do+ let xs =+ [ User 1 (Just "alice") (Just 30)+ , User 2 Nothing (Just 40)+ , User 3 (Just "carol") Nothing+ , User 4 Nothing Nothing+ ]+ df = D.fromRecords xs+ case D.toRecords df :: Either T.Text [User] of+ Left e -> assertFailure (T.unpack e)+ Right ys -> assertEqual "nullable round-trip" xs ys++identityRoundTrip :: Test+identityRoundTrip = TestCase $ do+ let xs = [Account 1 "alice", Account 2 "bob"]+ df = D.fromRecords xs+ assertEqual+ "identity column names"+ ["accountId", "accountName"]+ (D.columnNames df)+ case D.toRecords df :: Either T.Text [Account] of+ Left e -> assertFailure (T.unpack e)+ Right ys -> assertEqual "identity round-trip" xs ys++schemaNameOverride :: Test+schemaNameOverride = TestCase $ do+ let xs = [Item 1 9.99, Item 2 19.99]+ df :: DT.TypedDataFrame ItemsSchema+ df = DT.fromRecordsTyped xs+ case DT.toRecordsTyped df of+ Left e -> assertFailure (T.unpack e)+ Right ys -> assertEqual "schema-name override round-trip" xs ys++typeMismatchError :: Test+typeMismatchError = TestCase $ do+ let badDf =+ D.fromNamedColumns+ [ ("order_id", DI.fromList ([1, 2, 3] :: [Int])) -- wrong: Int, not Int64+ , ("region", DI.fromList (["us", "eu", "ap"] :: [T.Text]))+ , ("amount", DI.fromList ([10.0, 20.5, 30.0] :: [Double]))+ ]+ case D.toRecords badDf :: Either T.Text [Order] of+ Right _ -> assertFailure "expected error on type mismatch"+ Left e ->+ assertBool+ ("error should mention order_id, got: " ++ T.unpack e)+ ("order_id" `T.isInfixOf` e)++missingColumnError :: Test+missingColumnError = TestCase $ do+ let badDf =+ D.fromNamedColumns+ [ ("region", DI.fromList (["us"] :: [T.Text]))+ , ("amount", DI.fromList ([10.0] :: [Double]))+ ]+ case D.toRecords badDf :: Either T.Text [Order] of+ Right _ -> assertFailure "expected error on missing column"+ Left e ->+ assertBool+ ("error should mention order_id, got: " ++ T.unpack e)+ ("order_id" `T.isInfixOf` e)++wideRoundTrip :: Test+wideRoundTrip = TestCase $ do+ let xs =+ [ Wide 1 2 3 4 5 6 7 8+ , Wide 9 10 11 12 13 14 15 16+ ]+ df = D.fromRecords xs+ case D.toRecords df :: Either T.Text [Wide] of+ Left e -> assertFailure (T.unpack e)+ Right ys -> assertEqual "wide round-trip" xs ys++genericRoundTrip :: Test+genericRoundTrip = TestCase $ do+ let xs = [Foo 1 "a" 1.0, Foo 2 "b" 2.5]+ df = D.fromRecords xs+ case D.toRecords df :: Either T.Text [Foo] of+ Left e -> assertFailure (T.unpack e)+ Right ys -> assertEqual "generic round-trip" xs ys++genericColumnNames :: Test+genericColumnNames = TestCase $ do+ let df = D.fromRecords [Foo 1 "a" 1.0]+ assertEqual+ "generic snake-cased column names"+ ["foo_id", "foo_name", "foo_value"]+ (D.columnNames df)++deriveSchemaSplice :: Test+deriveSchemaSplice = TestCase $ do+ assertEqual+ "orderSchema column names"+ ["amount", "order_id", "region"]+ (M.keys (IS.elements orderSchema))+ assertEqual+ "order_id is Int64"+ (Just (IS.schemaType @Int64))+ (M.lookup "order_id" (IS.elements orderSchema))+ assertEqual+ "region is Text"+ (Just (IS.schemaType @T.Text))+ (M.lookup "region" (IS.elements orderSchema))+ assertEqual+ "amount is Double"+ (Just (IS.schemaType @Double))+ (M.lookup "amount" (IS.elements orderSchema))++deriveSchemaNullable :: Test+deriveSchemaNullable = TestCase $ do+ assertEqual+ "userSchema column names"+ ["user_age", "user_id", "user_name"]+ (M.keys (IS.elements userSchema))+ assertEqual+ "user_id is Int64"+ (Just (IS.schemaType @Int64))+ (M.lookup "user_id" (IS.elements userSchema))+ assertEqual+ "user_name is Maybe Text"+ (Just (IS.schemaType @(Maybe T.Text)))+ (M.lookup "user_name" (IS.elements userSchema))+ assertEqual+ "user_age is Maybe Int"+ (Just (IS.schemaType @(Maybe Int)))+ (M.lookup "user_age" (IS.elements userSchema))++deriveSchemaWide :: Test+deriveSchemaWide = TestCase $ do+ assertEqual+ "wideSchema has 8 keys"+ 8+ (M.size (IS.elements wideSchema))+ assertEqual+ "f1 is Int"+ (Just (IS.schemaType @Int))+ (M.lookup "f1" (IS.elements wideSchema))+ assertEqual+ "f8 is Int"+ (Just (IS.schemaType @Int))+ (M.lookup "f8" (IS.elements wideSchema))++deriveSchemaReadsCsv :: Test+deriveSchemaReadsCsv = TestCase $ do+ let csv =+ T.unlines+ [ "order_id,region,amount"+ , "1,us,10.0"+ , "2,eu,20.5"+ , "3,ap,30.0"+ ]+ tmp = "/tmp/dataframe_test_deriveSchema.csv"+ TIO.writeFile tmp csv+ df <- D.readCsvWithSchema orderSchema tmp+ assertEqual+ "deriveSchema-driven readCsvWithSchema column names"+ ["order_id", "region", "amount"]+ (D.columnNames df)+ case D.toRecords df :: Either T.Text [Order] of+ Left e -> assertFailure (T.unpack e)+ Right xs ->+ assertEqual+ "deriveSchema-driven CSV parses back to records"+ [Order 1 "us" 10.0, Order 2 "eu" 20.5, Order 3 "ap" 30.0]+ xs++deriveSchemaAccessorFilter :: Test+deriveSchemaAccessorFilter = TestCase $ do+ let df =+ D.fromRecords+ [ Order 1 "us" 20.0+ , Order 2 "eu" 20.5+ , Order 3 "ap" 30.0+ , Order 4 "us" 25.0+ ]+ big =+ D.filterWhere+ (orderAmount .>. F.lit @Double 15.0 .&&. orderRegion .==. F.lit @T.Text "us")+ df+ case D.toRecords big :: Either T.Text [Order] of+ Left e -> assertFailure (T.unpack e)+ Right xs ->+ assertEqual+ "accessor drives D.filter (amount > 15.0 && region == \"us\")"+ [Order 1 "us" 20.0, Order 4 "us" 25.0]+ xs++deriveSchemaAccessorDerive :: Test+deriveSchemaAccessorDerive = TestCase $ do+ let df =+ D.fromRecords+ [ Order 1 "us" 10.0+ , Order 2 "eu" 20.0+ ]+ df' = D.derive "double_amount" (orderAmount + orderAmount) df+ assertEqual+ "accessor composes in derive expression"+ [20.0, 40.0]+ (D.columnAsList (D.col @Double "double_amount") df')++tests :: [Test]+tests =+ [ TestLabel "basicTypedRoundTrip" basicTypedRoundTrip+ , TestLabel "basicUntypedRoundTrip" basicUntypedRoundTrip+ , TestLabel "emptyRoundTrip" emptyRoundTrip+ , TestLabel "nullableRoundTrip" nullableRoundTrip+ , TestLabel "identityRoundTrip" identityRoundTrip+ , TestLabel "schemaNameOverride" schemaNameOverride+ , TestLabel "typeMismatchError" typeMismatchError+ , TestLabel "missingColumnError" missingColumnError+ , TestLabel "wideRoundTrip" wideRoundTrip+ , TestLabel "genericRoundTrip" genericRoundTrip+ , TestLabel "genericColumnNames" genericColumnNames+ , TestLabel "deriveSchemaSplice" deriveSchemaSplice+ , TestLabel "deriveSchemaNullable" deriveSchemaNullable+ , TestLabel "deriveSchemaWide" deriveSchemaWide+ , TestLabel "deriveSchemaReadsCsv" deriveSchemaReadsCsv+ , TestLabel "deriveSchemaAccessorFilter" deriveSchemaAccessorFilter+ , TestLabel "deriveSchemaAccessorDerive" deriveSchemaAccessorDerive+ ]
tests/Operations/Sort.hs view
@@ -89,6 +89,72 @@ (print $ D.sortBy [D.Asc (F.col @Int "test0")] testData) ) +compoundTestData :: D.DataFrame+compoundTestData =+ D.fromNamedColumns+ [ ("a", DI.fromList ([3, 1, 4, 1, 5] :: [Int]))+ , ("b", DI.fromList ([10, 40, 20, 30, 50] :: [Int]))+ ]++sortByCompoundExpression :: Test+sortByCompoundExpression =+ TestCase+ ( assertEqual+ "Sorting by Asc (a + b) orders rows by the sum without leaking a synthetic column"+ ( D.fromNamedColumns+ [ ("a", DI.fromList ([3, 4, 1, 1, 5] :: [Int]))+ , ("b", DI.fromList ([10, 20, 30, 40, 50] :: [Int]))+ ]+ )+ (D.sortBy [D.Asc (F.col @Int "a" + F.col @Int "b")] compoundTestData)+ )++sortByCompoundExpressionDescending :: Test+sortByCompoundExpressionDescending =+ TestCase+ ( assertEqual+ "Sorting by Desc (b - a) orders rows by descending difference"+ ( D.fromNamedColumns+ [ ("a", DI.fromList ([5, 1, 1, 4, 3] :: [Int]))+ , ("b", DI.fromList ([50, 40, 30, 20, 10] :: [Int]))+ ]+ )+ (D.sortBy [D.Desc (F.col @Int "b" - F.col @Int "a")] compoundTestData)+ )++sortByCompoundMixedWithBareColumn :: Test+sortByCompoundMixedWithBareColumn =+ TestCase+ ( assertEqual+ "Mixing a compound Asc key with a bare Desc tie-breaker works"+ ( D.fromNamedColumns+ [ ("a", DI.fromList ([1, 1, 3, 4, 5] :: [Int]))+ , ("b", DI.fromList ([40, 30, 10, 20, 50] :: [Int]))+ ]+ )+ ( D.sortBy+ [D.Asc (F.col @Int "a" * 2), D.Desc (F.col @Int "b")]+ compoundTestData+ )+ )++sortByCompoundMissingColumn :: Test+sortByCompoundMissingColumn =+ TestCase+ ( assertExpectException+ "[Error Case]"+ ( D.columnsNotFound+ ["nope"]+ "sortBy"+ (D.columnNames compoundTestData)+ )+ ( print $+ D.sortBy+ [D.Asc (F.col @Int "nope" + F.col @Int "a")]+ compoundTestData+ )+ )+ tests :: [Test] tests = [ TestLabel "sortByAscendingWAI" sortByAscendingWAI@@ -96,4 +162,10 @@ , TestLabel "sortByColumnDoesNotExist" sortByColumnDoesNotExist , TestLabel "sortByTwoColumns" sortByTwoColumns , TestLabel "sortByOneColumnAscOneColumnDesc" sortByOneColumnAscOneColumnDesc+ , TestLabel "sortByCompoundExpression" sortByCompoundExpression+ , TestLabel+ "sortByCompoundExpressionDescending"+ sortByCompoundExpressionDescending+ , TestLabel "sortByCompoundMixedWithBareColumn" sortByCompoundMixedWithBareColumn+ , TestLabel "sortByCompoundMissingColumn" sortByCompoundMissingColumn ]
tests/Parquet.hs view
@@ -1,3 +1,4 @@+{-# LANGUAGE OverloadedRecordDot #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TypeApplications #-} @@ -8,22 +9,25 @@ import qualified DataFrame as D import qualified DataFrame.Functions as F import qualified DataFrame.IO.Parquet as DP+import qualified DataFrame.IO.Parquet.Schema as PS import ParquetTestData (allTypes, mtCarsDataset, tinyPagesLast10, transactions) import qualified Data.ByteString as BS-import Data.Int+import Data.Int (Int32, Int64)+import Data.Maybe (fromMaybe) import qualified Data.Set as S import qualified Data.Text as T import Data.Word import DataFrame.IO.Parquet.Thrift (- columnMetaData,- columnPathInSchema,- columnStatistics,- rowGroupColumns,- rowGroups,+ cc_meta_data,+ cmd_path_in_schema,+ cmd_statistics,+ rg_columns,+ row_groups, schema,+ stats_null_count,+ unField, )-import DataFrame.IO.Parquet.Types (columnNullCount) import DataFrame.Internal.Binary ( littleEndianWord32, littleEndianWord64,@@ -370,6 +374,11 @@ -- Group 2: Compression codecs (unsupported → error tests) -- --------------------------------------------------------------------------- +-- TODO: LZ4 and LZ4_RAW compression are not yet implemented. When support+-- is added via a Haskell lz4 binding, hadoopLz4Compressed,+-- hadoopLz4CompressedLarger, nonHadoopLz4Compressed, lz4RawCompressed, and+-- lz4RawCompressedLarger should all change from assertExpectException to+-- assertEqual checking their respective row/column dimensions. hadoopLz4Compressed :: Test hadoopLz4Compressed = TestCase@@ -415,15 +424,26 @@ (D.readParquet "./tests/data/lz4_raw_compressed_larger.parquet") ) +-- Was: assertExpectException "concatenatedGzipMembers" "12" ...+-- The old parser failed with a ZLIB size error. The new decompressor+-- handles concatenated gzip members correctly. concatenatedGzipMembers :: Test concatenatedGzipMembers = TestCase- ( assertExpectException+ ( assertEqual "concatenatedGzipMembers"- "12"- (D.readParquet "./tests/data/concatenated_gzip_members.parquet")+ (513, 1)+ ( unsafePerformIO+ ( fmap+ D.dimensions+ (D.readParquet "./tests/data/concatenated_gzip_members.parquet")+ )+ ) ) +-- TODO: BROTLI compression is not yet implemented. When a Haskell brotli+-- binding is added, change this to assertEqual checking the actual+-- dimensions of large_string_map.brotli.parquet. largeBrotliMap :: Test largeBrotliMap = TestCase@@ -437,66 +457,114 @@ -- Group 3: Delta / RLE encodings (unsupported → error tests) -- --------------------------------------------------------------------------- +-- Was: assertExpectException "deltaBinaryPacked" "EDELTA_BINARY_PACKED" ...+-- The new parser's error includes the encoding name "DELTA_BINARY_PACKED"+-- without the old "E" prefix used in the previous error format.+-- TODO: When DELTA_BINARY_PACKED (encoding id=5) is implemented, change+-- this to assertEqual checking actual dimensions. The encoding stores+-- integer data as bit-packed deltas and is common for monotonically+-- increasing columns (row IDs, timestamps):+-- https://parquet.apache.org/docs/file-format/data-pages/encodings/#delta-encoding-delta_binary_packed--5 deltaBinaryPacked :: Test deltaBinaryPacked = TestCase ( assertExpectException "deltaBinaryPacked"- "EDELTA_BINARY_PACKED"+ "DELTA_BINARY_PACKED" (D.readParquet "./tests/data/delta_binary_packed.parquet") ) +-- Was: assertExpectException "deltaByteArray" "EDELTA_BYTE_ARRAY" ...+-- Same reason as deltaBinaryPacked: new error format drops the "E" prefix.+-- TODO: When DELTA_BYTE_ARRAY (encoding id=7) is implemented, change this+-- to assertEqual checking actual dimensions. The encoding prefix-differences+-- consecutive string values, reducing storage for sorted byte arrays:+-- https://parquet.apache.org/docs/file-format/data-pages/encodings/#delta-strings-delta_byte_array--7 deltaByteArray :: Test deltaByteArray = TestCase ( assertExpectException "deltaByteArray"- "EDELTA_BYTE_ARRAY"+ "DELTA_BYTE_ARRAY" (D.readParquet "./tests/data/delta_byte_array.parquet") ) +-- Was: assertExpectException "deltaEncodingOptionalColumn" "EDELTA_BINARY_PACKED" ...+-- The first column that errors in this file uses DELTA_BYTE_ARRAY encoding,+-- so we match the broader "unsupported encoding" substring instead.+-- TODO: Once DELTA_BINARY_PACKED and DELTA_BYTE_ARRAY are both implemented,+-- change this to assertEqual checking the actual row count of+-- delta_encoding_optional_column.parquet. deltaEncodingOptionalColumn :: Test deltaEncodingOptionalColumn = TestCase ( assertExpectException "deltaEncodingOptionalColumn"- "EDELTA_BINARY_PACKED"+ "unsupported encoding" (D.readParquet "./tests/data/delta_encoding_optional_column.parquet") ) +-- Was: assertExpectException "deltaEncodingRequiredColumn" "EDELTA_BINARY_PACKED" ...+-- Same as deltaEncodingOptionalColumn: first failing column uses DELTA_BYTE_ARRAY.+-- TODO: Same as deltaEncodingOptionalColumn — change to assertEqual once+-- DELTA_BINARY_PACKED and DELTA_BYTE_ARRAY encodings are both supported. deltaEncodingRequiredColumn :: Test deltaEncodingRequiredColumn = TestCase ( assertExpectException "deltaEncodingRequiredColumn"- "EDELTA_BINARY_PACKED"+ "unsupported encoding" (D.readParquet "./tests/data/delta_encoding_required_column.parquet") ) +-- Was: assertExpectException "deltaLengthByteArray" "ZSTD" ...+-- The old parser failed during ZSTD decompression. The new parser+-- detects the unsupported DELTA_LENGTH_BYTE_ARRAY encoding before decompression.+-- TODO: When DELTA_LENGTH_BYTE_ARRAY (encoding id=6) is implemented, change+-- this to assertEqual checking actual dimensions. The encoding stores a+-- delta-encoded list of byte-array lengths followed by the raw concatenated+-- values:+-- https://parquet.apache.org/docs/file-format/data-pages/encodings/#delta-length-byte-array-delta_length_byte_array--6 deltaLengthByteArray :: Test deltaLengthByteArray = TestCase ( assertExpectException "deltaLengthByteArray"- "ZSTD"+ "DELTA_LENGTH_BYTE_ARRAY" (D.readParquet "./tests/data/delta_length_byte_array.parquet") ) +-- Was: assertExpectException "rleBooleanEncoding" "Zlib" ...+-- The old parser failed during Zlib decompression. The new parser+-- detects the unsupported RLE boolean encoding before reaching decompression.+-- TODO: When RLE/Bit-Packing Hybrid (encoding id=3, bit-width=1) is+-- implemented for BOOLEAN columns, change this to assertEqual checking the+-- actual decoded boolean values. The encoding is spec-valid for BOOLEAN:+-- https://parquet.apache.org/docs/file-format/data-pages/encodings/#run-length-encoding--bit-packing-hybrid-rle--3 rleBooleanEncoding :: Test rleBooleanEncoding = TestCase ( assertExpectException "rleBooleanEncoding"- "Zlib"+ "unsupported encoding RLE" (D.readParquet "./tests/data/rle_boolean_encoding.parquet") ) +-- Was: assertExpectException "dictPageOffsetZero" "Unknown kv" ...+-- The old parser reported "Unknown kv" for a bad key-value field. The new+-- Pinch-based page-header parser reports "Field 1 is absent" for the+-- malformed page header in this file.+-- TODO: Investigate whether dict-page-offset-zero.parquet can be read+-- successfully with a more lenient page-header parser. If the missing+-- mandatory field can be treated as a per-page soft error rather than+-- aborting the whole read, this test would change to assertEqual+-- checking actual dimensions. dictPageOffsetZero :: Test dictPageOffsetZero = TestCase ( assertExpectException "dictPageOffsetZero"- "Unknown kv"+ "Field 1 is absent" (D.readParquet "./tests/data/dict-page-offset-zero.parquet") ) @@ -504,31 +572,64 @@ -- Group 4: Data Page V2 (unsupported → error tests) -- --------------------------------------------------------------------------- +-- Was: assertExpectException "datapageV2Snappy" "InvalidOffset" ...+-- The old parser failed with an offset validation error. The new parser+-- first encounters the unsupported RLE encoding used by data-page-v2.+-- TODO: Full Data Page V2 support requires two changes:+-- 1. RLE/Bit-Packing Hybrid (id=3, bit-width=1) for BOOLEAN values+-- (shared with rleBooleanEncoding above).+-- 2. Parsing DataPageHeaderV2's in-line level streams: in v2, definition+-- and repetition levels are stored uncompressed before the (optionally+-- compressed) value bytes, with lengths given by+-- definition_levels_byte_length and repetition_levels_byte_length.+-- Once both are done, change to assertEqual checking actual dimensions:+-- https://parquet.apache.org/docs/file-format/data-pages/ datapageV2Snappy :: Test datapageV2Snappy = TestCase ( assertExpectException "datapageV2Snappy"- "InvalidOffset"+ "unsupported encoding RLE" (D.readParquet "./tests/data/datapage_v2.snappy.parquet") ) +-- Was: assertExpectException "datapageV2EmptyDatapage" "UnexpectedEOF" ...+-- The old Snappy decompressor raised "UnexpectedEOF". The new Snappy+-- library raises "EmptyInput" when given zero-length compressed data.+-- The v2 page structure is parsed correctly: readLevelsV2V strips the+-- in-line level streams before decompression, leaving an empty value+-- payload (BS.empty) for a page with 0 values. The Snappy decompressor+-- then raises "EmptyInput" because it is handed zero bytes.+-- TODO: An empty data page (0 values) is valid and should contribute+-- 0 rows without raising an error. The fix is a single guard in the+-- DATA_PAGE_V2 branch of readPages (Page.hs): short-circuit+-- decompressData when compValBytes is empty, returning BS.empty+-- directly. Once fixed, change this to assertEqual checking the+-- total expected row count of the file. datapageV2EmptyDatapage :: Test datapageV2EmptyDatapage = TestCase ( assertExpectException "datapageV2EmptyDatapage"- "UnexpectedEOF"+ "EmptyInput" (D.readParquet "./tests/data/datapage_v2_empty_datapage.snappy.parquet") ) +-- Was: assertExpectException "pageV2EmptyCompressed" "10" ...+-- The old parser failed on empty compressed page-v2 blocks. The new parser+-- treats empty compressed data as zero-value pages and reads all 10 rows. pageV2EmptyCompressed :: Test pageV2EmptyCompressed = TestCase- ( assertExpectException+ ( assertEqual "pageV2EmptyCompressed"- "10"- (D.readParquet "./tests/data/page_v2_empty_compressed.parquet")+ (10, 1)+ ( unsafePerformIO+ ( fmap+ D.dimensions+ (D.readParquet "./tests/data/page_v2_empty_compressed.parquet")+ )+ ) ) -- ---------------------------------------------------------------------------@@ -591,6 +692,12 @@ ) ) +-- TODO: CRC checksum validation is not yet implemented; corrupt page+-- checksums are silently ignored. When validation is added, consider a+-- validateChecksums :: Bool field in ParquetReadOptions (default False)+-- so callers can opt in. Once implemented, datapageV1CorruptChecksum and+-- rleDictUncompressedCorruptChecksum should change to assertExpectException+-- checking for a checksum mismatch error. datapageV1CorruptChecksum :: Test datapageV1CorruptChecksum = TestCase@@ -726,22 +833,44 @@ ) ) +-- Was: assertExpectException "fixedLengthDecimal" "FIXED_LEN_BYTE_ARRAY" ...+-- The old parser recognised FIXED_LEN_BYTE_ARRAY as a physical type but+-- had no page decoder for it; reading data from such a column threw an+-- error at the decoding stage. The new parser's fixedLenByteArrayDecoder+-- reads the raw bytes and surfaces them as a text column.+-- TODO: When the DECIMAL logical type is properly decoded for+-- FIXED_LEN_BYTE_ARRAY columns, replace this dimension-only check with a+-- value-level assertion verifying the actual decimal values (e.g. as+-- Scientific or Double). The raw-byte Text column should become a typed+-- numeric column. fixedLengthDecimal :: Test fixedLengthDecimal = TestCase- ( assertExpectException+ ( assertEqual "fixedLengthDecimal"- "FIXED_LEN_BYTE_ARRAY"- (D.readParquet "./tests/data/fixed_length_decimal.parquet")+ (24, 1)+ ( unsafePerformIO+ (fmap D.dimensions (D.readParquet "./tests/data/fixed_length_decimal.parquet"))+ ) ) +-- Was: assertExpectException "fixedLengthDecimalLegacy" "FIXED_LEN_BYTE_ARRAY" ...+-- Same as fixedLengthDecimal: the old parser had no page decoder for+-- FIXED_LEN_BYTE_ARRAY; the new parser's fixedLenByteArrayDecoder handles it.+-- TODO: Same as fixedLengthDecimal — add a value-level assertion once+-- DECIMAL decoding over FIXED_LEN_BYTE_ARRAY is implemented. fixedLengthDecimalLegacy :: Test fixedLengthDecimalLegacy = TestCase- ( assertExpectException+ ( assertEqual "fixedLengthDecimalLegacy"- "FIXED_LEN_BYTE_ARRAY"- (D.readParquet "./tests/data/fixed_length_decimal_legacy.parquet")+ (24, 1)+ ( unsafePerformIO+ ( fmap+ D.dimensions+ (D.readParquet "./tests/data/fixed_length_decimal_legacy.parquet")+ )+ ) ) -- ---------------------------------------------------------------------------@@ -773,13 +902,18 @@ ) ) +-- Was: assertExpectException "fixedLengthByteArray" "FIXED_LEN_BYTE_ARRAY" ...+-- Same as fixedLengthDecimal: the old parser had no page decoder for+-- FIXED_LEN_BYTE_ARRAY; the new parser's fixedLenByteArrayDecoder handles it. fixedLengthByteArray :: Test fixedLengthByteArray = TestCase- ( assertExpectException+ ( assertEqual "fixedLengthByteArray"- "FIXED_LEN_BYTE_ARRAY"- (D.readParquet "./tests/data/fixed_length_byte_array.parquet")+ (1000, 1)+ ( unsafePerformIO+ (fmap D.dimensions (D.readParquet "./tests/data/fixed_length_byte_array.parquet"))+ ) ) -- ---------------------------------------------------------------------------@@ -801,13 +935,21 @@ -- Group 10: Metadata / index / bloom filters -- --------------------------------------------------------------------------- +-- Was: assertExpectException "columnChunkKeyValueMetadata" "Unknown page header field" ...+-- The old parser rejected extra fields in page headers. Pinch ignores+-- unknown fields gracefully. This file contains 0 data rows. columnChunkKeyValueMetadata :: Test columnChunkKeyValueMetadata = TestCase- ( assertExpectException+ ( assertEqual "columnChunkKeyValueMetadata"- "Unknown page header field"- (D.readParquet "./tests/data/column_chunk_key_value_metadata.parquet")+ (0, 2)+ ( unsafePerformIO+ ( fmap+ D.dimensions+ (D.readParquet "./tests/data/column_chunk_key_value_metadata.parquet")+ )+ ) ) dataIndexBloomEncodingStats :: Test@@ -838,64 +980,117 @@ ) ) +-- Was: assertEqual "sortColumns" (3, 2) ...+-- The file contains two row groups, each storing 3 rows (6 rows total).+-- DuckDB's parquet-metadata output shows row_group_num_rows=3, which is+-- the count *per row group*, not the file total.row group*, not the file total.row group*, not the file total.row group*, not the file total.+-- https://github.com/apache/parquet-testing/blob/master/data/README.md#:~:text=sort_columns.parquet+-- The above link is to the repository the test parquet files comes from.+-- The table describes sort_columns.parquet as having two row groups.+-- The old parser only read the first row group (a bug). The new parser+-- reads all row groups and returns (6, 2) correctly. sortColumns :: Test sortColumns = TestCase ( assertEqual "sortColumns"- (3, 2)+ (6, 2) ( unsafePerformIO (fmap D.dimensions (D.readParquet "./tests/data/sort_columns.parquet")) ) ) +-- Was: assertExpectException "overflowI16PageCnt" "UNIMPLEMENTED" ...+-- The old parser used Int16 for page counts and overflowed on this file.+-- The new parser uses Int32 and reads all 40,000 rows correctly. overflowI16PageCnt :: Test overflowI16PageCnt = TestCase- ( assertExpectException+ ( assertEqual "overflowI16PageCnt"- "UNIMPLEMENTED"- (D.readParquet "./tests/data/overflow_i16_page_cnt.parquet")+ (40000, 1)+ ( unsafePerformIO+ (fmap D.dimensions (D.readParquet "./tests/data/overflow_i16_page_cnt.parquet"))+ ) ) -- --------------------------------------------------------------------------- -- Group 11: Nested / complex types and byte-stream-split -- --------------------------------------------------------------------------- +-- Was: assertExpectException "byteStreamSplitZstd" "EBYTE_STREAM_SPLIT" ...+-- The new parser's error includes the encoding name "BYTE_STREAM_SPLIT"+-- without the old "E" prefix used in the previous error format.+-- TODO: When BYTE_STREAM_SPLIT (encoding id=9) is implemented, change this+-- to assertEqual checking actual dimensions. The encoding interleaves the+-- individual byte streams of multi-byte scalars to improve compression for+-- floating-point and other structured data:+-- https://parquet.apache.org/docs/file-format/data-pages/encodings/#byte-stream-split-byte_stream_split--9 byteStreamSplitZstd :: Test byteStreamSplitZstd = TestCase ( assertExpectException "byteStreamSplitZstd"- "EBYTE_STREAM_SPLIT"+ "BYTE_STREAM_SPLIT" (D.readParquet "./tests/data/byte_stream_split.zstd.parquet") ) +-- Was: assertExpectException "byteStreamSplitExtendedGzip" "FIXED_LEN_BYTE_ARRAY" ...+-- The old parser had no page decoder for FIXED_LEN_BYTE_ARRAY and threw+-- before ever inspecting the encoding. The new parser handles the physical+-- type but the BYTE_STREAM_SPLIT encoding used for values is not yet+-- implemented, so the error message shifts from the type to the encoding.+-- TODO: Same as byteStreamSplitZstd — change to assertEqual once+-- BYTE_STREAM_SPLIT encoding is supported. byteStreamSplitExtendedGzip :: Test byteStreamSplitExtendedGzip = TestCase ( assertExpectException "byteStreamSplitExtendedGzip"- "FIXED_LEN_BYTE_ARRAY"+ "BYTE_STREAM_SPLIT" (D.readParquet "./tests/data/byte_stream_split_extended.gzip.parquet") ) +-- Was: assertExpectException "float16NonzerosAndNans" "PFIXED_LEN_BYTE_ARRAY" ...+-- The "PFIXED_LEN_BYTE_ARRAY" in the old error was the Show of the old+-- parser's ParquetType enum hitting a catch-all dispatch branch — it+-- recognised the physical type but had no decoder for it. The new parser's+-- fixedLenByteArrayDecoder reads 2-byte FIXED_LEN_BYTE_ARRAY (float16)+-- columns as raw-byte text; proper float16 value decoding is not yet+-- implemented.+-- TODO: When IEEE 754 half-precision (float16) decoding is implemented,+-- add a value-level assertion using hasElemType @Float (or a dedicated+-- Float16 type if one is introduced). Verify that the decoded values match+-- the known reference values for float16_nonzeros_and_nans.parquet.+-- The column should no longer be exposed as raw-byte Text. float16NonzerosAndNans :: Test float16NonzerosAndNans = TestCase- ( assertExpectException+ ( assertEqual "float16NonzerosAndNans"- "PFIXED_LEN_BYTE_ARRAY"- (D.readParquet "./tests/data/float16_nonzeros_and_nans.parquet")+ (8, 1)+ ( unsafePerformIO+ ( fmap+ D.dimensions+ (D.readParquet "./tests/data/float16_nonzeros_and_nans.parquet")+ )+ ) ) +-- Was: assertExpectException "float16ZerosAndNans" "PFIXED_LEN_BYTE_ARRAY" ...+-- Same as float16NonzerosAndNans: old parser had no decoder for the+-- FIXED_LEN_BYTE_ARRAY physical type; new parser reads raw bytes as text.+-- TODO: Same as float16NonzerosAndNans — add a value-level assertion once+-- float16 decoding is implemented. float16ZerosAndNans :: Test float16ZerosAndNans = TestCase- ( assertExpectException+ ( assertEqual "float16ZerosAndNans"- "PFIXED_LEN_BYTE_ARRAY"- (D.readParquet "./tests/data/float16_zeros_and_nans.parquet")+ (3, 1)+ ( unsafePerformIO+ (fmap D.dimensions (D.readParquet "./tests/data/float16_zeros_and_nans.parquet"))+ ) ) nestedListsSnappy :: Test@@ -1011,12 +1206,20 @@ ) ) +-- Was: assertExpectException "unknownLogicalType" "Unknown logical type" ...+-- The old parser raised a custom "Unknown logical type" message. The new+-- Pinch-based metadata parser raises "Field 16 is absent" for the+-- unrecognised LogicalType variant in this file.+-- TODO: If Pinch is extended to support forward-compatible decoding of+-- unknown union variants (treating unrecognised logical-type IDs as absent+-- rather than raising an error), change this to assertEqual where the file+-- parses successfully and the column falls back to its physical type. unknownLogicalType :: Test unknownLogicalType = TestCase ( assertExpectException "unknownLogicalType"- "Unknown logical type"+ "Field 16 is absent" (D.readParquet "./tests/data/unknown-logical-type.parquet") ) @@ -1024,13 +1227,24 @@ -- Group 12: Malformed files -- --------------------------------------------------------------------------- +-- Was: assertExpectException "nationDictMalformed" "dict index count mismatch" ...+-- The old parser validated the dictionary entry count against data-page+-- indices and raised "dict index count mismatch". The new parser does not+-- replicate that check; the dictionary bytes happen to decode correctly+-- despite the metadata discrepancy, returning the complete 25-row dataset.+-- TODO: If a stricter dictionary-validation pass is added (checking that+-- the number of decoded entries matches num_values in the dictionary page+-- header), revert this to assertExpectException with a count-mismatch+-- substring. nationDictMalformed :: Test nationDictMalformed = TestCase- ( assertExpectException+ ( assertEqual "nationDictMalformed"- "dict index count mismatch"- (D.readParquet "./tests/data/nation.dict-malformed.parquet")+ (25, 4)+ ( unsafePerformIO+ (fmap D.dimensions (D.readParquet "./tests/data/nation.dict-malformed.parquet"))+ ) ) shardedNullableSchema :: Test@@ -1038,22 +1252,26 @@ TestCase $ do metas <- mapM- (fmap fst . DP.readMetadataFromPath)+ DP.readMetadataFromPath ["data/sharded/part-0.parquet", "data/sharded/part-1.parquet"] let nullableCols = S.fromList [ last (map T.pack colPath) | meta <- metas- , rg <- rowGroups meta- , cc <- rowGroupColumns rg- , let cm = columnMetaData cc- colPath = columnPathInSchema cm+ , rg <- unField meta.row_groups+ , cc <- unField rg.rg_columns+ , Just cm <- [unField cc.cc_meta_data]+ , let colPath = map T.unpack (unField cm.cmd_path_in_schema) , not (null colPath)- , columnNullCount (columnStatistics cm) > 0+ , let nc :: Int64+ nc = case unField cm.cmd_statistics of+ Nothing -> 0+ Just stats -> fromMaybe 0 (unField stats.stats_null_count)+ , nc > 0 ] df = foldl- (\acc meta -> acc <> F.schemaToEmptyDataFrame nullableCols (schema meta))+ (\acc meta -> acc <> PS.schemaToEmptyDataFrame nullableCols (unField meta.schema)) D.empty metas assertBool "id should be nullable" (hasMissing (unsafeGetColumn "id" df))@@ -1063,18 +1281,22 @@ singleShardNoNulls :: Test singleShardNoNulls = TestCase $ do- (meta, _) <- DP.readMetadataFromPath "data/sharded/part-0.parquet"+ meta <- DP.readMetadataFromPath "data/sharded/part-0.parquet" let nullableCols = S.fromList [ last (map T.pack colPath)- | rg <- rowGroups meta- , cc <- rowGroupColumns rg- , let cm = columnMetaData cc- colPath = columnPathInSchema cm+ | rg <- unField meta.row_groups+ , cc <- unField rg.rg_columns+ , Just cm <- [unField cc.cc_meta_data]+ , let colPath = map T.unpack (unField cm.cmd_path_in_schema) , not (null colPath)- , columnNullCount (columnStatistics cm) > 0+ , let nc :: Int64+ nc = case unField cm.cmd_statistics of+ Nothing -> 0+ Just stats -> fromMaybe 0 (unField stats.stats_null_count)+ , nc > 0 ]- df = F.schemaToEmptyDataFrame nullableCols (schema meta)+ df = PS.schemaToEmptyDataFrame nullableCols (unField meta.schema) assertBool "id should NOT be nullable" (not (hasMissing (unsafeGetColumn "id" df)))